Better Polygons from Masks

supervision/detection/utils.py

@@ -166,9 +166,63 @@ def mask_to_xyxy(masks: np.ndarray) -> np.ndarray:
     return bboxes
 
 
+def is_clockwise(contour):
+    value = 0
+    num = len(contour)
+    for i, point in enumerate(contour):
+        p1 = contour[i]
+        if i < num - 1:
+            p2 = contour[i + 1]
+        else:
+            p2 = contour[0]
+        value += (p2[0][0] - p1[0][0]) * (p2[0][1] + p1[0][1])
+    return value < 0
+
+
+def get_merge_point_idx(contour1, contour2):
+    idx1 = 0
+    idx2 = 0
+    distance_min = -1
+    for i, p1 in enumerate(contour1):
+        for j, p2 in enumerate(contour2):
+            distance = pow(p2[0][0] - p1[0][0], 2) + pow(p2[0][1] - p1[0][1], 2)
+            if distance_min < 0:
+                distance_min = distance
+                idx1 = i
+                idx2 = j
+            elif distance < distance_min:
+                distance_min = distance
+                idx1 = i
+                idx2 = j
+    return idx1, idx2
+
+
+def merge_contours(contour1, contour2, idx1, idx2):
+    contour = []
+    for i in list(range(0, idx1 + 1)):
+        contour.append(contour1[i])
+    for i in list(range(idx2, len(contour2))):
+        contour.append(contour2[i])
+    for i in list(range(0, idx2 + 1)):
+        contour.append(contour2[i])
+    for i in list(range(idx1, len(contour1))):
+        contour.append(contour1[i])
+    contour = np.array(contour)
+    return contour
+
+
+def merge_with_parent(contour_parent, contour):
+    if not is_clockwise(contour_parent):
+        contour_parent = contour_parent[::-1]
+    if is_clockwise(contour):
+        contour = contour[::-1]
+    idx1, idx2 = get_merge_point_idx(contour_parent, contour)
+    return merge_contours(contour_parent, contour, idx1, idx2)
+
+
 def mask_to_polygons(mask: np.ndarray) -> List[np.ndarray]:
     """
-    Converts a binary mask to a list of polygons.
+    Converts a binary mask to a list of connected polygons.
 
     Parameters:
         mask (np.ndarray): A binary mask represented as a 2D NumPy array of
@@ -181,13 +235,41 @@ def mask_to_polygons(mask: np.ndarray) -> List[np.ndarray]:
             of the points. Polygons with fewer points than `MIN_POLYGON_POINT_COUNT = 3`
             are excluded from the output.
     """
-
-    contours, _ = cv2.findContours(
-        mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
+    contours, hierarchies = cv2.findContours(
+        mask.astype(np.uint8), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_TC89_KCOS
     )
+
+    contours_approx = []
+    for contour in contours:
+        epsilon = 0.001 * cv2.arcLength(contour, True)
+        contour_approx = cv2.approxPolyDP(contour, epsilon, True)
+        contours_approx.append(contour_approx)
+
+    contours_parent = []
+    for i, contour in enumerate(contours_approx):
+        parent_idx = hierarchies[0][i][3]
+        if parent_idx < 0 and len(contour) >= 3:
+            contours_parent.append(contour)
+        else:
+            contours_parent.append([])
+
+    for i, contour in enumerate(contours_approx):
+        parent_idx = hierarchies[0][i][3]
+        if parent_idx >= 0 and len(contour) >= 3:
+            contour_parent = contours_parent[parent_idx]
+            if len(contour_parent) == 0:
+                continue
+            contours_parent[parent_idx] = merge_with_parent(contour_parent, contour)
+
+    contours_parent_tmp = []
+    for contour in contours_parent:
+        if len(contour) == 0:
+            continue
+        contours_parent_tmp.append(contour)
+
     return [
         np.squeeze(contour, axis=1)
-        for contour in contours
+        for contour in contours_parent_tmp
         if contour.shape[0] >= MIN_POLYGON_POINT_COUNT
     ]
 

supervision/draw/utils.py

@@ -82,7 +82,12 @@ def draw_filled_rectangle(scene: np.ndarray, rect: Rect, color: Color) -> np.nda
 
 
 def draw_polygon(
-    scene: np.ndarray, polygon: np.ndarray, color: Color, thickness: int = 2
+    scene: np.ndarray,
+    polygon: np.ndarray,
+    color: Color,
+    thickness: int = 2,
+    fill: bool = True,
+    opacity: float = 0.5,
 ) -> np.ndarray:
     """Draw a polygon on a scene.
 
@@ -98,6 +103,10 @@ def draw_polygon(
     cv2.polylines(
         scene, [polygon], isClosed=True, color=color.as_bgr(), thickness=thickness
     )
+    if fill:
+        colored_mask = np.array(scene, copy=True, dtype=np.uint8)
+        cv2.fillPoly(colored_mask, [polygon], color=color.as_bgr())
+        scene = cv2.addWeighted(colored_mask, opacity, scene, 1 - opacity, 0)
     return scene