Merge pull request #38 from robmarkcole/add-face

add face beginnings
robmarkcole · Jan 3, 2021 · 275aa33 · 275aa33
2 parents f9b4ae4 + d88f911
commit 275aa33
Show file tree

Hide file tree

Showing 5 changed files with 195 additions and 108 deletions.
diff --git a/README.md b/README.md
@@ -1,5 +1,5 @@
 # deepstack-ui
-UI for working with [Deepstack](https://python.deepstack.cc/). Allows uploading an image and performing object detection with Deepstack. The effect of various parameters can be explored, including filtering the classes of object detected, filtering by minimum confidence, and spatial filtering using a region of interest (ROI).
+UI for working with [Deepstack](https://python.deepstack.cc/). Allows uploading an image and performing object detection or face recognition with Deepstack. The effect of various parameters can be explored, including filtering the classes of object detected, filtering by minimum confidence, and spatial filtering using a region of interest (ROI).
 
 <p align="center">
 <img src="https://github.com/robmarkcole/deepstack-ui/blob/master/usage.png" width="1000">

diff --git a/app/couple.jpg b/app/couple.jpg
diff --git a/app/deepstack-ui.py b/app/deepstack-ui.py
@@ -9,8 +9,13 @@
 import utils
 import const
 
-DEFAULT_CONFIDENCE_THRESHOLD = 0.01
-TEST_IMAGE = "street.jpg"
+DEFAULT_CONFIDENCE_THRESHOLD = 0.45
+MIN_CONFIDENCE_THRESHOLD = 0.1
+MAX_CONFIDENCE_THRESHOLD = 1.0
+OBJECT_TEST_IMAGE = "street.jpg"
+FACE_TEST_IMAGE = "faces.jpg"
+FACE = "Face"
+OBJECT = "Object"
 
 DEFAULT_ROI_Y_MIN = 0.0
 DEFAULT_ROI_Y_MAX = 1.0
@@ -34,128 +39,186 @@
 
 
 @st.cache
-def process_image(pil_image, dsobject):
+def process_image_object(pil_image, dsobject):
     image_bytes = utils.pil_image_to_byte_array(pil_image)
     predictions = dsobject.detect(image_bytes)
     return predictions
 
 
-## Setup sidebar
-st.title("Deepstack Object detection")
-if not DEEPSTACK_CUSTOM_MODEL:
-    st.text("Using default model")
-else:
-    st.text(f"Using custom model named {DEEPSTACK_CUSTOM_MODEL}")
+@st.cache
+def process_image_face(pil_image, dsface):
+    image_bytes = utils.pil_image_to_byte_array(pil_image)
+    predictions = dsface.recognize(image_bytes)
+    return predictions
 
-img_file_buffer = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
 
-st.sidebar.title("Parameters")
-st.text("Adjust parameters to select what is displayed")
-CONFIDENCE_THRESHOLD = st.sidebar.slider(
-    "Confidence threshold", DEFAULT_CONFIDENCE_THRESHOLD, 1.0
-)
+deepstack_mode = st.selectbox("Select Deepstack mode:", [OBJECT, FACE])
 
-if not DEEPSTACK_CUSTOM_MODEL:
-    CLASSES_TO_INCLUDE = st.sidebar.multiselect(
-        "Select object classes to include",
-        options=const.CLASSES,
-        default=const.CLASSES,
-    )
+if deepstack_mode == FACE:
+    st.title("Deepstack Face recogntion")
 
-# Get ROI info
-st.sidebar.title("ROI")
-ROI_X_MIN = st.sidebar.slider("x_min", 0.0, 1.0, DEFAULT_ROI_X_MIN)
-ROI_Y_MIN = st.sidebar.slider("y_min", 0.0, 1.0, DEFAULT_ROI_Y_MIN)
-ROI_X_MAX = st.sidebar.slider("x_max", 0.0, 1.0, DEFAULT_ROI_X_MAX)
-ROI_Y_MAX = st.sidebar.slider("y_max", 0.0, 1.0, DEFAULT_ROI_Y_MAX)
-ROI_TUPLE = (
-    ROI_Y_MIN,
-    ROI_X_MIN,
-    ROI_Y_MAX,
-    ROI_X_MAX,
-)
-ROI_DICT = {
-    "x_min": ROI_X_MIN,
-    "y_min": ROI_Y_MIN,
-    "x_max": ROI_X_MAX,
-    "y_max": ROI_Y_MAX,
-}
-
-## Process image
-if img_file_buffer is not None:
-    pil_image = Image.open(img_file_buffer)
-
-else:
-    pil_image = Image.open(TEST_IMAGE)
-
-if not DEEPSTACK_CUSTOM_MODEL:
-    dsobject = ds.DeepstackObject(
+    img_file_buffer = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
+    ## Process image
+    if img_file_buffer is not None:
+        pil_image = Image.open(img_file_buffer)
+
+    else:
+        pil_image = Image.open(FACE_TEST_IMAGE)
+
+    dsface = ds.DeepstackFace(
         ip=DEEPSTACK_IP,
         port=DEEPSTACK_PORT,
         api_key=DEEPSTACK_API_KEY,
         timeout=DEEPSTACK_TIMEOUT,
-        min_confidence=DEFAULT_CONFIDENCE_THRESHOLD,
+        min_confidence=MIN_CONFIDENCE_THRESHOLD,
     )
-else:
-    dsobject = ds.DeepstackObject(
-        ip=DEEPSTACK_IP,
-        port=DEEPSTACK_PORT,
-        api_key=DEEPSTACK_API_KEY,
-        timeout=DEEPSTACK_TIMEOUT,
-        min_confidence=DEFAULT_CONFIDENCE_THRESHOLD,
-        custom_model=DEEPSTACK_CUSTOM_MODEL,
+    predictions = process_image_face(pil_image, dsface)
+    faces = utils.get_faces(predictions, pil_image.width, pil_image.height)
+
+    # Draw object boxes
+    draw = ImageDraw.Draw(pil_image)
+    for face in faces:
+        name = face["name"]
+        confidence = face["confidence"]
+        box = face["bounding_box"]
+        box_label = f"{name}"
+
+        utils.draw_box(
+            draw,
+            (box["y_min"], box["x_min"], box["y_max"], box["x_max"]),
+            pil_image.width,
+            pil_image.height,
+            text=box_label,
+            color=const.YELLOW,
+        )
+    st.image(
+        np.array(pil_image), caption=f"Processed image", use_column_width=True,
     )
-
-predictions = process_image(pil_image, dsobject)
-objects = utils.get_objects(predictions, pil_image.width, pil_image.height)
-all_objects_names = set([obj["name"] for obj in objects])
-
-# Filter objects for display
-objects = [obj for obj in objects if obj["confidence"] > CONFIDENCE_THRESHOLD]
-objects = [obj for obj in objects if utils.object_in_roi(ROI_DICT, obj["centroid"])]
-if not DEEPSTACK_CUSTOM_MODEL:
-    objects = [obj for obj in objects if obj["name"] in CLASSES_TO_INCLUDE]
-
-# Draw object boxes
-draw = ImageDraw.Draw(pil_image)
-for obj in objects:
-    name = obj["name"]
-    confidence = obj["confidence"]
-    box = obj["bounding_box"]
-    box_label = f"{name}"
-
-    utils.draw_box(
-        draw,
-        (box["y_min"], box["x_min"], box["y_max"], box["x_max"]),
-        pil_image.width,
-        pil_image.height,
-        text=box_label,
-        color=const.YELLOW,
+    st.subheader("All faces")
+    st.write(faces)
+
+
+elif deepstack_mode == OBJECT:
+    ## Setup main
+    st.title("Deepstack Object detection")
+    if not DEEPSTACK_CUSTOM_MODEL:
+        st.text("Using default model")
+    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,
     )
 
-# Draw ROI box
-if ROI_TUPLE != DEFAULT_ROI:
-    utils.draw_box(
-        draw,
-        ROI_TUPLE,
-        pil_image.width,
-        pil_image.height,
-        text="ROI",
-        color=const.GREEN,
+    # Get ROI info
+    st.sidebar.title("ROI")
+    ROI_X_MIN = st.sidebar.slider("x_min", 0.0, 1.0, DEFAULT_ROI_X_MIN)
+    ROI_Y_MIN = st.sidebar.slider("y_min", 0.0, 1.0, DEFAULT_ROI_Y_MIN)
+    ROI_X_MAX = st.sidebar.slider("x_max", 0.0, 1.0, DEFAULT_ROI_X_MAX)
+    ROI_Y_MAX = st.sidebar.slider("y_max", 0.0, 1.0, DEFAULT_ROI_Y_MAX)
+    ROI_TUPLE = (
+        ROI_Y_MIN,
+        ROI_X_MIN,
+        ROI_Y_MAX,
+        ROI_X_MAX,
     )
+    ROI_DICT = {
+        "x_min": ROI_X_MIN,
+        "y_min": ROI_Y_MIN,
+        "x_max": ROI_X_MAX,
+        "y_max": ROI_Y_MAX,
+    }
+
+    img_file_buffer = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
+
+    if not DEEPSTACK_CUSTOM_MODEL:
+        CLASSES_TO_INCLUDE = st.sidebar.multiselect(
+            "Select object classes to include",
+            options=const.CLASSES,
+            default=const.CLASSES,
+        )
+
+    ## Process image
+    if img_file_buffer is not None:
+        pil_image = Image.open(img_file_buffer)
+
+    else:
+        pil_image = Image.open(OBJECT_TEST_IMAGE)
+
+    if not DEEPSTACK_CUSTOM_MODEL:
+        dsobject = ds.DeepstackObject(
+            ip=DEEPSTACK_IP,
+            port=DEEPSTACK_PORT,
+            api_key=DEEPSTACK_API_KEY,
+            timeout=DEEPSTACK_TIMEOUT,
+            min_confidence=MIN_CONFIDENCE_THRESHOLD,
+        )
+    else:
+        dsobject = ds.DeepstackObject(
+            ip=DEEPSTACK_IP,
+            port=DEEPSTACK_PORT,
+            api_key=DEEPSTACK_API_KEY,
+            timeout=DEEPSTACK_TIMEOUT,
+            min_confidence=MIN_CONFIDENCE_THRESHOLD,
+            custom_model=DEEPSTACK_CUSTOM_MODEL,
+        )
+
+    predictions = process_image_object(pil_image, dsobject)
+    objects = utils.get_objects(predictions, pil_image.width, pil_image.height)
+    all_objects_names = set([obj["name"] for obj in objects])
+
+    # Filter objects for display
+    objects = [obj for obj in objects if obj["confidence"] > CONFIDENCE_THRESHOLD]
+    objects = [obj for obj in objects if utils.object_in_roi(ROI_DICT, obj["centroid"])]
+    if not DEEPSTACK_CUSTOM_MODEL:
+        objects = [obj for obj in objects if obj["name"] in CLASSES_TO_INCLUDE]
+
+    # Draw object boxes
+    draw = ImageDraw.Draw(pil_image)
+    for obj in objects:
+        name = obj["name"]
+        confidence = obj["confidence"]
+        box = obj["bounding_box"]
+        box_label = f"{name}"
+
+        utils.draw_box(
+            draw,
+            (box["y_min"], box["x_min"], box["y_max"], box["x_max"]),
+            pil_image.width,
+            pil_image.height,
+            text=box_label,
+            color=const.YELLOW,
+        )
+
+    # Draw ROI box
+    if ROI_TUPLE != DEFAULT_ROI:
+        utils.draw_box(
+            draw,
+            ROI_TUPLE,
+            pil_image.width,
+            pil_image.height,
+            text="ROI",
+            color=const.GREEN,
+        )
+
+    # Display image and results
+    st.image(
+        np.array(pil_image), caption=f"Processed image", use_column_width=True,
+    )
+    st.subheader("All discovered objects")
+    st.write(all_objects_names)
 
-# Display image and results
-st.image(
-    np.array(pil_image), caption=f"Processed image", use_column_width=True,
-)
-st.subheader("All discovered objects")
-st.write(all_objects_names)
-
-st.subheader("Filtered object count")
-obj_types = list(set([obj["name"] for obj in objects]))
-for obj_type in obj_types:
-    obj_type_count = len([obj for obj in objects if obj["name"] == obj_type])
-    st.write(f"{obj_type} : {obj_type_count}")
+    st.subheader("Filtered object count")
+    obj_types = list(set([obj["name"] for obj in objects]))
+    for obj_type in obj_types:
+        obj_type_count = len([obj for obj in objects if obj["name"] == obj_type])
+        st.write(f"{obj_type} : {obj_type_count}")
 
-st.subheader("All filtered objects")
-st.write(objects)
+    st.subheader("All filtered objects")
+    st.write(objects)
diff --git a/app/idris.jpg b/app/idris.jpg
diff --git a/app/utils.py b/app/utils.py
@@ -63,6 +63,30 @@ def get_objects(predictions: list, img_width: int, img_height: int):
     return objects
 
 
+def get_faces(predictions: list, img_width: int, img_height: int):
+    """Return faces info."""
+    faces = []
+    decimal_places = 3
+    for pred in predictions:
+        box_width = pred["x_max"] - pred["x_min"]
+        box_height = pred["y_max"] - pred["y_min"]
+        name = pred["userid"]
+        confidence = pred["confidence"]
+        box = {
+            "height": round(box_height / img_height, decimal_places),
+            "width": round(box_width / img_width, decimal_places),
+            "y_min": round(pred["y_min"] / img_height, decimal_places),
+            "x_min": round(pred["x_min"] / img_width, decimal_places),
+            "y_max": round(pred["y_max"] / img_height, decimal_places),
+            "x_max": round(pred["x_max"] / img_width, decimal_places),
+        }
+
+        faces.append(
+            {"name": name, "confidence": confidence, "bounding_box": box,}
+        )
+    return faces
+
+
 def draw_box(
     draw: ImageDraw,
     box: Tuple[float, float, float, float],