|
|
import os |
|
|
os.system("pip install git+https://github.com/elliottzheng/face-detection.git@master") |
|
|
os.system("git clone https://github.com/thohemp/6DRepNet") |
|
|
|
|
|
import sys |
|
|
sys.path.append("6DRepNet") |
|
|
|
|
|
import numpy as np |
|
|
import gradio as gr |
|
|
import torch |
|
|
from huggingface_hub import hf_hub_download |
|
|
|
|
|
from face_detection import RetinaFace |
|
|
from model import SixDRepNet |
|
|
import utils |
|
|
import cv2 |
|
|
from PIL import Image |
|
|
|
|
|
snapshot_path = hf_hub_download(repo_id="osanseviero/6DRepNet_300W_LP_AFLW2000", filename="model.pth") |
|
|
|
|
|
model = SixDRepNet(backbone_name='RepVGG-B1g2', |
|
|
backbone_file='', |
|
|
deploy=True, |
|
|
pretrained=False) |
|
|
|
|
|
detector = RetinaFace(0) |
|
|
saved_state_dict = torch.load(os.path.join( |
|
|
snapshot_path), map_location='cpu') |
|
|
|
|
|
if 'model_state_dict' in saved_state_dict: |
|
|
model.load_state_dict(saved_state_dict['model_state_dict']) |
|
|
else: |
|
|
model.load_state_dict(saved_state_dict) |
|
|
model.cuda(0) |
|
|
model.eval() |
|
|
|
|
|
def predict(frame): |
|
|
faces = detector(frame) |
|
|
for box, landmarks, score in faces: |
|
|
|
|
|
if score < .95: |
|
|
continue |
|
|
x_min = int(box[0]) |
|
|
y_min = int(box[1]) |
|
|
x_max = int(box[2]) |
|
|
y_max = int(box[3]) |
|
|
bbox_width = abs(x_max - x_min) |
|
|
bbox_height = abs(y_max - y_min) |
|
|
|
|
|
x_min = max(0,x_min-int(0.2*bbox_height)) |
|
|
y_min = max(0,y_min-int(0.2*bbox_width)) |
|
|
x_max = x_max+int(0.2*bbox_height) |
|
|
y_max = y_max+int(0.2*bbox_width) |
|
|
|
|
|
img = frame[y_min:y_max,x_min:x_max] |
|
|
img = cv2.resize(img, (244, 244))/255.0 |
|
|
img = img.transpose(2, 0, 1) |
|
|
img = torch.from_numpy(img).type(torch.FloatTensor) |
|
|
img = torch.Tensor(img).cuda(0) |
|
|
img=img.unsqueeze(0) |
|
|
R_pred = model(img) |
|
|
euler = utils.compute_euler_angles_from_rotation_matrices( |
|
|
R_pred)*180/np.pi |
|
|
p_pred_deg = euler[:, 0].cpu() |
|
|
y_pred_deg = euler[:, 1].cpu() |
|
|
r_pred_deg = euler[:, 2].cpu() |
|
|
return utils.plot_pose_cube(frame, y_pred_deg, p_pred_deg, r_pred_deg, x_min + int(.5*(x_max-x_min)), y_min + int(.5*(y_max-y_min)), size = bbox_width) |
|
|
|
|
|
title = "6D Rotation Representation for Unconstrained Head Pose Estimation" |
|
|
description = "Gradio demo for 6DRepNet. To use it, simply click the camera picture. Read more at the links below." |
|
|
article = "<div style='text-align: center;'><a href='https://github.com/thohemp/6DRepNet' target='_blank'>Github Repo</a> | <a href='https://arxiv.org/abs/2202.12555' target='_blank'>Paper</a></div>" |
|
|
|
|
|
image_flip_css = """ |
|
|
.input-image .image-preview img{ |
|
|
-webkit-transform: scaleX(-1); |
|
|
transform: scaleX(-1) !important; |
|
|
} |
|
|
|
|
|
.output-image img { |
|
|
-webkit-transform: scaleX(-1); |
|
|
transform: scaleX(-1) !important; |
|
|
} |
|
|
""" |
|
|
|
|
|
iface = gr.Interface( |
|
|
fn=predict, |
|
|
inputs=gr.inputs.Image(label="Input Image", source="webcam"), |
|
|
outputs='image', |
|
|
live=True, |
|
|
title=title, |
|
|
description=description, |
|
|
article=article, |
|
|
css = image_flip_css |
|
|
) |
|
|
|
|
|
iface.launch() |
