File size: 3,711 Bytes
319d3b5 1fed86e 319d3b5 1fed86e 319d3b5 1fed86e 319d3b5 1fed86e 319d3b5 1fed86e 319d3b5 1fed86e 319d3b5 1011fc1 319d3b5 1fed86e 319d3b5 1011fc1 319d3b5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 |
import os
import cv2
import imghdr
import shutil
import warnings
import numpy as np
import gradio as gr
from dataclasses import dataclass
from mivolo.predictor import Predictor
from utils import is_url, download_file, get_jpg_files, MODEL_DIR
TMP_DIR = "./__pycache__"
@dataclass
class Cfg:
detector_weights: str
checkpoint: str
device: str = "cpu"
with_persons: bool = True
disable_faces: bool = False
draw: bool = True
class ValidImgDetector:
predictor = None
def __init__(self):
detector_path = f"{MODEL_DIR}/yolov8x_person_face.pt"
age_gender_path = f"{MODEL_DIR}/model_imdb_cross_person_4.22_99.46.pth.tar"
predictor_cfg = Cfg(detector_path, age_gender_path)
self.predictor = Predictor(predictor_cfg)
def _detect(
self,
image: np.ndarray,
score_threshold: float,
iou_threshold: float,
mode: str,
predictor: Predictor,
) -> np.ndarray:
predictor.detector.detector_kwargs["conf"] = score_threshold
predictor.detector.detector_kwargs["iou"] = iou_threshold
if mode == "Use persons and faces":
use_persons = True
disable_faces = False
elif mode == "Use persons only":
use_persons = True
disable_faces = True
elif mode == "Use faces only":
use_persons = False
disable_faces = False
predictor.age_gender_model.meta.use_persons = use_persons
predictor.age_gender_model.meta.disable_faces = disable_faces
detected_objects, out_im = predictor.recognize(image)
has_child, has_female, has_male = False, False, False
if len(detected_objects.ages) > 0:
has_child = min(detected_objects.ages) < 18
has_female = "female" in detected_objects.genders
has_male = "male" in detected_objects.genders
return out_im[:, :, ::-1], has_child, has_female, has_male
def valid_video(self, video_path):
cap = cv2.VideoCapture(video_path)
results = []
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
# Process each frame
out_frame, has_child, has_female, has_male = self._detect(frame, 0.4, 0.7, "Use persons and faces", self.predictor)
results.append((out_frame, has_child, has_female, has_male))
cap.release()
return results
def infer(video_path: str):
if is_url(video_path):
if os.path.exists(TMP_DIR):
shutil.rmtree(TMP_DIR)
video_path = download_file(video_path, f"{TMP_DIR}/download.mp4")
detector = ValidImgDetector()
if not video_path or not os.path.exists(video_path):
return None, None, None, "Please input the video correctly"
results = detector.valid_video(video_path)
# Process results to only return first frame's result as an example
if results:
first_frame_result = results[0]
return first_frame_result
else:
return None, None, None, "No frames detected in video."
if __name__ == "__main__":
with gr.Blocks() as iface:
warnings.filterwarnings("ignore")
with gr.Tab("Upload Mode"):
gr.Interface(
fn=infer,
inputs=gr.Video(label="Upload Video"),
outputs=[
gr.Image(label="Detection Result", type="numpy"),
gr.Textbox(label="Has Child"),
gr.Textbox(label="Has Female"),
gr.Textbox(label="Has Male"),
],
allow_flagging="never",
)
iface.launch()
|