updates

app.py

@@ -12,13 +12,48 @@ import gradio as gr
 import face_detection
 
 
-def detect_faces(img: np.ndarray):
-    frame = np.array(img) 
+def compare_detect_faces(img: np.ndarray, 
+                         confidence_threshold, 
+                         nms_threshold,
+                         kornia_toggle,
+                         retina_toggle,
+                         retina_mobile_toggle,
+                         dsfd_toggle
+                         ):
     
-    kornia_detections = kornia_detect(frame)
-    retina_detections = retina_detect(frame)
-    retina_mobile_detections = retina_mobilenet_detect(frame)
-    dsfd_detections = dsfd_detect(frame)
+    detections = []
+    
+    if kornia_toggle=="On":
+        kornia_detections = kornia_detect(img, 
+                                          confidence_threshold=confidence_threshold, 
+                                          nms_threshold=nms_threshold)
+    else:
+        kornia_detections = None
+        
+    if retina_toggle=="On":
+        retina_detections = retina_detect(img, 
+                                          confidence_threshold=confidence_threshold, 
+                                          nms_threshold=nms_threshold)
+        detections.append(retina_detections)
+    else:
+        retina_detections = None
+    
+    if retina_mobile_toggle=="On":
+        retina_mobile_detections = retina_mobilenet_detect(img, 
+                                          confidence_threshold=confidence_threshold, 
+                                          nms_threshold=nms_threshold)
+        detections.append(retina_mobile_detections)
+    else:
+        retina_mobile_detections = None
+
+    if dsfd_toggle=="On":
+        dsfd_detections = dsfd_detect(img, 
+                                          confidence_threshold=confidence_threshold, 
+                                          nms_threshold=nms_threshold)
+        detections.append(dsfd_detections)
+    else:
+        dsfd_detections = None
+
     
     return kornia_detections, retina_detections, retina_mobile_detections, dsfd_detections
 
@@ -43,27 +78,27 @@ def base_detect(detector, img):
     return img_vis
 
 
-def retina_detect(img):
+def retina_detect(img, confidence_threshold, nms_threshold):
     detector = face_detection.build_detector(
-    "RetinaNetResNet50", confidence_threshold=.5, nms_iou_threshold=.3)
+    "RetinaNetResNet50", confidence_threshold=confidence_threshold, nms_iou_threshold=nms_threshold)
     
     img_vis = base_detect(detector, img)
     
     return img_vis
 
 
-def retina_mobilenet_detect(img):
+def retina_mobilenet_detect(img, confidence_threshold, nms_threshold):
     detector = face_detection.build_detector(
-    "RetinaNetMobileNetV1", confidence_threshold=.5, nms_iou_threshold=.3)
+    "RetinaNetMobileNetV1", confidence_threshold=confidence_threshold, nms_iou_threshold=nms_threshold)
     
     img_vis = base_detect(detector, img)
     
     return img_vis
 
 
-def dsfd_detect(img):
+def dsfd_detect(img, confidence_threshold, nms_threshold):
     detector = face_detection.build_detector(
-    "DSFDDetector", confidence_threshold=.5, nms_iou_threshold=.3)
+    "DSFDDetector", confidence_threshold=confidence_threshold, nms_iou_threshold=nms_threshold)
     
     img_vis = base_detect(detector, img)
     
@@ -71,10 +106,9 @@ def dsfd_detect(img):
     
 
 
-def kornia_detect(img):
+def kornia_detect(img, confidence_threshold, nms_threshold):
     # select the device
     device = torch.device('cpu')
-    vis_threshold = 0.6
 
     # load the image and scale
     img_raw = scale_image(img, 400)
@@ -84,7 +118,8 @@ def kornia_detect(img):
     img = K.color.bgr_to_rgb(img.float())
 
     # create the detector and find the faces !
-    face_detection = FaceDetector().to(device)
+    face_detection = FaceDetector(confidence_threshold=confidence_threshold, 
+                                  nms_threshold=nms_threshold).to(device)
 
     with torch.no_grad():
         dets = face_detection(img)
@@ -95,8 +130,6 @@ def kornia_detect(img):
     img_vis = img_raw.copy()
 
     for b in dets:
-        if b.score < vis_threshold:
-            continue
 
         # draw face bounding box
         img_vis = cv2.rectangle(img_vis, 
@@ -107,29 +140,37 @@ def kornia_detect(img):
     
     return img_vis
     
-
 input_image = gr.components.Image()
-
 image_kornia = gr.components.Image(label="Kornia YuNet")
 image_retina = gr.components.Image(label="RetinaFace")
 image_retina_mobile = gr.components.Image(label="Retina Mobilenet")
 image_dsfd = gr.components.Image(label="DSFD")
 
-
-confidence_slider = gr.components.Slider(minimum=0.1, maximum=0.9, value=0.5, label="Confidence Threshold")
-nms_slider = gr.components.Slider(minimum=0.1, maximum=0.9, value=0.5, label="Min Number of Neighbours")
-# scale_slider = gr.components.Slider(minimum=1.1, maximum=2.0, value=1.3, step=0.1, label="Scale Factor")
-# classifier_radio = gr.components.Radio(s)
+confidence_slider = gr.components.Slider(minimum=0.1, maximum=0.95, value=0.5, step=0.05, label="Confidence Threshold")
+nms_slider = gr.components.Slider(minimum=0.1, maximum=0.95, value=0.3, step=0.05, label="Non Maximum Supression (NMS) Threshold")
+    
+    
+kornia_radio = gr.Radio(["On", "Off"], value="On", label="Kornia YuNet")
+retinanet_radio = gr.Radio(["On", "Off"], value="On", label="RetinaFace")
+retina_mobile_radio = gr.Radio(["On", "Off"], value="On", label="Retina Mobilenets")
+dsfd_radio = gr.Radio(["On", "Off"], value="On", label="DSFD")
 
 #methods_dropdown = gr.components.Dropdown(["Kornia YuNet", "RetinaFace", "RetinaMobile", "DSFD"], value="Kornia YuNet", label="Choose a method")
 
-description = """Face Detection"""
+description = """This space let's you compare different face detection algorithms, based on Convolutional Neural Networks (CNNs).
 
+The models used here are:
+* Kornia YuNet:
+* RetinaFace:
+* RetinaMobileNet:
+* DSFD:
+"""
 
-Iface = gr.Interface(
-    fn=detect_faces,
-    inputs=[input_image],#, size_slider, neighbour_slider, scale_slider],
+compare_iface = gr.Interface(
+    fn=compare_detect_faces,
+    inputs=[input_image, confidence_slider, nms_slider, kornia_radio, retinanet_radio, retina_mobile_radio, dsfd_radio],#, size_slider, neighbour_slider, scale_slider],
     outputs=[image_kornia, image_retina, image_retina_mobile, image_dsfd],
     examples=[["data/9_Press_Conference_Press_Conference_9_86.jpg"], ["data/12_Group_Group_12_Group_Group_12_39.jpg"], ["data/31_Waiter_Waitress_Waiter_Waitress_31_55.jpg"]],
-    title="Face Detection",
+    title="Face Detections",
+    description=description
 ).launch()