Spaces:
Running
Running
| import cv2 | |
| import numpy as np | |
| import gradio as gr | |
| from dataclasses import dataclass | |
| from mivolo.predictor import Predictor | |
| from utils import * | |
| import warnings | |
| warnings.filterwarnings("ignore") | |
| 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): | |
| # Use a pipeline as a high-level helper | |
| detector_file = "yolov8x_person_face.pt" | |
| age_gender_file = "model_imdb_cross_person_4.22_99.46.pth.tar" | |
| detector_path = f"./model/{detector_file}" | |
| age_gender_path = f"./model/{age_gender_file}" | |
| domain = "https://huggingface.co/monet-joe/human-detector/resolve/main/" | |
| download_file(f"{domain}{detector_file}", detector_path) | |
| download_file(f"{domain}{age_gender_file}", age_gender_path) | |
| 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: | |
| # input is rgb image, output must be rgb too | |
| 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 | |
| # image = image[:, :, ::-1] # RGB -> BGR | |
| 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_img(self, img_path): | |
| image = cv2.imread(img_path) | |
| return self._detect( | |
| image, | |
| 0.4, | |
| 0.7, | |
| "Use persons and faces", | |
| self.predictor | |
| ) | |
| def inference(photo): | |
| detector = ValidImgDetector() | |
| return detector.valid_img(photo) | |
| iface = gr.Interface( | |
| fn=inference, | |
| inputs=gr.Image(label='Upload photo', type='filepath'), | |
| outputs=[ | |
| gr.Image(label='Detection result', type='numpy'), | |
| gr.Textbox(label='Has child'), | |
| gr.Textbox(label='Has female'), | |
| gr.Textbox(label='Has male') | |
| ], | |
| examples=get_jpg_files('./examples') | |
| ) | |
| iface.launch() | |