From 35d1827aeed5da63d135745d0f9ba2ccb83d44ea Mon Sep 17 00:00:00 2001 From: DancinParrot Date: Sat, 25 May 2024 16:39:35 +0800 Subject: [PATCH 1/3] feat: merge multi mask with bitwise or If a segmentation list of an annotation in a COCO dataset contains more than 1 nested list, seperate binary masks are created for each nested list (polygon). Then, the first polygon is designated as a parent and subsequent polygons in the list will be merged with the parent using numpy logical or. --- supervision/dataset/formats/coco.py | 3 +++ supervision/dataset/utils.py | 26 ++++++++++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/supervision/dataset/formats/coco.py b/supervision/dataset/formats/coco.py index 353e33f5e..7e83105f7 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_masks, 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_masks(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..f6dbe407f 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,31 @@ def approximate_mask_with_polygons( ] +def merge_masks(segmentations, resolution_wh): + parent = None + for s in segmentations: + if parent is None: + parent = polygon_to_mask( + polygon=np.reshape( + np.asarray(s, dtype=np.int32), + (-1, 2), + ), + resolution_wh=resolution_wh, + ) + else: + mask = polygon_to_mask( + polygon=np.reshape( + np.asarray(s, dtype=np.int32), + (-1, 2), + ), + resolution_wh=resolution_wh, + ) + + parent = np.logical_or(parent, mask) + + return parent + + def merge_class_lists(class_lists: List[List[str]]) -> List[str]: unique_classes = set() From 5b37308c37d384374bf5f632cba0658c2c26fd22 Mon Sep 17 00:00:00 2001 From: DancinParrot Date: Sun, 26 May 2024 00:00:07 +0800 Subject: [PATCH 2/3] test: unit test for merge_polygons (renamed from merge_masks) The function merge_masks() has been renamed to merge_polygons(). It has been refactored to include explicit function parameters and return value. --- supervision/dataset/formats/coco.py | 4 +- supervision/dataset/utils.py | 70 ++++++++++++++++++++--------- test/dataset/test_utils.py | 58 ++++++++++++++++++++++++ 3 files changed, 109 insertions(+), 23 deletions(-) diff --git a/supervision/dataset/formats/coco.py b/supervision/dataset/formats/coco.py index 7e83105f7..6b7d2bcf5 100644 --- a/supervision/dataset/formats/coco.py +++ b/supervision/dataset/formats/coco.py @@ -11,7 +11,7 @@ approximate_mask_with_polygons, map_detections_class_id, mask_to_rle, - merge_masks, + merge_polygons, rle_to_mask, ) from supervision.detection.core import Detections @@ -75,7 +75,7 @@ def coco_annotations_to_masks( resolution_wh=resolution_wh, ) if image_annotation["iscrowd"] - else merge_masks(image_annotation["segmentation"], resolution_wh) + else merge_polygons(image_annotation["segmentation"], resolution_wh) if len(image_annotation["segmentation"]) > 1 else polygon_to_mask( polygon=np.reshape( diff --git a/supervision/dataset/utils.py b/supervision/dataset/utils.py index f6dbe407f..33ca85b84 100644 --- a/supervision/dataset/utils.py +++ b/supervision/dataset/utils.py @@ -47,29 +47,57 @@ def approximate_mask_with_polygons( ] -def merge_masks(segmentations, resolution_wh): - parent = None - for s in segmentations: - if parent is None: - parent = polygon_to_mask( - polygon=np.reshape( - np.asarray(s, dtype=np.int32), - (-1, 2), - ), - resolution_wh=resolution_wh, - ) - else: - mask = polygon_to_mask( - polygon=np.reshape( - np.asarray(s, dtype=np.int32), - (-1, 2), - ), - resolution_wh=resolution_wh, - ) +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 = np.logical_or(parent, mask) + parent_mask = np.logical_or(parent_mask, child_mask) - return parent + return parent_mask.astype(float) def merge_class_lists(class_lists: List[List[str]]) -> List[str]: 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", [ From 8a137befcc2864976ce99ea9e06c33eb771e0448 Mon Sep 17 00:00:00 2001 From: DancinParrot Date: Sun, 26 May 2024 00:03:44 +0800 Subject: [PATCH 3/3] docs: add entry for merge_polygons() function in mkdocs --- docs/datasets/utils.md | 6 ++++++ 1 file changed, 6 insertions(+) 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