hysts/insightface-SCRFD

1#!/usr/bin/env python2 3from __future__ import annotations4 5import functools6import os7import pathlib8import shlex9import subprocess10import sys11import urllib.request12 13if os.environ.get('SYSTEM') == 'spaces':14    import mim15    mim.install('mmcv-full==1.4', is_yes=True)16 17    subprocess.call(shlex.split('pip uninstall -y opencv-python'))18    subprocess.call(shlex.split('pip uninstall -y opencv-python-headless'))19    subprocess.call(20        shlex.split('pip install opencv-python-headless==4.5.5.64'))21    subprocess.call(shlex.split('pip install terminaltables==3.1.0'))22    subprocess.call(shlex.split('pip install mmpycocotools==12.0.3'))23 24    subprocess.call(shlex.split('pip install insightface==0.6.2'))25    subprocess.call(shlex.split('sed -i 23,26d __init__.py'),26                    cwd='insightface/detection/scrfd/mmdet')27 28import cv229import gradio as gr30import huggingface_hub31import numpy as np32import torch33import torch.nn as nn34 35sys.path.insert(0, 'insightface/detection/scrfd')36 37from mmdet.apis import inference_detector, init_detector, show_result_pyplot38 39TITLE = 'insightface Face Detection (SCRFD)'40DESCRIPTION = 'This is an unofficial demo for https://github.com/deepinsight/insightface/tree/master/detection/scrfd.'41 42HF_TOKEN = os.getenv('HF_TOKEN')43 44 45def load_model(model_size: str, device) -> nn.Module:46    ckpt_path = huggingface_hub.hf_hub_download(47        'hysts/insightface',48        f'models/scrfd_{model_size}/model.pth',49        use_auth_token=HF_TOKEN)50    scrfd_dir = 'insightface/detection/scrfd'51    config_path = f'{scrfd_dir}/configs/scrfd/scrfd_{model_size}.py'52    model = init_detector(config_path, ckpt_path, device.type)53    return model54 55 56def update_test_pipeline(model: nn.Module, mode: int):57    cfg = model.cfg58    pipelines = cfg.data.test.pipeline59    for pipeline in pipelines:60        if pipeline.type == 'MultiScaleFlipAug':61            if mode == 0:  # 640 scale62                pipeline.img_scale = (640, 640)63                if hasattr(pipeline, 'scale_factor'):64                    del pipeline.scale_factor65            elif mode == 1:  # for single scale in other pages66                pipeline.img_scale = (1100, 1650)67                if hasattr(pipeline, 'scale_factor'):68                    del pipeline.scale_factor69            elif mode == 2:  # original scale70                pipeline.img_scale = None71                pipeline.scale_factor = 1.072            transforms = pipeline.transforms73            for transform in transforms:74                if transform.type == 'Pad':75                    if mode != 2:76                        transform.size = pipeline.img_scale77                        if hasattr(transform, 'size_divisor'):78                            del transform.size_divisor79                    else:80                        transform.size = None81                        transform.size_divisor = 3282 83 84def detect(image: np.ndarray, model_size: str, mode: int,85           face_score_threshold: float,86           detectors: dict[str, nn.Module]) -> np.ndarray:87    model = detectors[model_size]88    update_test_pipeline(model, mode)89 90    # RGB -> BGR91    image = image[:, :, ::-1]92    preds = inference_detector(model, image)93    boxes = preds[0]94 95    res = image.copy()96    for box in boxes:97        box, score = box[:4], box[4]98        if score < face_score_threshold:99            continue100        box = np.round(box).astype(int)101 102        line_width = max(2, int(3 * (box[2:] - box[:2]).max() / 256))103        cv2.rectangle(res, tuple(box[:2]), tuple(box[2:]), (0, 255, 0),104                      line_width)105 106    res = cv2.cvtColor(res, cv2.COLOR_BGR2RGB)107    return res108 109 110device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')111 112model_sizes = [113    '500m',114    '1g',115    '2.5g',116    '10g',117    '34g',118]119detectors = {120    model_size: load_model(model_size, device=device)121    for model_size in model_sizes122}123modes = [124    '(640, 640)',125    '(1100, 1650)',126    'original',127]128 129func = functools.partial(detect, detectors=detectors)130 131image_path = pathlib.Path('selfie.jpg')132if not image_path.exists():133    url = 'https://raw.githubusercontent.com/peiyunh/tiny/master/data/demo/selfie.jpg'134    urllib.request.urlretrieve(url, image_path)135examples = [[image_path.as_posix(), '10g', modes[0], 0.3]]136 137gr.Interface(138    fn=func,139    inputs=[140        gr.Image(label='Input', type='numpy'),141        gr.Radio(label='Model', choices=model_sizes, type='value',142                 value='10g'),143        gr.Radio(label='Mode', choices=modes, type='index', value=modes[0]),144        gr.Slider(label='Face Score Threshold',145                  minimum=0,146                  maximum=1,147                  step=0.05,148                  default=0.3),149    ],150    outputs=gr.Image(label='Output', type='numpy'),151    examples=examples,152    title=TITLE,153    description=DESCRIPTION,154).queue().launch(show_api=False)155