`mask_to_polygons` and `filter_polygons_by_area` added

supervision/dataset/formats/pascal_voc.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Tuple
 from xml.dom.minidom import parseString
 from xml.etree.ElementTree import Element, SubElement, tostring
 
@@ -9,9 +9,7 @@ def detections_to_pascal_voc(
         detections: Detections,
         classes: List[str],
         filename: str,
-        width: int,
-        height: int,
-        depth: int = 3,
+        image_shape: Tuple[int, int, int]
 ) -> str:
     """
     Converts Detections object to Pascal VOC XML format.
@@ -20,12 +18,11 @@ def detections_to_pascal_voc(
         detections (Detections): A Detections object containing bounding boxes, class ids, and other relevant information.
         classes (List[str]): A list of class names corresponding to the class ids in the Detections object.
         filename (str): The name of the image file associated with the detections.
-        width (int): The width of the image in pixels.
-        height (int): The height of the image in pixels.
-        depth (int, optional): The number of color channels in the image. Defaults to 3 for RGB images.
+        image_shape (Tuple[int, int, int]): The shape of the image file associated with the detections.
     Returns:
         str: An XML string in Pascal VOC format representing the detections.
     """
+    height, width, depth = image_shape
 
     # Create root element
     annotation = Element("annotation")
@@ -35,8 +32,8 @@ def detections_to_pascal_voc(
     folder.text = "VOC"
 
     # Add filename element
-    fname = SubElement(annotation, "filename")
-    fname.text = filename
+    file_name = SubElement(annotation, "filename")
+    file_name.text = filename
 
     # Add source element
     source = SubElement(annotation, "source")

supervision/detection/utils.py

@@ -1,9 +1,12 @@
-from typing import Tuple
+from typing import Tuple, List, Optional
 
 import cv2
 import numpy as np
 
 
+MIN_POLYGON_POINT_COUNT = 3
+
+
 def generate_2d_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.ndarray:
     """Generate a 2D mask from a polygon.
 
@@ -146,3 +149,43 @@ def mask_to_xyxy(masks: np.ndarray) -> np.ndarray:
             bboxes[i, :] = [x_min, y_min, x_max, y_max]
 
     return bboxes
+
+
+def mask_to_polygons(mask: np.ndarray) -> List[np.ndarray]:
+    contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
+    return [
+        np.squeeze(contour, axis=1)
+        for contour
+        in contours
+        if contour.shape[0] >= MIN_POLYGON_POINT_COUNT
+    ]
+
+
+def filter_polygons_by_area(polygons: List[np.ndarray], min_area: Optional[float], max_area: Optional[float]) -> List[np.ndarray]:
+    """
+    Filters a list of polygons based on their area.
+
+    Parameters:
+        polygons (List[np.ndarray]): A list of polygons, where each polygon is represented by a NumPy array of shape (N, 2),
+            containing the x, y coordinates of the points.
+        min_area (Optional[float]): The minimum area threshold. Only polygons with an area greater than or equal to this value
+            will be included in the output. If set to None, no minimum area constraint will be applied.
+        max_area (Optional[float]): The maximum area threshold. Only polygons with an area less than or equal to this value
+            will be included in the output. If set to None, no maximum area constraint will be applied.
+
+    Returns:
+        List[np.ndarray]: A new list of polygons containing only those with areas within the specified thresholds.
+    """
+    if min_area is None and max_area is None:
+        return polygons
+    ares = [
+        cv2.contourArea(polygon)
+        for polygon
+        in polygons
+    ]
+    return [
+        polygon
+        for polygon, area
+        in zip(polygons, ares)
+        if (min_area is None or area>= min_area) and (max_area is None or area<= max_area)
+    ]

test/detection/test_utils.py

@@ -1,11 +1,11 @@
 from contextlib import ExitStack as DoesNotRaise
-from typing import Optional, Tuple
+from typing import Optional, Tuple, List
 
 import pytest
 
 import numpy as np
 
-from supervision.detection.utils import non_max_suppression, clip_boxes
+from supervision.detection.utils import non_max_suppression, clip_boxes, filter_polygons_by_area
 
 
 @pytest.mark.parametrize(
@@ -186,3 +186,93 @@ def test_non_max_suppression(
 def test_clip_boxes(boxes_xyxy: np.ndarray, frame_resolution_wh: Tuple[int, int], expected_result: np.ndarray) -> None:
     result = clip_boxes(boxes_xyxy=boxes_xyxy, frame_resolution_wh=frame_resolution_wh)
     assert np.array_equal(result, expected_result)
+
+
+@pytest.mark.parametrize(
+    "polygons, min_area, max_area, expected_result, exception",
+    [
+        (
+            [np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
+            None,
+            None,
+            [np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
+            DoesNotRaise()
+        ),  # single polygon without area constraints
+        (
+            [np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
+            50,
+            None,
+            [np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
+            DoesNotRaise()
+        ),  # single polygon with min_area constraint
+        (
+            [np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
+            None,
+            50,
+            [],
+            DoesNotRaise()
+        ),  # single polygon with max_area constraint
+        (
+            [
+                np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
+                np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
+            ],
+            200,
+            None,
+            [np.array([[0, 0], [0, 20], [20, 20], [20, 0]])],
+            DoesNotRaise()
+        ),  # two polygons with min_area constraint
+        (
+            [
+                np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
+                np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
+            ],
+            None,
+            200,
+            [np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
+            DoesNotRaise()
+        ),  # two polygons with max_area constraint
+        (
+            [
+                np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
+                np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
+            ],
+            200,
+            200,
+            [],
+            DoesNotRaise()
+        ),  # two polygons with both area constraints
+        (
+            [
+                np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
+                np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
+            ],
+            100,
+            100,
+            [np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
+            DoesNotRaise()
+        ),  # two polygons with min_area and max_area equal to the area of the first polygon
+        (
+            [
+                np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
+                np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
+            ],
+            400,
+            400,
+            [np.array([[0, 0], [0, 20], [20, 20], [20, 0]])],
+            DoesNotRaise()
+        ),  # two polygons with min_area and max_area equal to the area of the second polygon
+    ]
+)
+def test_filter_polygons_by_area(
+        polygons: List[np.ndarray],
+        min_area: Optional[float],
+        max_area: Optional[float],
+        expected_result: List[np.ndarray],
+        exception: Exception
+) -> None:
+    with exception:
+        result = filter_polygons_by_area(polygons=polygons, min_area=min_area, max_area=max_area)
+        assert len(result) == len(expected_result)
+        for result_polygon, expected_result_polygon in zip(result, expected_result):
+            assert np.array_equal(result_polygon, expected_result_polygon)