app.py

import uuid

from fire import Fire
from head_detector import HeadDetector
import gradio as gr

detector = HeadDetector()


def detect_faces(image):
    image = detector(image).draw()
    return image


iface = gr.Interface(
    fn=detect_faces,
    inputs=gr.inputs.Image(type="numpy", label="Upload an Image"),
    outputs=gr.outputs.Image(type="numpy", label="Image with Faces"),
    title="Demo of: VGGHeads: A Large-Scale Synthetic Dataset for 3D Human Heads",
    description="Upload an image. The model performs simultaneous heads detection and head meshes reconstruction from a single image in a single step.",
    allow_flagging='never',
)


def main():
    iface.launch()


if __name__ == '__main__':
    Fire(main)