From af900e7e1ae181e1459c0731f9d55ee57d33d98c Mon Sep 17 00:00:00 2001
From: Robin Cole <robmarkcole@gmail.com>
Date: Sun, 3 Jan 2021 13:57:50 +0000
Subject: [PATCH] Update deepstack-ui.py

---
 app/deepstack-ui.py | 36 +++++++++++++++++++++++-------------
 1 file changed, 23 insertions(+), 13 deletions(-)

diff --git a/app/deepstack-ui.py b/app/deepstack-ui.py
index 0166a73..9e17176 100644
--- a/app/deepstack-ui.py
+++ b/app/deepstack-ui.py
@@ -53,6 +53,15 @@ def process_image_face(pil_image, dsface):
 
 
 deepstack_mode = st.selectbox("Select Deepstack mode:", [OBJECT, FACE])
+## Setup sidebar
+st.sidebar.title("Parameters")
+st.text("Adjust parameters to select what is displayed")
+CONFIDENCE_THRESHOLD = st.sidebar.slider(
+    "Confidence threshold",
+    MIN_CONFIDENCE_THRESHOLD,
+    MAX_CONFIDENCE_THRESHOLD,
+    DEFAULT_CONFIDENCE_THRESHOLD,
+)
 
 if deepstack_mode == FACE:
     st.title("Deepstack Face recogntion")
@@ -74,14 +83,21 @@ def process_image_face(pil_image, dsface):
     )
     predictions = process_image_face(pil_image, dsface)
     faces = utils.get_faces(predictions, pil_image.width, pil_image.height)
+    recognised_faces = [
+        face for face in faces if face["confidence"] > CONFIDENCE_THRESHOLD
+    ]
 
     # Draw object boxes
     draw = ImageDraw.Draw(pil_image)
     for face in faces:
-        name = face["name"]
         confidence = face["confidence"]
-        box = face["bounding_box"]
+        name = face["name"]
         box_label = f"{name}"
+        box = face["bounding_box"]
+        if confidence < CONFIDENCE_THRESHOLD or name == "unknown":
+            box_colour = const.YELLOW
+        else:
+            box_colour = const.GREEN
 
         utils.draw_box(
             draw,
@@ -89,11 +105,15 @@ def process_image_face(pil_image, dsface):
             pil_image.width,
             pil_image.height,
             text=box_label,
-            color=const.YELLOW,
+            color=box_colour,
         )
     st.image(
         np.array(pil_image), caption=f"Processed image", use_column_width=True,
     )
+
+    st.subheader("All recognised faces")
+    st.write(recognised_faces)
+
     st.subheader("All faces")
     st.write(faces)
 
@@ -123,16 +143,6 @@ def process_image_face(pil_image, dsface):
     else:
         st.text(f"Using custom model named {DEEPSTACK_CUSTOM_MODEL}")
 
-    ## Setup sidebar
-    st.sidebar.title("Parameters")
-    st.text("Adjust parameters to select what is displayed")
-    CONFIDENCE_THRESHOLD = st.sidebar.slider(
-        "Confidence threshold",
-        MIN_CONFIDENCE_THRESHOLD,
-        MAX_CONFIDENCE_THRESHOLD,
-        DEFAULT_CONFIDENCE_THRESHOLD,
-    )
-
     # Get ROI info
     st.sidebar.title("ROI")
     ROI_X_MIN = st.sidebar.slider("x_min", 0.0, 1.0, DEFAULT_ROI_X_MIN)