diff --git a/docs/detection/utils/converters.md b/docs/detection/utils/converters.md index b6b1e2af6c..354ce9308d 100644 --- a/docs/detection/utils/converters.md +++ b/docs/detection/utils/converters.md @@ -47,6 +47,12 @@ status: new :::supervision.detection.utils.converters.mask_to_polygons +
+

polygons_to_mask

+
+ +:::supervision.detection.utils.converters.polygons_to_mask +

polygon_to_mask

diff --git a/supervision/__init__.py b/supervision/__init__.py index 00820076fe..f099afb360 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -61,8 +61,8 @@ from supervision.detection.utils.converters import ( mask_to_polygons, mask_to_xyxy, - polygon_to_mask, polygon_to_xyxy, + polygons_to_mask, xcycwh_to_xyxy, xywh_to_xyxy, xyxy_to_mask, @@ -244,8 +244,8 @@ "pillow_to_cv2", "plot_image", "plot_images_grid", - "polygon_to_mask", "polygon_to_xyxy", + "polygons_to_mask", "process_video", "resize_image", "rle_to_mask", diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 26762e086f..6920b35be8 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -26,7 +26,7 @@ from supervision.detection.utils.boxes import clip_boxes, spread_out_boxes from supervision.detection.utils.converters import ( mask_to_polygons, - polygon_to_mask, + polygons_to_mask, xyxy_to_polygons, ) from supervision.draw.base import ImageType @@ -2936,7 +2936,7 @@ def _mask_from_xyxy(scene: np.ndarray, detections: Detections) -> np.ndarray: polygons = xyxy_to_polygons(detections.xyxy) for polygon in polygons: - polygon_mask = polygon_to_mask(polygon, resolution_wh=resolution_wh) + polygon_mask = polygons_to_mask(polygon, resolution_wh=resolution_wh) mask |= polygon_mask.astype(np.bool_) return mask @@ -2949,7 +2949,7 @@ def _mask_from_obb(scene: np.ndarray, detections: Detections) -> np.ndarray: resolution_wh = scene.shape[1], scene.shape[0] for polygon in detections.data[ORIENTED_BOX_COORDINATES]: - polygon_mask = polygon_to_mask(polygon, resolution_wh=resolution_wh) + polygon_mask = polygons_to_mask(polygon, resolution_wh=resolution_wh) mask |= polygon_mask.astype(np.bool_) return mask diff --git a/supervision/dataset/formats/coco.py b/supervision/dataset/formats/coco.py index b4827f29d4..d95efd834c 100644 --- a/supervision/dataset/formats/coco.py +++ b/supervision/dataset/formats/coco.py @@ -13,7 +13,7 @@ rle_to_mask, ) from supervision.detection.core import Detections -from supervision.detection.utils.converters import polygon_to_mask +from supervision.detection.utils.converters import polygons_to_mask from supervision.detection.utils.masks import contains_holes, contains_multiple_segments from supervision.utils.file import read_json_file, save_json_file @@ -68,17 +68,13 @@ def coco_annotations_to_masks( ) -> npt.NDArray[np.bool_]: return np.array( [ - rle_to_mask( - rle=np.array(image_annotation["segmentation"]["counts"]), - resolution_wh=resolution_wh, - ) - if image_annotation["iscrowd"] - else polygon_to_mask( - polygon=np.reshape( - np.asarray(image_annotation["segmentation"], dtype=np.int32), - (-1, 2), - ), - resolution_wh=resolution_wh, + ( + rle_to_mask( + rle=np.array(image_annotation["segmentation"]["counts"]), + resolution_wh=resolution_wh, + ) + if image_annotation["iscrowd"] + else polygons_to_mask(image_annotation["segmentation"], resolution_wh) ) for image_annotation in image_annotations ], diff --git a/supervision/dataset/formats/pascal_voc.py b/supervision/dataset/formats/pascal_voc.py index c5172c33da..d5b2a54872 100644 --- a/supervision/dataset/formats/pascal_voc.py +++ b/supervision/dataset/formats/pascal_voc.py @@ -11,7 +11,7 @@ from supervision.dataset.utils import approximate_mask_with_polygons from supervision.detection.core import Detections -from supervision.detection.utils.converters import polygon_to_mask, polygon_to_xyxy +from supervision.detection.utils.converters import polygon_to_xyxy, polygons_to_mask from supervision.utils.file import list_files_with_extensions @@ -246,8 +246,8 @@ def detections_from_xml_obj( # https://github.com/roboflow/supervision/issues/144 polygon -= 1 - mask_from_polygon = polygon_to_mask( - polygon=polygon, + mask_from_polygon = polygons_to_mask( + polygons=polygon, resolution_wh=resolution_wh, ) masks.append(mask_from_polygon) diff --git a/supervision/dataset/formats/yolo.py b/supervision/dataset/formats/yolo.py index 8af010fa92..fc34ad7051 100644 --- a/supervision/dataset/formats/yolo.py +++ b/supervision/dataset/formats/yolo.py @@ -10,7 +10,7 @@ from supervision.config import ORIENTED_BOX_COORDINATES from supervision.dataset.utils import approximate_mask_with_polygons from supervision.detection.core import Detections -from supervision.detection.utils.converters import polygon_to_mask, polygon_to_xyxy +from supervision.detection.utils.converters import polygon_to_xyxy, polygons_to_mask from supervision.utils.file import ( list_files_with_extensions, read_txt_file, @@ -51,7 +51,7 @@ def _polygons_to_masks( ) -> np.ndarray: return np.array( [ - polygon_to_mask(polygon=polygon, resolution_wh=resolution_wh) + polygons_to_mask(polygons=[polygon], resolution_wh=resolution_wh) for polygon in polygons ], dtype=bool, diff --git a/supervision/detection/tools/polygon_zone.py b/supervision/detection/tools/polygon_zone.py index 952da5141b..540bb93097 100644 --- a/supervision/detection/tools/polygon_zone.py +++ b/supervision/detection/tools/polygon_zone.py @@ -9,7 +9,7 @@ from supervision import Detections from supervision.detection.utils.boxes import clip_boxes -from supervision.detection.utils.converters import polygon_to_mask +from supervision.detection.utils.converters import polygons_to_mask from supervision.draw.color import Color from supervision.draw.utils import draw_filled_polygon, draw_polygon, draw_text from supervision.geometry.core import Position @@ -73,8 +73,8 @@ def __init__( x_max, y_max = np.max(polygon, axis=0) self.frame_resolution_wh = (x_max + 1, y_max + 1) - self.mask = polygon_to_mask( - polygon=polygon, resolution_wh=(x_max + 2, y_max + 2) + self.mask = polygons_to_mask( + polygons=[polygon], resolution_wh=(x_max + 2, y_max + 2) ) def trigger(self, detections: Detections) -> npt.NDArray[np.bool_]: diff --git a/supervision/detection/utils/converters.py b/supervision/detection/utils/converters.py index 4aef2dc87c..457fd2da41 100644 --- a/supervision/detection/utils/converters.py +++ b/supervision/detection/utils/converters.py @@ -1,6 +1,8 @@ import cv2 import numpy as np +from supervision.utils.internal import deprecated + MIN_POLYGON_POINT_COUNT = 3 @@ -23,6 +25,60 @@ def xyxy_to_polygons(box: np.ndarray) -> np.ndarray: return polygon +def polygons_to_mask( + polygons: list[np.ndarray], resolution_wh: tuple[int, int] +) -> np.ndarray: + """Generate a mask from an array of polygon(s). + + Args: + polygons (List[np.ndarray]): A nested list of vertices, with each + 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.polygons_to_mask([[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], + # ]) + ``` + """ + + width, height = map(int, resolution_wh) + parent_polygon = np.reshape( + np.asarray(polygons[0], dtype=np.int32), + (-1, 2), + ) + parent_mask = np.zeros((height, width), dtype=np.uint8) + cv2.fillPoly(parent_mask, [parent_polygon.astype(np.int32)], color=1) + + for p in polygons[1:]: + child_polygon = np.reshape( + np.asarray(p, dtype=np.int32), + (-1, 2), + ) + child_mask = np.zeros((height, width), dtype=np.uint8) + cv2.fillPoly(child_mask, [child_polygon.astype(np.int32)], color=1) + + parent_mask = np.logical_or(parent_mask, child_mask) + + return parent_mask.astype(float) + + +@deprecated( + "`polygon_to_mask` is deprecated and will be removed in " + "`supervision-0.27.0`. Use `polygons_to_mask` instead." +) def polygon_to_mask(polygon: np.ndarray, resolution_wh: tuple[int, int]) -> np.ndarray: """Generate a mask from a polygon. diff --git a/supervision/detection/utils/internal.py b/supervision/detection/utils/internal.py index f5d6dc9fbf..eb46997f62 100644 --- a/supervision/detection/utils/internal.py +++ b/supervision/detection/utils/internal.py @@ -8,7 +8,7 @@ import numpy.typing as npt from supervision.config import CLASS_NAME_DATA_FIELD -from supervision.detection.utils.converters import polygon_to_mask +from supervision.detection.utils.converters import polygons_to_mask from supervision.geometry.core import Vector @@ -99,7 +99,9 @@ def process_roboflow_result( polygon = np.array( [[point["x"], point["y"]] for point in prediction["points"]], dtype=int ) - mask = polygon_to_mask(polygon, resolution_wh=(image_width, image_height)) + mask = polygons_to_mask( + [polygon], resolution_wh=(image_width, image_height) + ) xyxy.append([x_min, y_min, x_max, y_max]) class_id.append(prediction["class_id"]) class_name.append(prediction["class"]) @@ -152,12 +154,14 @@ def is_metadata_equal(metadata_a: dict[str, Any], metadata_b: dict[str, Any]) -> True if the metadata payloads are equal, False otherwise. """ return set(metadata_a.keys()) == set(metadata_b.keys()) and all( - np.array_equal(metadata_a[key], metadata_b[key]) - if ( - isinstance(metadata_a[key], np.ndarray) - and isinstance(metadata_b[key], np.ndarray) + ( + np.array_equal(metadata_a[key], metadata_b[key]) + if ( + isinstance(metadata_a[key], np.ndarray) + and isinstance(metadata_b[key], np.ndarray) + ) + else metadata_a[key] == metadata_b[key] ) - else metadata_a[key] == metadata_b[key] for key in metadata_a ) diff --git a/supervision/detection/utils/iou_and_nms.py b/supervision/detection/utils/iou_and_nms.py index 8a444cf16e..50ed046948 100644 --- a/supervision/detection/utils/iou_and_nms.py +++ b/supervision/detection/utils/iou_and_nms.py @@ -5,7 +5,7 @@ import numpy as np import numpy.typing as npt -from supervision.detection.utils.converters import polygon_to_mask +from supervision.detection.utils.converters import polygons_to_mask from supervision.detection.utils.masks import resize_masks @@ -383,11 +383,11 @@ def oriented_box_iou_batch( mask_true = np.zeros((boxes_true.shape[0], max_height, max_width)) for i, box_true in enumerate(boxes_true): - mask_true[i] = polygon_to_mask(box_true, (max_width, max_height)) + mask_true[i] = polygons_to_mask(box_true, (max_width, max_height)) mask_detection = np.zeros((boxes_detection.shape[0], max_height, max_width)) for i, box_detection in enumerate(boxes_detection): - mask_detection[i] = polygon_to_mask(box_detection, (max_width, max_height)) + mask_detection[i] = polygons_to_mask(box_detection, (max_width, max_height)) ious = mask_iou_batch(mask_true, mask_detection) return ious diff --git a/supervision/detection/vlm.py b/supervision/detection/vlm.py index 97988c9f09..31768f4edf 100644 --- a/supervision/detection/vlm.py +++ b/supervision/detection/vlm.py @@ -12,7 +12,7 @@ from PIL import Image from supervision.detection.utils.boxes import denormalize_boxes -from supervision.detection.utils.converters import polygon_to_mask, polygon_to_xyxy +from supervision.detection.utils.converters import polygon_to_xyxy, polygons_to_mask from supervision.utils.internal import deprecated from supervision.validators import validate_resolution @@ -537,7 +537,7 @@ def from_florence_2( for polygons_of_same_class in result["polygons"]: for polygon in polygons_of_same_class: polygon = np.reshape(polygon, (-1, 2)).astype(np.int32) - mask = polygon_to_mask(polygon, resolution_wh).astype(bool) + mask = polygons_to_mask([polygon], resolution_wh).astype(bool) masks_list.append(mask) xyxy = polygon_to_xyxy(polygon) xyxy_list.append(xyxy) diff --git a/test/detection/utils/test_converters.py b/test/detection/utils/test_converters.py index 52a3b52004..7d446076ed 100644 --- a/test/detection/utils/test_converters.py +++ b/test/detection/utils/test_converters.py @@ -1,9 +1,12 @@ from __future__ import annotations +from contextlib import ExitStack as DoesNotRaise + import numpy as np import pytest from supervision.detection.utils.converters import ( + polygons_to_mask, xcycwh_to_xyxy, xywh_to_xyxy, xyxy_to_mask, @@ -301,3 +304,64 @@ def test_xyxy_to_mask(boxes: np.ndarray, resolution_wh, expected: np.ndarray) -> assert result.dtype == np.bool_ assert result.shape == expected.shape np.testing.assert_array_equal(result, expected) + + +@pytest.mark.parametrize( + "polygons, resolution_wh, expected_result, exception", + [ + ( + [np.array([1, 1, 3, 1, 3, 3, 1, 3])], + (5, 5), + np.array( + [ + [0, 0, 0, 0, 0], + [0, 1, 1, 1, 0], + [0, 1, 1, 1, 0], + [0, 1, 1, 1, 0], + [0, 0, 0, 0, 0], + ] + ), + DoesNotRaise(), + ), # one polygon + ( + [np.array([1, 1, 3, 1, 3, 3, 1, 3]), np.array([1, 1, 3, 1, 3, 3, 1, 3])], + (5, 5), + np.array( + [ + [0, 0, 0, 0, 0], + [0, 1, 1, 1, 0], + [0, 1, 1, 1, 0], + [0, 1, 1, 1, 0], + [0, 0, 0, 0, 0], + ] + ), + DoesNotRaise(), + ), # overlapping polygons + ( + [ + np.array([0, 0, 1, 0, 1, 1, 0, 1]), + np.array([3, 3, 4, 3, 4, 4, 3, 4]), + ], + (5, 5), + np.array( + [ + [1, 1, 0, 0, 0], + [1, 1, 0, 0, 0], + [0, 0, 0, 0, 0], + [0, 0, 0, 1, 1], + [0, 0, 0, 1, 1], + ] + ), + DoesNotRaise(), + ), # multiple polygons + ], +) +def test_polygons_to_mask( + polygons: list[np.ndarray], + resolution_wh: tuple[int, int], + expected_result: np.ndarray, + exception: Exception, +) -> None: + with exception: + result = polygons_to_mask(polygons=polygons, resolution_wh=resolution_wh) + assert np.array_equal(result, expected_result)