diff --git a/docs/datasets/utils.md b/docs/datasets/utils.md
index 6be56303f..c4cf6503e 100644
--- a/docs/datasets/utils.md
+++ b/docs/datasets/utils.md
@@ -16,3 +16,9 @@ status: new
:::supervision.dataset.utils.mask_to_rle
+
+
+
+:::supervision.dataset.utils.merge_polygons
diff --git a/supervision/dataset/formats/coco.py b/supervision/dataset/formats/coco.py
index 353e33f5e..6b7d2bcf5 100644
--- a/supervision/dataset/formats/coco.py
+++ b/supervision/dataset/formats/coco.py
@@ -11,6 +11,7 @@
approximate_mask_with_polygons,
map_detections_class_id,
mask_to_rle,
+ merge_polygons,
rle_to_mask,
)
from supervision.detection.core import Detections
@@ -74,6 +75,8 @@ def coco_annotations_to_masks(
resolution_wh=resolution_wh,
)
if image_annotation["iscrowd"]
+ else merge_polygons(image_annotation["segmentation"], resolution_wh)
+ if len(image_annotation["segmentation"]) > 1
else polygon_to_mask(
polygon=np.reshape(
np.asarray(image_annotation["segmentation"], dtype=np.int32),
diff --git a/supervision/dataset/utils.py b/supervision/dataset/utils.py
index 32ece6bf1..33ca85b84 100644
--- a/supervision/dataset/utils.py
+++ b/supervision/dataset/utils.py
@@ -13,6 +13,7 @@
approximate_polygon,
filter_polygons_by_area,
mask_to_polygons,
+ polygon_to_mask,
)
T = TypeVar("T")
@@ -46,6 +47,59 @@ def approximate_mask_with_polygons(
]
+def merge_polygons(
+ polygons: List[np.ndarray], resolution_wh: Tuple[int, int]
+) -> np.ndarray:
+ """
+ Merge polygons (in the form of vertices) within the segmentation list
+ of a COCO annotation.
+
+ Args:
+ polygons (List[np.ndarray]): A nested list of vertices, with each
+ nested list representing one polygon
+ resolution_wh (Tuple[int, int]): The width (w) and height (h)
+ of the desired binary mask.
+
+ Returns:
+ np.ndarray: The generated 2D mask, where the polygon is marked with
+ `1`'s and the rest is filled with `0`'s.
+
+ Examples:
+ ```python
+ import supervision as sv
+
+ sv.merge_polygons([[1.0, 1.0, 2.0, 3.0], [1.0, 1.0, 3.0, 2.0]], (8, 4))
+ # array([
+ # [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+ # [0.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+ # [0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0],
+ # [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+ # ])
+ ```
+ """
+ parent_polygon = polygons[0]
+ parent_mask = polygon_to_mask(
+ polygon=np.reshape(
+ np.asarray(parent_polygon, dtype=np.int32),
+ (-1, 2),
+ ),
+ resolution_wh=resolution_wh,
+ )
+
+ for p in polygons[1:]:
+ child_mask = polygon_to_mask(
+ polygon=np.reshape(
+ np.asarray(p, dtype=np.int32),
+ (-1, 2),
+ ),
+ resolution_wh=resolution_wh,
+ )
+
+ parent_mask = np.logical_or(parent_mask, child_mask)
+
+ return parent_mask.astype(float)
+
+
def merge_class_lists(class_lists: List[List[str]]) -> List[str]:
unique_classes = set()
diff --git a/test/dataset/test_utils.py b/test/dataset/test_utils.py
index 41e1da5bc..569873a45 100644
--- a/test/dataset/test_utils.py
+++ b/test/dataset/test_utils.py
@@ -12,6 +12,7 @@
map_detections_class_id,
mask_to_rle,
merge_class_lists,
+ merge_polygons,
rle_to_mask,
train_test_split,
)
@@ -125,6 +126,63 @@ def test_merge_class_maps(
assert result == expected_result
+@pytest.mark.parametrize(
+ "polygons, resolution_wh, expected_result, exception",
+ [
+ (
+ np.array([[1.0, 1.0, 2.0, 3.0]]),
+ (8, 4),
+ np.array(
+ [
+ [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+ [0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+ [0.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+ ]
+ ),
+ DoesNotRaise(),
+ ), # single polygon
+ (
+ np.array([[1.0, 1.0, 2.0, 3.0], [1.0, 1.0, 3.0, 2.0]]),
+ (8, 4),
+ np.array(
+ [
+ [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+ [0.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+ [0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+ ]
+ ),
+ DoesNotRaise(),
+ ), # two polygons
+ (
+ np.array(
+ [[1.0, 0.0, 2.0, 3.0], [1.0, 1.0, 3.0, 2.0], [1.0, 2.0, 1.0, 2.0]]
+ ),
+ (8, 4),
+ np.array(
+ [
+ [0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+ [0.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+ [0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+ ]
+ ),
+ DoesNotRaise(),
+ ), # multiple polygons
+ ],
+)
+def test_merge_polygons(
+ polygons: List[np.ndarray],
+ resolution_wh: Tuple[int, int],
+ expected_result: np.ndarray,
+ exception: Exception,
+) -> None:
+ with exception:
+ result = merge_polygons(polygons=polygons, resolution_wh=resolution_wh)
+ assert np.array_equal(result, expected_result)
+
+
@pytest.mark.parametrize(
"source_classes, target_classes, expected_result, exception",
[