Contours (#37)

* Use contours to calculate bounding boxes

* Fix bug

* Add option to draw bounding boxes which is enabled by default

* Add option to docs

* Formatting

Author: NickM-27

docs/config.md

@@ -52,6 +52,8 @@ cameras:
         # OPTIONAL: but highly recommended, setting the default url for a snapshot to be
         # processed by this camera. This is required for auto detection (Default: none).
         url: "http://ip.ad.dr.ess/jpg"
+        # OPTIONAL: Whether or not to draw bounding boxes for confirmed objects in the snapshots (Default: shown below).
+        bounding_box: true
         # OPTIONAL: Whether or not to save the snapshots of confirmed detections (Default: shown below).
         save_detections: true
         # OPTIONAL: Whether or not to save the snapshots of missed detections (Default: shown below).

swatch/config.py

@@ -29,6 +29,10 @@ class SnapshotConfig(SwatchBaseModel):
 
     url: str = Field(title="Camera Snapshot Url.", default=None)
     mode: SnapshotModeEnum = Field(title="Snapshot mode.", default=SnapshotModeEnum.ALL)
+    bounding_box: bool = Field(
+        title="Write bounding boxes for detected objects on the snapshot.",
+        default=True,
+    )
     save_detections: bool = Field(
         title="Save snapshots of detections that are found.", default=True
     )

swatch/detection.py

@@ -83,10 +83,10 @@ def __handle_detections__(self, detection_result: Dict[str, Any]) -> None:
                 self.obj_data[non_unique_id]["zone_name"] = zone_name
                 self.obj_data[non_unique_id]["variant"] = object_result["variant"]
 
-                if object_result["area"] > self.obj_data[non_unique_id].get(
-                    "top_area", 0
-                ):
-                    self.obj_data[non_unique_id]["top_area"] = object_result["area"]
+                top_area = max([d["area"] for d in object_result["objects"]])
+
+                if top_area > self.obj_data[non_unique_id].get("top_area", 0):
+                    self.obj_data[non_unique_id]["top_area"] = top_area
 
                     # save snapshot with best area
                     self.snap_processor.save_snapshot(

swatch/image.py

@@ -1,7 +1,7 @@
 """ImageProcessor for getting detectable info from images."""
 
 import datetime
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, Optional, Set, Tuple
 from colorthief import ColorThief
 import cv2
 import numpy as np
@@ -42,6 +42,32 @@ def __mask_image__(crop: Any, color_variant: ColorVariantConfig) -> Tuple[Any, i
     return (output, matches)
 
 
+def __detect_objects__(mask: Any, object: ObjectConfig) -> Set[Dict[str, Any]]:
+    """Detect objects and return list of bounding boxes."""
+    # get gray image
+    gray = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
+
+    # calculate contours
+    _, thresh = cv2.threshold(gray, 1, 255, 0)
+    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+    detected = []
+
+    for contour in contours:
+        x, y, w, h = cv2.boundingRect(contour)
+        area = w * h
+
+        if object.min_area < area < object.max_area:
+            detected.append(
+                {
+                    "box": [x, y, x + w, y + h],
+                    "area": area,
+                }
+            )
+
+    return detected
+
+
 class ImageProcessor:
     """Processing images with swatch config data."""
 
@@ -76,8 +102,22 @@ def __check_image__(
                 continue
 
             output, matches = __mask_image__(crop, color_variant)
-
-            if detectable.min_area < matches < detectable.max_area:
+            detected_objects = __detect_objects__(crop, detectable)
+
+            if detected_objects:
+
+                # draw bounding boxes on image if enabled
+                if snapshot.bounding_box:
+                    for obj in detected_objects:
+                        cv2.rectangle(
+                            output,
+                            (obj["box"][0], obj["box"][1]),
+                            (obj["box"][2], obj["box"][3]),
+                            (0, 255, 0),
+                            4,
+                        )
+
+                # save the snapshot if enabled
                 if snapshot.save_detections and snapshot.mode in [
                     SnapshotModeEnum.ALL,
                     SnapshotModeEnum.MASK,
@@ -91,9 +131,9 @@ def __check_image__(
 
                 return {
                     "result": True,
-                    "area": matches,
                     "variant": variant_name,
                     "camera_name": camera_name,
+                    "objects": detected_objects,
                 }
 
             if matches > best_fail.get("area", 0):