app.py update, change of outputs

app.py

@@ -1,5 +1,5 @@
 import os
-os.system('pip install gradio==4.29.0')
+os.system('pip install gradio==4.29.0') # as gradio==4.29.0 doesn't work in requirements.txt
 
 import random
 from dataclasses import dataclass
@@ -132,17 +132,6 @@ def grounded_segmentation(image: Union[Image.Image, str], labels: List[str], thr
     detections = segment(image, detections, polygon_refinement, segmenter_id)
     return np.array(image), detections
 
-def extract_insect_masks(image: np.ndarray, detections: List[DetectionResult]) -> List[np.ndarray]:
-    return [detection.mask for detection in detections if detection.mask is not None]
-
-def put_masks_on_yellow_background(image_shape: Tuple[int, int], masks: List[np.ndarray]) -> np.ndarray:
-    yellow_background = np.full((image_shape[0], image_shape[1], 3), (0, 255, 255), dtype=np.uint8)
-    for mask in masks:
-        mask_rgb = cv2.cvtColor(mask, cv2.COLOR_GRAY2RGB)
-        for c in range(3):
-            yellow_background[:,:,c] = cv2.bitwise_or(yellow_background[:,:,c], mask_rgb[:,:,c])
-    return yellow_background
-
 def mask_to_min_max(mask: np.ndarray) -> Tuple[int, int, int, int]:
     y, x = np.where(mask)
     return x.min(), y.min(), x.max(), y.max()
@@ -168,33 +157,16 @@ 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: np.ndarray, detections: List[DetectionResult]) -> np.ndarray:
-    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, (0, 0, 0), 2)
-    return image_with_insects
-
 def process_image(image):
     labels = ["insect"]
     original_image, detections = grounded_segmentation(image, labels, threshold=0.3, polygon_refinement=True)
-    masked_image = plot_detections(original_image, detections)
-    insect_masks = extract_insect_masks(original_image, detections)
-    yellow_background_with_masks = put_masks_on_yellow_background(original_image.shape[:2], insect_masks)
-    yellow_background_with_insects = create_yellow_background_with_insects(original_image, detections)
-    yellow_background_with_boxes = draw_classification_boxes(yellow_background_with_insects, detections)
-    return masked_image, yellow_background_with_masks, yellow_background_with_boxes
+    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
 
 gr.Interface(
     fn=process_image,
     inputs=gr.Image(type="pil"),
-    outputs=[gr.Image(type="numpy"), gr.Image(type="numpy"), gr.Image(type="numpy")],
-    title="🐞 Insect Detection and Masking"
+    outputs=[gr.Image(type="numpy"), gr.Image(type="numpy")],
+    title="Insect Detection and Masking"
 ).launch()