From c1ccff474f0df7efdfc823300060115ead9f617c Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Mon, 10 Apr 2023 23:15:35 +0200 Subject: [PATCH] Documentation for PolygonZone and PolygonZoneAnnotator --- docs/detection/tools/polygon_zone.md | 7 +++ mkdocs.yml | 2 + supervision/__init__.py | 2 +- supervision/detection/annotate.py | 12 ++--- supervision/detection/tools/__init__.py | 0 .../detection/{ => tools}/polygon_zone.py | 44 +++++++++++++++++++ 6 files changed, 60 insertions(+), 7 deletions(-) create mode 100644 docs/detection/tools/polygon_zone.md create mode 100644 supervision/detection/tools/__init__.py rename supervision/detection/{ => tools}/polygon_zone.py (57%) diff --git a/docs/detection/tools/polygon_zone.md b/docs/detection/tools/polygon_zone.md new file mode 100644 index 00000000..54a66da9 --- /dev/null +++ b/docs/detection/tools/polygon_zone.md @@ -0,0 +1,7 @@ +## PolygonZone + +:::supervision.detection.tools.polygon_zone.PolygonZone + +## PolygonZoneAnnotator + +:::supervision.detection.tools.polygon_zone.PolygonZoneAnnotator \ No newline at end of file diff --git a/mkdocs.yml b/mkdocs.yml index 89d444ad..14d963f7 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -29,6 +29,8 @@ nav: - Core: detection/core.md - Annotate: detection/annotate.md - Utils: detection/utils.md + - Tools: + - Polygon Zone: detection/tools/polygon_zone.md - Draw: - Utils: draw/utils.md - Annotations: diff --git a/supervision/__init__.py b/supervision/__init__.py index 43278acb..eefd48c4 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -4,7 +4,7 @@ from supervision.annotation.voc import detections_to_voc_xml from supervision.detection.annotate import BoxAnnotator, MaskAnnotator from supervision.detection.core import Detections from supervision.detection.line_counter import LineZone, LineZoneAnnotator -from supervision.detection.polygon_zone import PolygonZone, PolygonZoneAnnotator +from supervision.detection.tools.polygon_zone import PolygonZone, PolygonZoneAnnotator from supervision.detection.utils import generate_2d_mask, mask_to_xyxy from supervision.draw.color import Color, ColorPalette from supervision.draw.utils import draw_filled_rectangle, draw_polygon, draw_text diff --git a/supervision/detection/annotate.py b/supervision/detection/annotate.py index 78dd899c..98c5244e 100644 --- a/supervision/detection/annotate.py +++ b/supervision/detection/annotate.py @@ -120,16 +120,16 @@ class BoxAnnotator: class MaskAnnotator: + """ + A class for overlaying masks on an image using detections provided. + + Attributes: + color (Union[Color, ColorPalette]): The color to fill the mask, can be a single color or a color palette + """ def __init__( self, color: Union[Color, ColorPalette] = ColorPalette.default(), ): - """ - A class for overlaying masks on an image using detections provided. - - Attributes: - color (Union[Color, ColorPalette]): The color to fill the mask, can be a single color or a color palette - """ self.color: Union[Color, ColorPalette] = color def annotate( diff --git a/supervision/detection/tools/__init__.py b/supervision/detection/tools/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/supervision/detection/polygon_zone.py b/supervision/detection/tools/polygon_zone.py similarity index 57% rename from supervision/detection/polygon_zone.py rename to supervision/detection/tools/polygon_zone.py index da60a62c..fb53ea30 100644 --- a/supervision/detection/polygon_zone.py +++ b/supervision/detection/tools/polygon_zone.py @@ -13,6 +13,16 @@ from supervision.geometry.utils import get_polygon_center class PolygonZone: + """ + A class for defining a polygon-shaped zone within a frame for detecting objects. + + Attributes: + polygon (np.ndarray): A numpy array defining the polygon vertices + frame_resolution_wh (Tuple[int, int]): The frame resolution (width, height) + triggering_position (Position): The position within the bounding box that triggers the zone (default: Position.BOTTOM_CENTER) + current_count (int): The current count of detected objects within the zone + mask (np.ndarray): The 2D bool mask for the polygon zone + """ def __init__( self, polygon: np.ndarray, @@ -30,6 +40,16 @@ class PolygonZone: ) def trigger(self, detections: Detections) -> np.ndarray: + """ + Determines if the detections are within the polygon zone. + + Parameters: + detections (Detections): The detections to be checked against the polygon zone + + Returns: + np.ndarray: A boolean numpy array indicating if each detection is within the polygon zone + """ + clipped_xyxy = clip_boxes( boxes_xyxy=detections.xyxy, frame_resolution_wh=self.frame_resolution_wh ) @@ -43,6 +63,20 @@ class PolygonZone: class PolygonZoneAnnotator: + """ + A class for annotating a polygon-shaped zone within a frame with a count of detected objects. + + Attributes: + zone (PolygonZone): The polygon zone to be annotated + color (Color): The color to draw the polygon lines + thickness (int): The thickness of the polygon lines, default is 2 + text_color (Color): The color of the text on the polygon, default is black + text_scale (float): The scale of the text on the polygon, default is 0.5 + text_thickness (int): The thickness of the text on the polygon, default is 1 + text_padding (int): The padding around the text on the polygon, default is 10 + font (int): The font type for the text on the polygon, default is cv2.FONT_HERSHEY_SIMPLEX + center (Tuple[int, int]): The center of the polygon for text placement + """ def __init__( self, zone: PolygonZone, @@ -64,6 +98,16 @@ class PolygonZoneAnnotator: self.center = get_polygon_center(polygon=zone.polygon) def annotate(self, scene: np.ndarray, label: Optional[str] = None) -> np.ndarray: + """ + Annotates the polygon zone within a frame with a count of detected objects. + + Parameters: + scene (np.ndarray): The image on which the polygon zone will be annotated + label (Optional[str]): An optional label for the count of detected objects within the polygon zone (default: None) + + Returns: + np.ndarray: The image with the polygon zone and count of detected objects + """ annotated_frame = draw_polygon( scene=scene, polygon=self.zone.polygon,