app.py

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