App update

app.py

@@ -5,16 +5,26 @@ from head_detector import HeadDetector
 import gradio as gr
 
 detector = HeadDetector()
+INVERSE_MAPPING = {
+    'Full': 'full',
+    'Head Boxes': 'bbox',
+    'Face Landmarks': 'points',
+    'Head Landmarks': 'landmarks',
+    'Head Pose': 'pose'
+}
 
 
-def detect_faces(image):
-    image = detector(image).draw()
+def detect_faces(image, method):
+    image = detector(image).draw(method=INVERSE_MAPPING[method])
     return image
 
 
 iface = gr.Interface(
     fn=detect_faces,
-    inputs=gr.inputs.Image(type="numpy", label="Upload an Image"),
+    inputs=[
+        gr.inputs.Image(type="numpy", label="Upload an Image"),
+        gr.inputs.Radio(['Full', 'Head Boxes', 'Face Landmarks', 'Head Landmarks', 'Head Pose'], label="Face Detection Method", default='Full')
+    ],
     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.",