Update app.py

app.py

@@ -49,35 +49,31 @@ class DetectionResult:
 def annotate(image: Union[Image.Image, np.ndarray], detection_results: List[DetectionResult]) -> np.ndarray:
     image_cv2 = np.array(image) if isinstance(image, Image.Image) else image
     image_cv2 = cv2.cvtColor(image_cv2, cv2.COLOR_RGB2BGR)
-    
-    # Create a completely yellow background
-    yellow_background = np.full(image_cv2.shape, (0, 255, 255), dtype=np.uint8)
-    
+
     for detection in detection_results:
+        label = detection.label
+        score = detection.score
         box = detection.box
         mask = detection.mask
-        
-        # Drawing bounding box with yellow fill
-        cv2.rectangle(yellow_background, (box.xmin, box.ymin), (box.xmax, box.ymax), (0, 255, 255), cv2.FILLED)
-        
+        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 = mask.astype(bool)
-            mask_3_channel = np.stack([mask, mask, mask], axis=-1)
-            
-            # Extract insect region using mask
-            insect_region = image_cv2 * mask_3_channel
-            
-            # Overlay insect region within the bounding box on the yellow background
-            yellow_background[mask_3_channel] = insect_region[mask_3_channel]
+            mask_uint8 = (mask * 255).astype(np.uint8)
+            contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+            cv2.drawContours(image_cv2, contours, -1, color, 2)
 
-    return cv2.cvtColor(yellow_background, cv2.COLOR_BGR2RGB)
+    return cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB)
 
 def plot_detections(image: Union[Image.Image, np.ndarray], detections: List[DetectionResult]) -> np.ndarray:
     annotated_image = annotate(image, detections)
     return annotated_image
 
 def load_image(image: Union[str, Image.Image]) -> Image.Image:
-    if isinstance(image, str) and image.startsWith("http"):
+    if isinstance(image, str) and image.startswith("http"):
         image = Image.open(requests.get(image, stream=True).raw).convert("RGB")
     elif isinstance(image, str):
         image = Image.open(image).convert("RGB")
@@ -110,7 +106,7 @@ def refine_masks(masks: torch.BoolTensor, polygon_refinement: bool = False) -> L
     return list(masks)
 
 @spaces.GPU
-def detect(image: Image.Image, labels: List[str], threshold: float = 0.3, detector_id: Optional[str] = None) -> List[DetectionResult]:
+def detect(image: Image.Image, labels: List[str], threshold: float = 0.3, detector_id: Optional[str] = None) -> List[Dict[str, Any]]:
     detector_id = detector_id if detector_id else "IDEA-Research/grounding-dino-base"
     object_detector = pipeline(model=detector_id, task="zero-shot-object-detection", device="cuda")
     labels = [label if label.endswith(".") else label+"." for label in labels]
@@ -163,7 +159,7 @@ def create_yellow_background_with_insects(image: np.ndarray, detections: List[De
     yellow_background = cv2.cvtColor(yellow_background, cv2.COLOR_BGR2RGB)
     return yellow_background
 
-def run_length_encoding(mask: np.ndarray) -> List[int]:
+def run_length_encoding(mask):
     pixels = mask.flatten()
     rle = []
     last_val = 0
@@ -180,7 +176,7 @@ def run_length_encoding(mask: np.ndarray) -> List[int]:
         rle.append(count)
     return rle
 
-def detections_to_json(detections: List[DetectionResult]) -> List[Dict[str, Any]]:
+def detections_to_json(detections):
     detections_list = []
     for detection in detections:
         detection_dict = {