| 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, _L, MODEL_DIR, TMP_DIR |
|
|
|
|
| @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 = _L("是") if min(detected_objects.ages) < 18 else _L("否") |
| has_female = _L("是") if "female" in detected_objects.genders else _L("否") |
| has_male = _L("是") if "male" in detected_objects.genders else _L("否") |
|
|
| 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 infer(photo: str): |
| status = "Success" |
| result = child = female = male = None |
| try: |
| if is_url(photo): |
| if os.path.exists(TMP_DIR): |
| shutil.rmtree(TMP_DIR) |
|
|
| photo = download_file(photo, f"{TMP_DIR}/download.jpg") |
|
|
| detector = ValidImgDetector() |
| if not photo or not os.path.exists(photo) or imghdr.what(photo) == None: |
| raise ValueError("请正确输入图片") |
|
|
| result, child, female, male = detector.valid_img(photo) |
|
|
| except Exception as e: |
| status = f"{e}" |
|
|
| return status, result, child, female, male |
|
|
|
|
| if __name__ == "__main__": |
| warnings.filterwarnings("ignore") |
| with gr.Blocks() as iface: |
| gr.Markdown(_L("# 性别年龄检测器")) |
| with gr.Tab(_L("上传模式")): |
| gr.Interface( |
| fn=infer, |
| inputs=gr.Image(label=_L("上传照片"), type="filepath"), |
| outputs=[ |
| gr.Textbox(label=_L("状态栏"), show_copy_button=True), |
| gr.Image( |
| label=_L("检测结果"), |
| type="numpy", |
| show_share_button=False, |
| ), |
| gr.Textbox(label=_L("存在儿童")), |
| gr.Textbox(label=_L("存在女性")), |
| gr.Textbox(label=_L("存在男性")), |
| ], |
| examples=get_jpg_files(f"{MODEL_DIR}/examples"), |
| flagging_mode="never", |
| cache_examples=False, |
| ) |
|
|
| with gr.Tab(_L("在线模式")): |
| gr.Interface( |
| fn=infer, |
| inputs=gr.Textbox( |
| label=_L("网络图片链接"), |
| show_copy_button=True, |
| ), |
| outputs=[ |
| gr.Textbox(label=_L("状态栏"), show_copy_button=True), |
| gr.Image( |
| label=_L("检测结果"), |
| type="numpy", |
| show_share_button=False, |
| ), |
| gr.Textbox(label=_L("存在儿童")), |
| gr.Textbox(label=_L("存在女性")), |
| gr.Textbox(label=_L("存在男性")), |
| ], |
| flagging_mode="never", |
| ) |
|
|
| iface.launch() |
|
|
