From 68412500fa10e6bd36f586f6ee8cdb6cc260e30e Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Wed, 19 Apr 2023 13:09:53 +0200 Subject: [PATCH] =?UTF-8?q?=F0=9F=93=84=20docs=20update?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/detection/utils.md | 22 +++++++++++++---- supervision/__init__.py | 10 +++++++- supervision/detection/tools/polygon_zone.py | 4 +-- supervision/detection/utils.py | 27 +++++++++++++++------ 4 files changed, 48 insertions(+), 15 deletions(-) diff --git a/docs/detection/utils.md b/docs/detection/utils.md index ace7e1f0..4e33f311 100644 --- a/docs/detection/utils.md +++ b/docs/detection/utils.md @@ -1,7 +1,3 @@ -## generate_2d_mask - -:::supervision.detection.utils.generate_2d_mask - ## box_iou_batch :::supervision.detection.utils.box_iou_batch @@ -10,6 +6,22 @@ :::supervision.detection.utils.non_max_suppression +## polygon_to_mask + +:::supervision.detection.utils.polygon_to_mask + ## mask_to_xyxy -:::supervision.detection.utils.mask_to_xyxy \ No newline at end of file +:::supervision.detection.utils.mask_to_xyxy + +## mask_to_polygons + +:::supervision.detection.utils.mask_to_polygons + +## polygon_to_xyxy + +:::supervision.detection.utils.polygon_to_xyxy + +## filter_polygons_by_area + +:::supervision.detection.utils.filter_polygons_by_area \ No newline at end of file diff --git a/supervision/__init__.py b/supervision/__init__.py index 49348033..e5458b95 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -5,7 +5,15 @@ from supervision.detection.annotate import BoxAnnotator, MaskAnnotator from supervision.detection.core import Detections from supervision.detection.line_counter import LineZone, LineZoneAnnotator from supervision.detection.tools.polygon_zone import PolygonZone, PolygonZoneAnnotator -from supervision.detection.utils import generate_2d_mask, mask_to_xyxy +from supervision.detection.utils import ( + box_iou_batch, + filter_polygons_by_area, + mask_to_polygons, + mask_to_xyxy, + non_max_suppression, + polygon_to_mask, + polygon_to_xyxy, +) from supervision.draw.color import Color, ColorPalette from supervision.draw.utils import draw_filled_rectangle, draw_polygon, draw_text from supervision.file import list_files_with_extensions diff --git a/supervision/detection/tools/polygon_zone.py b/supervision/detection/tools/polygon_zone.py index 48c3eaa0..5615422b 100644 --- a/supervision/detection/tools/polygon_zone.py +++ b/supervision/detection/tools/polygon_zone.py @@ -5,7 +5,7 @@ import cv2 import numpy as np from supervision import Detections -from supervision.detection.utils import clip_boxes, generate_2d_mask +from supervision.detection.utils import clip_boxes, polygon_to_mask from supervision.draw.color import Color from supervision.draw.utils import draw_polygon, draw_text from supervision.geometry.core import Position @@ -36,7 +36,7 @@ class PolygonZone: self.current_count = 0 width, height = frame_resolution_wh - self.mask = generate_2d_mask( + self.mask = polygon_to_mask( polygon=polygon, resolution_wh=(width + 1, height + 1) ) diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index fa5adf10..0e77c1c1 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -6,8 +6,8 @@ import numpy as np MIN_POLYGON_POINT_COUNT = 3 -def generate_2d_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.ndarray: - """Generate a 2D mask from a polygon. +def polygon_to_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.ndarray: + """Generate a mask from a polygon. Args: polygon (np.ndarray): The polygon for which the mask should be generated, given as a list of vertices. @@ -151,6 +151,19 @@ def mask_to_xyxy(masks: np.ndarray) -> np.ndarray: def mask_to_polygons(mask: np.ndarray) -> List[np.ndarray]: + """ + Converts a binary mask to a list of polygons. + + Parameters: + mask (np.ndarray): A binary mask represented as a 2D NumPy array of shape `(H, W)`, + where H and W are the height and width of the mask, respectively. + + Returns: + List[np.ndarray]: A list of polygons, where each polygon is represented by a NumPy array of shape `(N, 2)`, + containing the `x`, `y` coordinates of the points. Polygons with fewer points than `MIN_POLYGON_POINT_COUNT = 3` + are excluded from the output. + """ + contours, _ = cv2.findContours( mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE ) @@ -170,8 +183,8 @@ def filter_polygons_by_area( Filters a list of polygons based on their area. Parameters: - polygons (List[np.ndarray]): A list of polygons, where each polygon is represented by a NumPy array of shape (N, 2), - containing the x, y coordinates of the points. + polygons (List[np.ndarray]): A list of polygons, where each polygon is represented by a NumPy array of shape `(N, 2)`, + containing the `x`, `y` coordinates of the points. min_area (Optional[float]): The minimum area threshold. Only polygons with an area greater than or equal to this value will be included in the output. If set to None, no minimum area constraint will be applied. max_area (Optional[float]): The maximum area threshold. Only polygons with an area less than or equal to this value @@ -196,11 +209,11 @@ def polygon_to_xyxy(polygon: np.ndarray) -> np.ndarray: Converts a polygon represented by a NumPy array into a bounding box. Parameters: - polygon (np.ndarray): A polygon represented by a NumPy array of shape (N, 2), - containing the x, y coordinates of the points. + polygon (np.ndarray): A polygon represented by a NumPy array of shape `(N, 2)`, + containing the `x`, `y` coordinates of the points. Returns: - np.ndarray: A 1D NumPy array containing the bounding box (x_min, y_min, x_max, y_max) of the input polygon. + np.ndarray: A 1D NumPy array containing the bounding box `(x_min, y_min, x_max, y_max)` of the input polygon. """ x_min, y_min = np.min(polygon, axis=0) x_max, y_max = np.max(polygon, axis=0)