Spaces:
Running
on
Zero
Running
on
Zero
| #!/usr/bin/env python | |
| from __future__ import annotations | |
| import pathlib | |
| import sys | |
| import cv2 | |
| import gradio as gr | |
| import numpy as np | |
| import spaces | |
| import torch | |
| sys.path.insert(0, "face_detection") | |
| sys.path.insert(0, "face_alignment") | |
| from ibug.face_alignment import FANPredictor | |
| from ibug.face_detection import RetinaFacePredictor | |
| DESCRIPTION = "# [ibug-group/face_alignment](https://github.com/ibug-group/face_alignment)" | |
| device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") | |
| detector = RetinaFacePredictor(threshold=0.8, device="cpu", model=RetinaFacePredictor.get_model("mobilenet0.25")) | |
| detector.device = device | |
| detector.net.to(device) | |
| def load_model(model_name: str, device: torch.device) -> FANPredictor: | |
| model = FANPredictor( | |
| device="cpu", model=FANPredictor.get_model(model_name), config=FANPredictor.create_config(use_jit=False) | |
| ) | |
| model.device = device | |
| model.net.to(device) | |
| return model | |
| model_names = [ | |
| "2dfan2", | |
| "2dfan4", | |
| "2dfan2_alt", | |
| ] | |
| models = {name: load_model(name, device) for name in model_names} | |
| def predict(image: np.ndarray, model_name: str, max_num_faces: int, landmark_score_threshold: int) -> np.ndarray: | |
| model = models[model_name] | |
| # RGB -> BGR | |
| image = image[:, :, ::-1] | |
| faces = detector(image, rgb=False) | |
| if len(faces) == 0: | |
| raise RuntimeError("No face was found.") | |
| faces = sorted(list(faces), key=lambda x: -x[4])[:max_num_faces] | |
| faces = np.asarray(faces) | |
| landmarks, landmark_scores = model(image, faces, rgb=False) | |
| res = image.copy() | |
| for face, pts, scores in zip(faces, landmarks, landmark_scores): | |
| box = np.round(face[:4]).astype(int) | |
| cv2.rectangle(res, tuple(box[:2]), tuple(box[2:]), (0, 255, 0), 2) | |
| for pt, score in zip(np.round(pts).astype(int), scores): | |
| if score < landmark_score_threshold: | |
| continue | |
| cv2.circle(res, tuple(pt), 2, (0, 255, 0), cv2.FILLED) | |
| return res[:, :, ::-1] | |
| examples = [[path.as_posix(), model_names[0], 10, 0.2] for path in pathlib.Path("images").rglob("*.jpg")] | |
| 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=model_names[0], label="Model") | |
| max_num_faces = gr.Slider(1, 20, step=1, value=10, label="Max Number of Faces") | |
| landmark_score_thrshold = gr.Slider(0, 1, step=0.05, value=0.2, label="Landmark Score Threshold") | |
| run_button = gr.Button() | |
| with gr.Column(): | |
| result = gr.Image(label="Output") | |
| gr.Examples( | |
| examples=examples, | |
| inputs=[image, model_name, max_num_faces, landmark_score_thrshold], | |
| outputs=result, | |
| fn=predict, | |
| ) | |
| run_button.click( | |
| fn=predict, | |
| inputs=[image, model_name, max_num_faces, landmark_score_thrshold], | |
| outputs=result, | |
| api_name="predict", | |
| ) | |
| if __name__ == "__main__": | |
| demo.queue(max_size=20).launch() | |