Merge pull request #6 from robmarkcole/add-boxes

Add boxes

.gitignore

@@ -18,4 +18,5 @@ bin-release/
 # information for Eclipse / Flash Builder.
 
 venv
-.vscode
+.vscode
+__pycache__*

const.py

@@ -0,0 +1,8 @@
+BOX = "box"
+FILE = "file"
+OBJECT = "object"
+
+# rgb(red, green, blue)
+RED = (255, 0, 0)  # For objects within the ROI
+GREEN = (0, 255, 0)  # For ROI box
+YELLOW = (255, 255, 0)  # For objects outside the ROI

streamlit-ui.py

@@ -2,43 +2,41 @@
 import sys
 import io
 import streamlit as st
-from PIL import Image
+from PIL import Image, ImageDraw
 import numpy as np
 import os
 import deepstack.core as ds
+import utils
+import const
 
 ## Depstack setup
 DEEPSTACK_IP_ADDRESS = "localhost"
 DEEPSTACK_PORT = "5000"
 DEEPSTACK_API_KEY = ""
 DEEPSTACK_TIMEOUT = 20  # Default is 10
 
-DEFAULT_CONFIDENCE_THRESHOLD = 0.5
+DEFAULT_CONFIDENCE_THRESHOLD = 80
 TEST_IMAGE = "street.jpg"
 
-detections = ""
-
-
-def pil_image_to_byte_array(image):
-    imgByteArr = io.BytesIO()
-    image.save(imgByteArr, "PNG")
-    return imgByteArr.getvalue()
+predictions = None
 
 
 @st.cache
 def process_image(pil_image, dsobject):
     try:
-        image_bytes = pil_image_to_byte_array(pil_image)
+        image_bytes = utils.pil_image_to_byte_array(pil_image)
         dsobject.detect(image_bytes)
-        return dsobject.predictions
+        predictions = dsobject.predictions
+        summary = ds.get_objects_summary(dsobject.predictions)
+        return predictions, summary
     except Exception as exc:
         return exc
 
 
 st.title("Object detection with Deepstack")
 img_file_buffer = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
 confidence_threshold = st.slider(
-    "Confidence threshold", 0.0, 1.0, DEFAULT_CONFIDENCE_THRESHOLD, 0.05
+    "Confidence threshold", 0, 100, DEFAULT_CONFIDENCE_THRESHOLD, 1
 )
 
 if img_file_buffer is not None:
@@ -50,10 +48,32 @@ def process_image(pil_image, dsobject):
 dsobject = ds.DeepstackObject(
     DEEPSTACK_IP_ADDRESS, DEEPSTACK_PORT, DEEPSTACK_API_KEY, DEEPSTACK_TIMEOUT
 )
-detections = process_image(pil_image, dsobject)
+
+predictions, summary = process_image(pil_image, dsobject)
+objects = utils.get_objects(predictions, pil_image.width, pil_image.height)
+
+
+draw = ImageDraw.Draw(pil_image)
+for obj in objects:
+    name = obj["name"]
+    confidence = obj["confidence"]
+    box = obj["bounding_box"]
+    box_label = f"{name}: {confidence:.1f}%"
+
+    if confidence < confidence_threshold:
+        continue
+
+    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,
 )
-
-st.write(detections)
+st.write(summary)
+st.write(predictions)

utils.py

@@ -0,0 +1,101 @@
+import io
+from collections import namedtuple
+from PIL import Image, ImageDraw
+from typing import Tuple
+
+
+def pil_image_to_byte_array(image):
+    imgByteArr = io.BytesIO()
+    image.save(imgByteArr, "PNG")
+    return imgByteArr.getvalue()
+
+
+Box = namedtuple("Box", "y_min x_min y_max x_max")
+Point = namedtuple("Point", "y x")
+
+
+def point_in_box(box: Box, point: Point) -> bool:
+    """Return true if point lies in box"""
+    if (box.x_min <= point.x <= box.x_max) and (box.y_min <= point.y <= box.y_max):
+        return True
+    return False
+
+
+def object_in_roi(roi: dict, centroid: dict) -> bool:
+    """Convenience to convert dicts to the Point and Box."""
+    target_center_point = Point(centroid["y"], centroid["x"])
+    roi_box = Box(roi["y_min"], roi["x_min"], roi["y_max"], roi["x_max"])
+    return point_in_box(roi_box, target_center_point)
+
+
+def get_objects(predictions: list, img_width: int, img_height: int):
+    """Return objects with formatting and extra info."""
+    objects = []
+    decimal_places = 3
+    for pred in predictions:
+        box_width = pred["x_max"] - pred["x_min"]
+        box_height = pred["y_max"] - pred["y_min"]
+        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),
+        }
+        box_area = round(box["height"] * box["width"], decimal_places)
+        centroid = {
+            "x": round(box["x_min"] + (box["width"] / 2), decimal_places),
+            "y": round(box["y_min"] + (box["height"] / 2), decimal_places),
+        }
+        name = pred["label"]
+        confidence = round(pred["confidence"] * 100, decimal_places)
+
+        objects.append(
+            {
+                "bounding_box": box,
+                "box_area": box_area,
+                "centroid": centroid,
+                "name": name,
+                "confidence": confidence,
+            }
+        )
+    return objects
+
+
+def draw_box(
+    draw: ImageDraw,
+    box: Tuple[float, float, float, float],
+    img_width: int,
+    img_height: int,
+    text: str = "",
+    color: Tuple[int, int, int] = (255, 255, 0),
+) -> None:
+    """
+    Draw a bounding box on and image.
+    The bounding box is defined by the tuple (y_min, x_min, y_max, x_max)
+    where the coordinates are floats in the range [0.0, 1.0] and
+    relative to the width and height of the image.
+    For example, if an image is 100 x 200 pixels (height x width) and the bounding
+    box is `(0.1, 0.2, 0.5, 0.9)`, the upper-left and bottom-right coordinates of
+    the bounding box will be `(40, 10)` to `(180, 50)` (in (x,y) coordinates).
+    """
+
+    line_width = 3
+    font_height = 8
+    y_min, x_min, y_max, x_max = box
+    (left, right, top, bottom) = (
+        x_min * img_width,
+        x_max * img_width,
+        y_min * img_height,
+        y_max * img_height,
+    )
+    draw.line(
+        [(left, top), (left, bottom), (right, bottom), (right, top), (left, top)],
+        width=line_width,
+        fill=color,
+    )
+    if text:
+        draw.text(
+            (left + line_width, abs(top - line_width - font_height)), text, fill=color
+        )