bbox experiment with cv2

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], color=(0, 255, 255)) -> np.ndarray:
     image_cv2 = np.array(image) if isinstance(image, Image.Image) else image
     image_cv2 = cv2.cvtColor(image_cv2, cv2.COLOR_RGB2BGR)
 
@@ -55,7 +55,6 @@ def annotate(image: Union[Image.Image, np.ndarray], detection_results: List[Dete
         score = detection.score
         box = detection.box
         mask = detection.mask
-        color = (0, 255, 255)  # Yellow in BGR format
 
         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),
@@ -155,9 +154,9 @@ def create_yellow_background_with_insects(image: np.ndarray, detections: List[De
     for detection in detections:
         if detection.mask is not None:
             extract_and_paste_insect(image, detection, yellow_background)
-    # Convert back to RGB to match Gradio's expected input format
-    yellow_background = cv2.cvtColor(yellow_background, cv2.COLOR_BGR2RGB)
-    return yellow_background
+    # Draw bounding boxes and labels on yellow background
+    annotated_background = annotate(yellow_background, detections, color=(0, 255, 255))
+    return annotated_background
 
 def run_length_encoding(mask):
     pixels = mask.flatten()
@@ -196,16 +195,15 @@ def detections_to_json(detections):
 def process_image(image, include_json):
     labels = ["insect"]
     original_image, detections = grounded_segmentation(image, labels, threshold=0.3, polygon_refinement=True)
-    annotated_image = plot_detections(original_image, detections)
     yellow_background_with_insects = create_yellow_background_with_insects(np.array(original_image), detections)
     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 annotated_image, yellow_background_with_insects, json.dumps(detections_json, separators=(',', ':'))
+        return yellow_background_with_insects, json.dumps(detections_json, separators=(',', ':'))
     else:
-        return annotated_image, yellow_background_with_insects, None
+        return yellow_background_with_insects, None
 
 examples = [
     ["flower-night.jpg"]
@@ -214,7 +212,7 @@ examples = [
 gr.Interface(
     fn=process_image,
     inputs=[gr.Image(type="pil"), gr.Checkbox(label="Include JSON", value=False)],
-    outputs=[gr.Image(type="numpy"), gr.Image(type="numpy"), gr.Textbox()],
+    outputs=[gr.Image(type="numpy"), gr.Textbox()],
     title="InsectSAM 🐞",
     examples=examples
 ).launch()