Update app.py

app.py

@@ -50,16 +50,30 @@ def annotate(image: Union[Image.Image, np.ndarray], detection_results: List[Dete
     image_cv2 = np.array(image) if isinstance(image, Image.Image) else image
     image_cv2 = cv2.cvtColor(image_cv2, cv2.COLOR_RGB2BGR)
     
-    # Make the entire background yellow
+    # Create an empty yellow background
     yellow_background = np.full(image_cv2.shape, (0, 255, 255), dtype=np.uint8)
     
     for detection in detection_results:
         mask = detection.mask
         
         if mask is not None:
-            mask_expanded = np.stack([mask]*3, axis=-1) # Expand mask dimensions for color channels
-            insect_region = np.where(mask_expanded, image_cv2, yellow_background)
-            yellow_background = np.where(mask_expanded, insect_region, yellow_background)
+            mask_uint8 = (mask * 255).astype(np.uint8)
+            mask_3_channel = cv2.merge([mask_uint8, mask_uint8, mask_uint8])
+            
+            # Invert the mask to get the background part
+            inverted_mask = cv2.bitwise_not(mask_uint8)
+            inverted_mask_3_channel = cv2.merge([inverted_mask, inverted_mask, inverted_mask])
+            
+            # Extract insect region using mask
+            insect_region = cv2.bitwise_and(image_cv2, mask_3_channel)
+            yellow_background_region = cv2.bitwise_and(yellow_background, inverted_mask_3_channel)
+            
+            # Combine the insect region and the yellow background region
+            combined_region = cv2.add(insect_region, yellow_background_region)
+            
+            # Apply combined region onto the yellow background
+            yellow_background = cv2.bitwise_and(yellow_background, inverted_mask_3_channel)
+            yellow_background = cv2.add(yellow_background, combined_region)
 
     return cv2.cvtColor(yellow_background, cv2.COLOR_BGR2RGB)