Update app.py

app.py

@@ -43,8 +43,7 @@ class DetectionResult:
                 ymin=detection_dict['box']['ymin'],
                 xmax=detection_dict['box']['xmax'],
                 ymax=detection_dict['box']['ymax']
-            ),
-            mask=detection_dict.get('mask')
+            )
         )
 
 def annotate(image: Union[Image.Image, np.ndarray], detection_results: List[DetectionResult]) -> np.ndarray:
@@ -52,12 +51,20 @@ def annotate(image: Union[Image.Image, np.ndarray], detection_results: List[Dete
     image_cv2 = cv2.cvtColor(image_cv2, cv2.COLOR_RGB2BGR)
 
     for detection in detection_results:
+        label = detection.label
+        score = detection.score
+        box = detection.box
         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 mask is not None:
             mask_uint8 = (mask * 255).astype(np.uint8)
             contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-            # Drawing only the mask contours without any bounding box or label
-            cv2.drawContours(image_cv2, contours, -1, (0, 0, 0), 2)  # Black color for contours
+            cv2.drawContours(image_cv2, contours, -1, color, 2)
 
     return cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB)
 
@@ -188,25 +195,25 @@ 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)
-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 yellow_background_with_insects, json.dumps(detections_json, separators=(’,’, ‘:’))
-else:
-return yellow_background_with_insects, None
+    original_image, detections = grounded_segmentation(image, labels, threshold=0.3, polygon_refinement=True)
+    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 yellow_background_with_insects, json.dumps(detections_json, separators=(',', ':'))
+    else:
+        return yellow_background_with_insects, None
 
 examples = [
-[“flower-night.jpg”]
+    ["flower-night.jpg"]
 ]
 
 gr.Interface(
-fn=process_image,
-inputs=[gr.Image(type=“pil”), gr.Checkbox(label=“Include JSON”, value=False)],
-outputs=[gr.Image(type=“numpy”), gr.Textbox()],
-title=“InsectSAM 🐞”,
-examples=examples
+    fn=process_image,
+    inputs=[gr.Image(type="pil"), gr.Checkbox(label="Include JSON", value=False)],
+    outputs=[gr.Image(type="numpy"), gr.Textbox()],
+    title="InsectSAM 🐞",
+    examples=examples
 ).launch()