File size: 4,187 Bytes
a038241
 
 
 
384ccc2
a038241
 
 
 
 
 
 
f62a68b
a038241
86f1fe1
a038241
384ccc2
a038241
 
 
384ccc2
 
 
86f1fe1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
384ccc2
f62a68b
 
 
 
a038241
384ccc2
f62a68b
 
 
a038241
 
 
384ccc2
 
 
 
 
 
 
 
 
f62a68b
384ccc2
a038241
 
 
 
 
 
 
 
 
 
 
 
384ccc2
a038241
 
 
 
 
 
 
 
 
384ccc2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a038241
 
 
384ccc2
 
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
118
119
120
121
122
123
124
125
126
127
128
129
130
#!/usr/bin/env python

from __future__ import annotations

import os
import pathlib
import sys
import urllib.request

import cv2
import gradio as gr
import numpy as np
import spaces
import torch
from huggingface_hub import hf_hub_download

sys.path.insert(0, "face_detection")

from ibug.face_detection import RetinaFacePredictor, S3FDPredictor

DESCRIPTION = "# [ibug-group/face_detection](https://github.com/ibug-group/face_detection)"


def is_lfs_pointer_file(path: pathlib.Path) -> bool:
    try:
        with open(path, "r") as f:
            # Git LFS pointer files usually start with version line
            version_line = f.readline()
            if version_line.startswith("version https://git-lfs.github.com/spec/"):
                # Check for the presence of oid and size lines
                oid_line = f.readline()
                size_line = f.readline()
                if oid_line.startswith("oid sha256:") and size_line.startswith("size "):
                    return True
    except Exception as e:
        print(f"Error reading file {path}: {e}")
    return False


lfs_model_path = pathlib.Path("face_detection/ibug/face_detection/retina_face/weights/Resnet50_Final.pth")
if is_lfs_pointer_file(lfs_model_path):
    os.remove(lfs_model_path)
    out_path = hf_hub_download(
        "public-data/ibug-face-detection",
        filename=lfs_model_path.name,
        repo_type="model",
        subfolder="retina_face",
    )
    os.symlink(out_path, lfs_model_path)


def load_model(model_name: str, threshold: float, device: torch.device) -> RetinaFacePredictor | S3FDPredictor:
    if model_name == "s3fd":
        model = S3FDPredictor(threshold=threshold, device="cpu")
        model.device = device
        model.net.device = device
        model.net.to(device)
    else:
        model_name = model_name.replace("retinaface_", "")
        model = RetinaFacePredictor(threshold=threshold, device="cpu", model=RetinaFacePredictor.get_model(model_name))
        model.device = device
        model.net.to(device)
    return model


device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model_names = [
    "retinaface_mobilenet0.25",
    "retinaface_resnet50",
    "s3fd",
]
detectors = {name: load_model(name, threshold=0.8, device=device) for name in model_names}


@spaces.GPU
def detect(image: np.ndarray, model_name: str, face_score_threshold: float) -> np.ndarray:
    model = detectors[model_name]
    model.threshold = face_score_threshold

    # RGB -> BGR
    image = image[:, :, ::-1]
    preds = model(image, rgb=False)

    res = image.copy()
    for pred in preds:
        box = np.round(pred[:4]).astype(int)

        line_width = max(2, int(3 * (box[2:] - box[:2]).max() / 256))
        cv2.rectangle(res, tuple(box[:2]), tuple(box[2:]), (0, 255, 0), line_width)

        if len(pred) == 15:
            pts = pred[5:].reshape(-1, 2)
            for pt in np.round(pts).astype(int):
                cv2.circle(res, tuple(pt), line_width, (0, 255, 0), cv2.FILLED)

    return res[:, :, ::-1]


example_image_path = pathlib.Path("selfie.jpg")
if not example_image_path.exists():
    url = "https://raw.githubusercontent.com/peiyunh/tiny/master/data/demo/selfie.jpg"
    urllib.request.urlretrieve(url, example_image_path)

with gr.Blocks(css="style.css") as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        with gr.Column():
            image = gr.Image(type="numpy", label="Input")
            model_name = gr.Radio(model_names, type="value", value="retinaface_resnet50", label="Model")
            score_threshold = gr.Slider(minimum=0, maximum=1, step=0.05, value=0.8, label="Face Score Threshold")
            run_button = gr.Button()
        with gr.Column():
            result = gr.Image(label="Output")
    gr.Examples(
        examples=[[example_image_path.as_posix(), model_names[1], 0.8]],
        inputs=[image, model_name, score_threshold],
        outputs=result,
        fn=detect,
    )
    run_button.click(
        fn=detect,
        inputs=[image, model_name, score_threshold],
        outputs=result,
        api_name="detect",
    )


if __name__ == "__main__":
    demo.queue(max_size=20).launch()