133 lines
4.9 KiB
Python
133 lines
4.9 KiB
Python
from dataclasses import replace
|
|
from typing import Optional, Tuple
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from supervision import Detections
|
|
from supervision.detection.utils import clip_boxes, generate_2d_mask
|
|
from supervision.draw.color import Color
|
|
from supervision.draw.utils import draw_polygon, draw_text
|
|
from supervision.geometry.core import Position
|
|
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,
|
|
frame_resolution_wh: Tuple[int, int],
|
|
triggering_position: Position = Position.BOTTOM_CENTER,
|
|
):
|
|
self.polygon = polygon
|
|
self.frame_resolution_wh = frame_resolution_wh
|
|
self.triggering_position = triggering_position
|
|
self.current_count = 0
|
|
|
|
width, height = frame_resolution_wh
|
|
self.mask = generate_2d_mask(
|
|
polygon=polygon, resolution_wh=(width + 1, height + 1)
|
|
)
|
|
|
|
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
|
|
)
|
|
clipped_detections = replace(detections, xyxy=clipped_xyxy)
|
|
clipped_anchors = np.ceil(
|
|
clipped_detections.get_anchor_coordinates(anchor=self.triggering_position)
|
|
).astype(int)
|
|
is_in_zone = self.mask[clipped_anchors[:, 1], clipped_anchors[:, 0]]
|
|
self.current_count = np.sum(is_in_zone)
|
|
return is_in_zone.astype(bool)
|
|
|
|
|
|
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,
|
|
color: Color,
|
|
thickness: int = 2,
|
|
text_color: Color = Color.black(),
|
|
text_scale: float = 0.5,
|
|
text_thickness: int = 1,
|
|
text_padding: int = 10,
|
|
):
|
|
self.zone = zone
|
|
self.color = color
|
|
self.thickness = thickness
|
|
self.text_color = text_color
|
|
self.text_scale = text_scale
|
|
self.text_thickness = text_thickness
|
|
self.text_padding = text_padding
|
|
self.font = cv2.FONT_HERSHEY_SIMPLEX
|
|
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,
|
|
color=self.color,
|
|
thickness=self.thickness,
|
|
)
|
|
|
|
annotated_frame = draw_text(
|
|
scene=annotated_frame,
|
|
text=str(self.zone.current_count) if label is None else label,
|
|
text_anchor=self.center,
|
|
background_color=self.color,
|
|
text_color=self.text_color,
|
|
text_scale=self.text_scale,
|
|
text_thickness=self.text_thickness,
|
|
text_padding=self.text_padding,
|
|
text_font=self.font,
|
|
)
|
|
|
|
return annotated_frame
|