Spaces:
Runtime error
Runtime error
File size: 6,181 Bytes
f11983f fd57de8 f11983f fd57de8 f11983f fd57de8 f11983f fd57de8 f11983f fd57de8 f11983f fd57de8 7a6049c fd57de8 f11983f 967a6ae f11983f 967a6ae f11983f fd57de8 967a6ae f11983f fd57de8 f11983f fd57de8 f11983f fd57de8 f11983f fd57de8 967a6ae 739106a 967a6ae f11983f |
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 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 |
import os
import sys
from pathlib import Path
import uuid
import gradio as gr
from PIL import Image
import cv2
import torch
import numpy as np
import copy
import retinaface
device = 'cuda' if torch.cuda.is_available() else 'cpu'
try:
from spiga.demo.app import video_app
from spiga.inference.config import ModelConfig
from spiga.inference.framework import SPIGAFramework
from spiga.demo.visualize.plotter import Plotter
import spiga.demo.analyze.track.retinasort.config as cfg
except:
os.system("pip install -e ./SPIGA[demo]")
sys.path.append(os.path.abspath("./SPIGA"))
from spiga.demo.app import video_app
from spiga.inference.config import ModelConfig
from spiga.inference.framework import SPIGAFramework
from spiga.demo.visualize.plotter import Plotter
import spiga.demo.analyze.track.retinasort.config as cfg
face_processor = SPIGAFramework(ModelConfig('wflw'))
config = cfg.cfg_retinasort
face_detector = retinaface.RetinaFaceDetector(model=config['retina']['model_name'],
device=device,
extra_features=config['retina']['extra_features'],
cfg_postreat=config['retina']['postreat'])
def predict_video(video_in):
if video_in == None:
raise gr.Error("Please upload a video or image.")
video_in = Path(video_in)
output_path = Path("/tmp")
video_file_name = str(uuid.uuid4())
new_video_path = output_path / f"{video_file_name}{video_in.suffix}"
video_in.rename(new_video_path)
video_app(str(new_video_path),
# Choices=['wflw', '300wpublic', '300wprivate', 'merlrav']
spiga_dataset='wflw',
# Choices=['RetinaSort', 'RetinaSort_Res50']
tracker='RetinaSort',
save=True,
output_path=output_path,
visualize=False,
plot=['fps', 'face_id', 'landmarks', 'headpose'])
video_output_path = f"{output_path}/{new_video_path.name[:-4]}.mp4"
return video_output_path
def predict_image(image_in_video, image_in_img):
if image_in_video == None and image_in_img == None:
raise gr.Error("Please upload a video or image.")
if image_in_video or image_in_img:
print("image", image_in_video, image_in_img)
image = image_in_img if image_in_img else image_in_video
image = Image.open(image).convert("RGB")
cv2_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
face_detector.set_input_shape(image.size[1], image.size[0])
features = face_detector.inference(image)
if features:
bboxes = features['bbox']
bboxes_n = []
for bbox in bboxes:
x1, y1, x2, y2 = bbox[:4]
bbox_wh = [x1, y1, x2-x1, y2-y1]
bboxes_n.append(bbox_wh)
face_features = face_processor.inference(cv2_image, bboxes_n)
canvas = copy.deepcopy(cv2_image)
plotter = Plotter()
for landmarks, headpose, bbox in zip(face_features['landmarks'], face_features['headpose'], bboxes_n):
x0, y0, w, h = bbox
landmarks = np.array(landmarks)
headpose = np.array(headpose)
canvas = plotter.landmarks.draw_landmarks(canvas, landmarks)
canvas = plotter.hpose.draw_headpose(
canvas, [x0, y0, x0+w, y0+h], headpose[:3], headpose[3:], euler=True)
return Image.fromarray(cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB))
def toggle(choice):
if choice == "webcam":
return gr.update(visible=True, value=None), gr.update(visible=False, value=None)
else:
return gr.update(visible=False, value=None), gr.update(visible=True, value=None)
with gr.Blocks() as blocks:
gr.Markdown("""
# Unofficia Demo: SPIGA
## Shape Preserving Facial Landmarks with Graph Attention Networks.
* https://github.com/andresprados/SPIGA
* https://arxiv.org/abs/2210.07233
""")
with gr.Tab("Video") as tab:
with gr.Row():
with gr.Column():
video_or_file_opt = gr.Radio(["upload", "webcam"], value="upload",
label="How would you like to upload your video?")
video_in = gr.Video(source="upload", include_audio=False)
video_or_file_opt.change(fn=lambda s: gr.update(source=s, value=None), inputs=video_or_file_opt,
outputs=video_in, queue=False)
with gr.Column():
video_out = gr.Video()
run_btn = gr.Button("Run")
run_btn.click(fn=predict_video, inputs=[video_in], outputs=[video_out])
gr.Examples(fn=predict_video, examples=[["./examples/temp.mp4"]],
inputs=[video_in], outputs=[video_out],
cache_examples=True)
with gr.Tab("Image"):
with gr.Row():
with gr.Column():
image_or_file_opt = gr.Radio(["file", "webcam"], value="file",
label="How would you like to upload your image?")
image_in_video = gr.Image(
source="webcam", type="filepath", visible=False)
image_in_img = gr.Image(
source="upload", visible=True, type="filepath")
image_or_file_opt.change(fn=toggle, inputs=[image_or_file_opt],
outputs=[image_in_video, image_in_img], queue=False)
with gr.Column():
image_out = gr.Image()
run_btn = gr.Button("Run")
run_btn.click(fn=predict_image,
inputs=[image_in_img, image_in_video], outputs=[image_out])
gr.Examples(fn=predict_image, examples=[
["./examples/52_Photographers_photographertakingphoto_52_315.jpg", None],
["./examples/6_Funeral_Funeral_6_759.jpg", None]
],
inputs=[image_in_img, image_in_video], outputs=[image_out],
cache_examples=True
)
blocks.launch()
|