json output attempt

app.py

@@ -1,5 +1,5 @@
 import os
-os.system('pip install gradio==4.29.0')  # as gradio==4.29.0 doesn't work in requirements.txt
+os.system('pip install gradio==4.29.0')
 
 import random
 from dataclasses import dataclass
@@ -13,6 +13,7 @@ import matplotlib.pyplot as plt
 from transformers import AutoModelForMaskGeneration, AutoProcessor, pipeline
 import gradio as gr
 import spaces
+import json
 
 @dataclass
 class BoundingBox:
@@ -142,13 +143,11 @@ def extract_and_paste_insect(original_image: np.ndarray, detection: DetectionRes
     insect_crop = original_image[ymin:ymax, xmin:xmax]
     mask_crop = mask[ymin:ymax, xmin:xmax]
 
-    # Ensure that we keep the original colors of the insect
     insect = cv2.bitwise_and(insect_crop, insect_crop, mask=mask_crop)
 
     x_offset, y_offset = xmin, ymin
     x_end, y_end = x_offset + insect.shape[1], y_offset + insect.shape[0]
 
-    # Place the insect onto the yellow background
     background[y_offset:y_end, x_offset:x_end] = insect
 
 def create_yellow_background_with_insects(image: np.ndarray, detections: List[DetectionResult]) -> np.ndarray:
@@ -158,44 +157,54 @@ def create_yellow_background_with_insects(image: np.ndarray, detections: List[De
             extract_and_paste_insect(image, detection, yellow_background)
     return yellow_background
 
-def draw_classification_boxes(image_with_insects, detections):
+def run_length_encoding(mask):
+    pixels = mask.flatten()
+    rle = []
+    last_val = 0
+    count = 0
+    for pixel in pixels:
+        if pixel == last_val:
+            count += 1
+        else:
+            if count > 0:
+                rle.append(count)
+            count = 1
+            last_val = pixel
+    if count > 0:
+        rle.append(count)
+    return rle
+
+def detections_to_json(detections):
+    detections_list = []
     for detection in detections:
-        label = detection.label
-        score = detection.score
-        box = detection.box
-        color = (0, 255, 255)  # Yellow color for bounding box
-
-        cv2.rectangle(image_with_insects, (box.xmin, box.ymin), (box.xmax, box.ymax), color, 2)
-        (text_width, text_height), baseline = cv2.getTextSize(f"{label}: {score:.2f}", cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)
-        cv2.rectangle(
-            image_with_insects,
-            (box.xmin, box.ymin - text_height - baseline),
-            (box.xmin + text_width, box.ymin),
-            color,
-            thickness=cv2.FILLED
-        )
-        cv2.putText(
-            image_with_insects,
-            f"{label}: {score:.2f}",
-            (box.xmin, box.ymin - baseline),
-            cv2.FONT_HERSHEY_SIMPLEX,
-            0.5,
-            (255, 255, 255),
-            2
-        )
-    return image_with_insects
+        detection_dict = {
+            "score": detection.score,
+            "label": detection.label,
+            "box": {
+                "xmin": detection.box.xmin,
+                "ymin": detection.box.ymin,
+                "xmax": detection.box.xmax,
+                "ymax": detection.box.ymax
+            },
+            "mask": run_length_encoding(detection.mask) if detection.mask is not None else None
+        }
+        detections_list.append(detection_dict)
+    return detections_list
 
 def process_image(image):
     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)
-    yellow_background_with_boxes = draw_classification_boxes(yellow_background_with_insects.copy(), detections)
-    return annotated_image, yellow_background_with_boxes
+    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, json.dumps(detections_json, separators=(',', ':'))
 
 gr.Interface(
     fn=process_image,
     inputs=gr.Image(type="pil"),
-    outputs=[gr.Image(type="numpy"), gr.Image(type="numpy")],
+    outputs=[gr.Image(type="numpy"), gr.Textbox()],
     title="🐞 InsectSAM + GroundingDINO Inference",
-).launch()
+).launch()