| 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() |
