From 0cb722e91ee04780dccdefe81c32f6ca258ee45d Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Tue, 18 Apr 2023 14:48:25 +0200 Subject: [PATCH] `mask_to_polygons` and `filter_polygons_by_area` added --- supervision/dataset/formats/pascal_voc.py | 15 ++-- supervision/detection/utils.py | 45 ++++++++++- test/detection/test_utils.py | 94 ++++++++++++++++++++++- 3 files changed, 142 insertions(+), 12 deletions(-) diff --git a/supervision/dataset/formats/pascal_voc.py b/supervision/dataset/formats/pascal_voc.py index a9696d6a..9a6b89e8 100644 --- a/supervision/dataset/formats/pascal_voc.py +++ b/supervision/dataset/formats/pascal_voc.py @@ -1,4 +1,4 @@ -from typing import List +from typing import List, Tuple from xml.dom.minidom import parseString from xml.etree.ElementTree import Element, SubElement, tostring @@ -9,9 +9,7 @@ def detections_to_pascal_voc( detections: Detections, classes: List[str], filename: str, - width: int, - height: int, - depth: int = 3, + image_shape: Tuple[int, int, int] ) -> str: """ Converts Detections object to Pascal VOC XML format. @@ -20,12 +18,11 @@ def detections_to_pascal_voc( detections (Detections): A Detections object containing bounding boxes, class ids, and other relevant information. classes (List[str]): A list of class names corresponding to the class ids in the Detections object. filename (str): The name of the image file associated with the detections. - width (int): The width of the image in pixels. - height (int): The height of the image in pixels. - depth (int, optional): The number of color channels in the image. Defaults to 3 for RGB images. + image_shape (Tuple[int, int, int]): The shape of the image file associated with the detections. Returns: str: An XML string in Pascal VOC format representing the detections. """ + height, width, depth = image_shape # Create root element annotation = Element("annotation") @@ -35,8 +32,8 @@ def detections_to_pascal_voc( folder.text = "VOC" # Add filename element - fname = SubElement(annotation, "filename") - fname.text = filename + file_name = SubElement(annotation, "filename") + file_name.text = filename # Add source element source = SubElement(annotation, "source") diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 4db14be7..91b1f53d 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -1,9 +1,12 @@ -from typing import Tuple +from typing import Tuple, List, Optional import cv2 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. @@ -146,3 +149,43 @@ def mask_to_xyxy(masks: np.ndarray) -> np.ndarray: bboxes[i, :] = [x_min, y_min, x_max, y_max] return bboxes + + +def mask_to_polygons(mask: np.ndarray) -> List[np.ndarray]: + contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) + return [ + np.squeeze(contour, axis=1) + for contour + in contours + if contour.shape[0] >= MIN_POLYGON_POINT_COUNT + ] + + +def filter_polygons_by_area(polygons: List[np.ndarray], min_area: Optional[float], max_area: Optional[float]) -> List[np.ndarray]: + """ + 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. + 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 + will be included in the output. If set to None, no maximum area constraint will be applied. + + Returns: + List[np.ndarray]: A new list of polygons containing only those with areas within the specified thresholds. + """ + if min_area is None and max_area is None: + return polygons + ares = [ + cv2.contourArea(polygon) + for polygon + in polygons + ] + return [ + polygon + for polygon, area + in zip(polygons, ares) + if (min_area is None or area>= min_area) and (max_area is None or area<= max_area) + ] diff --git a/test/detection/test_utils.py b/test/detection/test_utils.py index 9ba3354a..cf7abfd3 100644 --- a/test/detection/test_utils.py +++ b/test/detection/test_utils.py @@ -1,11 +1,11 @@ from contextlib import ExitStack as DoesNotRaise -from typing import Optional, Tuple +from typing import Optional, Tuple, List import pytest import numpy as np -from supervision.detection.utils import non_max_suppression, clip_boxes +from supervision.detection.utils import non_max_suppression, clip_boxes, filter_polygons_by_area @pytest.mark.parametrize( @@ -186,3 +186,93 @@ def test_non_max_suppression( def test_clip_boxes(boxes_xyxy: np.ndarray, frame_resolution_wh: Tuple[int, int], expected_result: np.ndarray) -> None: result = clip_boxes(boxes_xyxy=boxes_xyxy, frame_resolution_wh=frame_resolution_wh) assert np.array_equal(result, expected_result) + + +@pytest.mark.parametrize( + "polygons, min_area, max_area, expected_result, exception", + [ + ( + [np.array([[0, 0], [0, 10], [10, 10], [10, 0]])], + None, + None, + [np.array([[0, 0], [0, 10], [10, 10], [10, 0]])], + DoesNotRaise() + ), # single polygon without area constraints + ( + [np.array([[0, 0], [0, 10], [10, 10], [10, 0]])], + 50, + None, + [np.array([[0, 0], [0, 10], [10, 10], [10, 0]])], + DoesNotRaise() + ), # single polygon with min_area constraint + ( + [np.array([[0, 0], [0, 10], [10, 10], [10, 0]])], + None, + 50, + [], + DoesNotRaise() + ), # single polygon with max_area constraint + ( + [ + np.array([[0, 0], [0, 10], [10, 10], [10, 0]]), + np.array([[0, 0], [0, 20], [20, 20], [20, 0]]) + ], + 200, + None, + [np.array([[0, 0], [0, 20], [20, 20], [20, 0]])], + DoesNotRaise() + ), # two polygons with min_area constraint + ( + [ + np.array([[0, 0], [0, 10], [10, 10], [10, 0]]), + np.array([[0, 0], [0, 20], [20, 20], [20, 0]]) + ], + None, + 200, + [np.array([[0, 0], [0, 10], [10, 10], [10, 0]])], + DoesNotRaise() + ), # two polygons with max_area constraint + ( + [ + np.array([[0, 0], [0, 10], [10, 10], [10, 0]]), + np.array([[0, 0], [0, 20], [20, 20], [20, 0]]) + ], + 200, + 200, + [], + DoesNotRaise() + ), # two polygons with both area constraints + ( + [ + np.array([[0, 0], [0, 10], [10, 10], [10, 0]]), + np.array([[0, 0], [0, 20], [20, 20], [20, 0]]) + ], + 100, + 100, + [np.array([[0, 0], [0, 10], [10, 10], [10, 0]])], + DoesNotRaise() + ), # two polygons with min_area and max_area equal to the area of the first polygon + ( + [ + np.array([[0, 0], [0, 10], [10, 10], [10, 0]]), + np.array([[0, 0], [0, 20], [20, 20], [20, 0]]) + ], + 400, + 400, + [np.array([[0, 0], [0, 20], [20, 20], [20, 0]])], + DoesNotRaise() + ), # two polygons with min_area and max_area equal to the area of the second polygon + ] +) +def test_filter_polygons_by_area( + polygons: List[np.ndarray], + min_area: Optional[float], + max_area: Optional[float], + expected_result: List[np.ndarray], + exception: Exception +) -> None: + with exception: + result = filter_polygons_by_area(polygons=polygons, min_area=min_area, max_area=max_area) + assert len(result) == len(expected_result) + for result_polygon, expected_result_polygon in zip(result, expected_result): + assert np.array_equal(result_polygon, expected_result_polygon)