Spaces:
Configuration error
Configuration error
| #!/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} | |
| 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() | |