app.py

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()