Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,116 Bytes
45bd6c9 54102f1 45bd6c9 6d6900d 45bd6c9 6d6900d 54102f1 6d6900d 54102f1 6d6900d 54102f1 6d6900d 45bd6c9 6d6900d 45bd6c9 6d6900d 45bd6c9 6d6900d 45bd6c9 6d6900d |
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 |
#!/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}
@spaces.GPU
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()
|