optional - Use retinaface for face detection (#7)

- optional - Use retinaface for face detection (c57f7aedf4dedc338af570caaecc4cee74d7f2bc)
- remove print (478af14d5268d7cec5efefcc368bee4c9e99b9f7)


Co-authored-by: Radamés Ajna <[email protected]>

app.py

@@ -1,27 +1,30 @@
 import gradio as gr
 import torch
-import dlib
 import numpy as np
 import PIL
 import base64
 from io import BytesIO
 from PIL import Image
-# Only used to convert to gray, could do it differently and remove this big dependency
-import cv2
+# import for face detection
+import retinaface
 
 from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
 from diffusers import UniPCMultistepScheduler
 
 from spiga.inference.config import ModelConfig
 from spiga.inference.framework import SPIGAFramework
+import spiga.demo.analyze.track.retinasort.config as cfg
 
 import matplotlib.pyplot as plt
 from matplotlib.path import Path
 import matplotlib.patches as patches
 
 # Bounding boxes
-face_detector = dlib.get_frontal_face_detector()
-
+config = cfg.cfg_retinasort
+face_detector = retinaface.RetinaFaceDetector(model=config['retina']['model_name'],
+                                              device='cuda' if torch.cuda.is_available() else 'cpu',
+                                              extra_features=config['retina']['extra_features'],
+                                              cfg_postreat=config['retina']['postreat'])
 # Landmark extraction
 spiga_extractor = SPIGAFramework(ModelConfig("300wpublic"))
 
@@ -59,14 +62,19 @@ async (image_in_img, prompt, image_file_live_opt, live_conditioning) => {
 }
 """
 
+
 def get_bounding_box(image):
-    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    faces = face_detector(gray)
-    if len(faces) == 0:
-        raise Exception("No face detected in image")
-    face = faces[0]
-    bbox = [face.left(), face.top(), face.width(), face.height()]
-    return bbox
+    pil_image = Image.fromarray(image)
+    face_detector.set_input_shape(pil_image.size[1], pil_image.size[0])
+    features = face_detector.inference(pil_image)
+
+    if (features is None) and (len(features['bbox']) <= 0):
+        raise Exception("No face detected")
+    # get the first face detected
+    bbox = features['bbox'][0]
+    x1, y1, x2, y2 = bbox[:4]
+    bbox_wh = [x1, y1, x2-x1, y2-y1]
+    return bbox_wh
 
 
 def get_landmarks(image, bbox):
@@ -145,7 +153,6 @@ def get_conditioning(image):
 def generate_images(image_in_img, prompt, image_file_live_opt='file', live_conditioning=None):
     if image_in_img is None and 'image' not in live_conditioning:
         raise gr.Error("Please provide an image")
-
     try:
         if image_file_live_opt == 'file':
             conditioning = get_conditioning(image_in_img)
@@ -166,29 +173,31 @@ def generate_images(image_in_img, prompt, image_file_live_opt='file', live_condi
     except Exception as e:
         raise gr.Error(str(e))
 
+
 def toggle(choice):
     if choice == "file":
         return gr.update(visible=True, value=None), gr.update(visible=False, value=None)
     elif choice == "webcam":
         return gr.update(visible=False, value=None), gr.update(visible=True, value=canvas_html)
 
+
 with gr.Blocks() as blocks:
     gr.Markdown("""
         ## Generate Uncanny Faces with ControlNet Stable Diffusion
         [Check out our blog to see how this was done (and train your own controlnet)](https://huggingface.co/blog/train-your-controlnet)
     """)
     with gr.Row():
-        live_conditioning  = gr.JSON(value={}, visible=False)
+        live_conditioning = gr.JSON(value={}, visible=False)
         with gr.Column():
             image_file_live_opt = gr.Radio(["file", "webcam"], value="file",
-                                         label="How would you like to upload your image?")
+                                           label="How would you like to upload your image?")
             image_in_img = gr.Image(source="upload", visible=True, type="pil")
             canvas = gr.HTML(None, elem_id="canvas_html", visible=False)
 
             image_file_live_opt.change(fn=toggle,
-                                     inputs=[image_file_live_opt],
-                                     outputs=[image_in_img, canvas],
-                                     queue=False)
+                                       inputs=[image_file_live_opt],
+                                       outputs=[image_in_img, canvas],
+                                       queue=False)
             prompt = gr.Textbox(
                 label="Enter your prompt",
                 max_lines=1,
@@ -198,7 +207,8 @@ with gr.Blocks() as blocks:
         with gr.Column():
             gallery = gr.Gallery().style(grid=[2], height="auto")
     run_button.click(fn=generate_images,
-                     inputs=[image_in_img, prompt, image_file_live_opt, live_conditioning],
+                     inputs=[image_in_img, prompt,
+                             image_file_live_opt, live_conditioning],
                      outputs=[gallery],
                      _js=get_js_image)
     blocks.load(None, None, None, _js=load_js)

requirements.txt

@@ -7,3 +7,4 @@ dlib
 opencv-python
 matplotlib
 Pillow
+retinaface-py>=0.0.2