diff --git a/supervision/__init__.py b/supervision/__init__.py index 15ebf2c7..ccd27293 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -87,8 +87,8 @@ from supervision.detection.utils.masks import ( calculate_masks_centroids, contains_holes, contains_multiple_segments, - move_masks, filter_segments_by_distance, + move_masks, ) from supervision.detection.utils.polygons import ( approximate_polygon, diff --git a/supervision/detection/utils/masks.py b/supervision/detection/utils/masks.py index 11943950..49b40311 100644 --- a/supervision/detection/utils/masks.py +++ b/supervision/detection/utils/masks.py @@ -396,4 +396,4 @@ def filter_segments_by_distance( else: raise ValueError("mode must be 'edge' or 'centroid'") - return keep_labels[labels] \ No newline at end of file + return keep_labels[labels] diff --git a/test/detection/utils/test_masks.py b/test/detection/utils/test_masks.py index 287dade7..3c61e409 100644 --- a/test/detection/utils/test_masks.py +++ b/test/detection/utils/test_masks.py @@ -10,7 +10,8 @@ from supervision.detection.utils.masks import ( calculate_masks_centroids, contains_holes, contains_multiple_segments, - move_masks, filter_segments_by_distance, + filter_segments_by_distance, + move_masks, ) @@ -507,98 +508,122 @@ def test_contains_multiple_segments( [ # single component, unchanged ( - np.array([ - [0, 0, 0, 0, 0, 0], - [0, 1, 1, 1, 0, 0], - [0, 1, 1, 1, 0, 0], - [0, 1, 1, 1, 0, 0], - [0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0], - ], dtype=bool), + np.array( + [ + [0, 0, 0, 0, 0, 0], + [0, 1, 1, 1, 0, 0], + [0, 1, 1, 1, 0, 0], + [0, 1, 1, 1, 0, 0], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + ], + dtype=bool, + ), 8, "edge", 2.0, None, - np.array([ - [0, 0, 0, 0, 0, 0], - [0, 1, 1, 1, 0, 0], - [0, 1, 1, 1, 0, 0], - [0, 1, 1, 1, 0, 0], - [0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0], - ], dtype=bool), + np.array( + [ + [0, 0, 0, 0, 0, 0], + [0, 1, 1, 1, 0, 0], + [0, 1, 1, 1, 0, 0], + [0, 1, 1, 1, 0, 0], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + ], + dtype=bool, + ), DoesNotRaise(), ), # two components, edge distance 2, kept with abs=1 ( - np.array([ - [0, 0, 0, 0, 0, 0], - [0, 1, 1, 1, 0, 1], - [0, 1, 1, 1, 0, 1], - [0, 1, 1, 1, 0, 1], - [0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0], - ], dtype=bool), + np.array( + [ + [0, 0, 0, 0, 0, 0], + [0, 1, 1, 1, 0, 1], + [0, 1, 1, 1, 0, 1], + [0, 1, 1, 1, 0, 1], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + ], + dtype=bool, + ), 8, "edge", 2.0, None, - np.array([ - [0, 0, 0, 0, 0, 0], - [0, 1, 1, 1, 0, 1], - [0, 1, 1, 1, 0, 1], - [0, 1, 1, 1, 0, 1], - [0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0], - ], dtype=bool), + np.array( + [ + [0, 0, 0, 0, 0, 0], + [0, 1, 1, 1, 0, 1], + [0, 1, 1, 1, 0, 1], + [0, 1, 1, 1, 0, 1], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + ], + dtype=bool, + ), DoesNotRaise(), ), # centroid mode, far centroids, dropped with small relative threshold ( - np.array([ - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 0, 0, 0], - [0, 0, 0, 1, 1, 1], - [0, 0, 0, 1, 1, 1], - ], dtype=bool), + np.array( + [ + [1, 1, 1, 0, 0, 0], + [1, 1, 1, 0, 0, 0], + [1, 1, 1, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 1, 1, 1], + [0, 0, 0, 1, 1, 1], + ], + dtype=bool, + ), 8, "centroid", None, 0.3, # diagonal ~8.49, threshold ~2.55, centroid gap ~4.24 - np.array([ - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0], - ], dtype=bool), + np.array( + [ + [1, 1, 1, 0, 0, 0], + [1, 1, 1, 0, 0, 0], + [1, 1, 1, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + ], + dtype=bool, + ), DoesNotRaise(), ), # centroid mode, larger relative threshold, kept ( - np.array([ - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 0, 0, 0], - [0, 0, 0, 1, 1, 1], - [0, 0, 0, 1, 1, 1], - ], dtype=bool), + np.array( + [ + [1, 1, 1, 0, 0, 0], + [1, 1, 1, 0, 0, 0], + [1, 1, 1, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 1, 1, 1], + [0, 0, 0, 1, 1, 1], + ], + dtype=bool, + ), 8, "centroid", None, 0.6, # diagonal ~8.49, threshold ~5.09, centroid gap ~4.24 - np.array([ - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [1, 1, 1, 0, 0, 0], - [0, 0, 0, 0, 0, 0], - [0, 0, 0, 1, 1, 1], - [0, 0, 0, 1, 1, 1], - ], dtype=bool), + np.array( + [ + [1, 1, 1, 0, 0, 0], + [1, 1, 1, 0, 0, 0], + [1, 1, 1, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 1, 1, 1], + [0, 0, 0, 1, 1, 1], + ], + dtype=bool, + ), DoesNotRaise(), ), # empty mask @@ -623,54 +648,65 @@ def test_contains_multiple_segments( ), # two components, pixel distance = 2, kept with abs=2 ( - np.array([ - [0, 0, 0, 0, 0, 0, 0, 0], - [0, 1, 1, 1, 0, 1, 1, 1], - [0, 1, 1, 1, 0, 1, 1, 1], - [0, 1, 1, 1, 0, 1, 1, 1], - [0, 0, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0, 0], - ], dtype=bool), + np.array( + [ + [0, 0, 0, 0, 0, 0, 0, 0], + [0, 1, 1, 1, 0, 1, 1, 1], + [0, 1, 1, 1, 0, 1, 1, 1], + [0, 1, 1, 1, 0, 1, 1, 1], + [0, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0], + ], + dtype=bool, + ), 8, "edge", 2.0, # was 1.0 None, - np.array([ - [0, 0, 0, 0, 0, 0, 0, 0], - [0, 1, 1, 1, 0, 1, 1, 1], - [0, 1, 1, 1, 0, 1, 1, 1], - [0, 1, 1, 1, 0, 1, 1, 1], - [0, 0, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0, 0], - ], dtype=bool), + np.array( + [ + [0, 0, 0, 0, 0, 0, 0, 0], + [0, 1, 1, 1, 0, 1, 1, 1], + [0, 1, 1, 1, 0, 1, 1, 1], + [0, 1, 1, 1, 0, 1, 1, 1], + [0, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0], + ], + dtype=bool, + ), DoesNotRaise(), ), - # two components, pixel distance = 3, dropped with abs=2 ( - np.array([ - [0, 0, 0, 0, 0, 0, 0, 0, 0], - [0, 1, 1, 1, 0, 0, 0, 1, 1], - [0, 1, 1, 1, 0, 0, 0, 1, 1], - [0, 1, 1, 1, 0, 0, 0, 1, 1], - [0, 0, 0, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0, 0, 0], - ], dtype=bool), + np.array( + [ + [0, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 1, 1, 1, 0, 0, 0, 1, 1], + [0, 1, 1, 1, 0, 0, 0, 1, 1], + [0, 1, 1, 1, 0, 0, 0, 1, 1], + [0, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0], + ], + dtype=bool, + ), 8, "edge", 2.0, # keep threshold below 3 so the right blob is removed None, - np.array([ - [0, 0, 0, 0, 0, 0, 0, 0, 0], - [0, 1, 1, 1, 0, 0, 0, 0, 0], - [0, 1, 1, 1, 0, 0, 0, 0, 0], - [0, 1, 1, 1, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0, 0, 0], - ], dtype=bool), + np.array( + [ + [0, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 1, 1, 1, 0, 0, 0, 0, 0], + [0, 1, 1, 1, 0, 0, 0, 0, 0], + [0, 1, 1, 1, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0], + ], + dtype=bool, + ), DoesNotRaise(), ), - ] + ], ) def test_filter_segments_by_distance_sweep( mask: npt.NDArray,