Update app.py

app.py

@@ -87,6 +87,13 @@ def get_boxes(detection_results: List[DetectionResult]) -> List[List[List[float]
         boxes.append(xyxy)
     return [boxes]
 
+def mask_to_polygon(mask: np.ndarray) -> np.ndarray:
+    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if len(contours) == 0:
+        return np.array([])
+    largest_contour = max(contours, key=cv2.contourArea)
+    return largest_contour
+
 def refine_masks(masks: torch.BoolTensor, polygon_refinement: bool = False) -> List[np.ndarray]:
     masks = masks.cpu().float().permute(0, 2, 3, 1).mean(axis=-1).numpy().astype(np.uint8)
     masks = (masks > 0).astype(np.uint8)
@@ -94,7 +101,7 @@ def refine_masks(masks: torch.BoolTensor, polygon_refinement: bool = False) -> L
         for idx, mask in enumerate(masks):
             shape = mask.shape
             polygon = mask_to_polygon(mask)
-            masks[idx] = polygon_to_mask(polygon, shape)
+            masks[idx] = cv2.fillPoly(np.zeros(shape, dtype=np.uint8), [polygon], 1)
     return list(masks)
 
 @spaces.GPU