Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,092 Bytes
a038241 384ccc2 a038241 86f1fe1 a038241 384ccc2 a038241 384ccc2 86f1fe1 384ccc2 a038241 384ccc2 a038241 384ccc2 a038241 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 |
#!/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 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=device)
else:
model_name = model_name.replace("retinaface_", "")
model = RetinaFacePredictor(
threshold=threshold, device=device, model=RetinaFacePredictor.get_model(model_name)
)
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}
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,
cache_examples=os.getenv("CACHE_EXAMPLES") == "1",
)
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()
|