bbox experiment

app.py

@@ -46,7 +46,7 @@ class DetectionResult:
             )
         )
 
-def annotate(image: Union[Image.Image, np.ndarray], detection_results: List[DetectionResult]) -> np.ndarray:
+def annotate(image: Union[Image.Image, np.ndarray], detection_results: List[DetectionResult], include_bboxes: bool = True) -> np.ndarray:
     image_cv2 = np.array(image) if isinstance(image, Image.Image) else image
     image_cv2 = cv2.cvtColor(image_cv2, cv2.COLOR_RGB2BGR)
 
@@ -57,9 +57,10 @@ def annotate(image: Union[Image.Image, np.ndarray], detection_results: List[Dete
         mask = detection.mask
         color = np.random.randint(0, 256, size=3).tolist()
 
-        cv2.rectangle(image_cv2, (box.xmin, box.ymin), (box.xmax, box.ymax), color, 2)
-        cv2.putText(image_cv2, f'{label}: {score:.2f}', (box.xmin, box.ymin - 10),
-                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+        if include_bboxes:
+            cv2.rectangle(image_cv2, (box.xmin, box.ymin), (box.xmax, box.ymax), color, 2)
+            cv2.putText(image_cv2, f'{label}: {score:.2f}', (box.xmin, box.ymin - 10),
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
 
         if mask is not None:
             mask_uint8 = (mask * 255).astype(np.uint8)
@@ -68,8 +69,8 @@ def annotate(image: Union[Image.Image, np.ndarray], detection_results: List[Dete
 
     return cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB)
 
-def plot_detections(image: Union[Image.Image, np.ndarray], detections: List[DetectionResult]) -> np.ndarray:
-    annotated_image = annotate(image, detections)
+def plot_detections(image: Union[Image.Image, np.ndarray], detections: List[DetectionResult], include_bboxes: bool = True) -> np.ndarray:
+    annotated_image = annotate(image, detections, include_bboxes)
     return annotated_image
 
 def load_image(image: Union[str, Image.Image]) -> Image.Image:
@@ -196,18 +197,16 @@ def detections_to_json(detections):
 def process_image(image, include_json, include_bboxes):
     labels = ["insect"]
     original_image, detections = grounded_segmentation(image, labels, threshold=0.3, polygon_refinement=True)
-    if not include_bboxes:
-        for detection in detections:
-            detection.box = None
     yellow_background_with_insects = create_yellow_background_with_insects(np.array(original_image), detections)
+    annotated_image = plot_detections(yellow_background_with_insects, detections, include_bboxes)
     if include_json:
         detections_json = detections_to_json(detections)
         json_output_path = "insect_detections.json"
         with open(json_output_path, 'w') as json_file:
             json.dump(detections_json, json_file, indent=4)
-        return yellow_background_with_insects, json.dumps(detections_json, separators=(',', ':'))
+        return annotated_image, json.dumps(detections_json, separators=(',', ':'))
     else:
-        return yellow_background_with_insects, None
+        return annotated_image, None
 
 examples = [
     ["flower-night.jpg"]