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 + +
+

merge_polygons

+
+ +:::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", [