test outputs

app.py

@@ -151,22 +151,50 @@ def extract_and_paste_insect(original_image: np.ndarray, detection: DetectionRes
     background[y_offset:y_end, x_offset:x_end] = combined
 
 def create_yellow_background_with_insects(image: np.ndarray, detections: List[DetectionResult]) -> np.ndarray:
-    yellow_background = np.full_like(image, (0, 255, 255), dtype=np.uint8)
+    yellow_background = np.full((image.shape[0], image.shape[1], 3), (0, 255, 255), dtype=np.uint8)
     for detection in detections:
         if detection.mask is not None:
             extract_and_paste_insect(image, detection, yellow_background)
     return yellow_background
 
+def draw_classification_boxes(image_with_insects, detections):
+    for detection in detections:
+        label = detection.label
+        score = detection.score
+        box = detection.box
+        color = np.random.randint(0, 256, size=3).tolist()
+
+        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
+
 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)
-    return annotated_image, yellow_background_with_insects
+    yellow_background_with_boxes = draw_classification_boxes(yellow_background_with_insects.copy(), detections)
+    return annotated_image, yellow_background_with_boxes
 
 gr.Interface(
     fn=process_image,
     inputs=gr.Image(type="pil"),
     outputs=[gr.Image(type="numpy"), gr.Image(type="numpy")],
     title="🐞 Insect Detection and Masking"
-).launch()
+).launch()