xiaoxue_admin
/
supervision
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #1448 from roboflow/feat/annotators-cleanup 

Annotator types and docs: MyPy, consistent 'optional', small tidy-ups
This commit is contained in:
LinasKo 2024-08-27 14:07:04 +03:00 committed by GitHub
parent e8ec271f0b a3262a8f9c
commit 1ae8fc798e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
21 changed files with 178 additions and 107 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
docs/detection/annotators.md
View File

@ -411,7 +411,7 @@ status: new
trace_annotator = sv.TraceAnnotator()
video_info = sv.VideoInfo.from_video_path(video_path='...')
frames_generator = get_video_frames_generator(source_path='...')
frames_generator = sv.get_video_frames_generator(source_path='...')
tracker = sv.ByteTrack()
with sv.VideoSink(target_path='...', video_info=video_info) as sink:
@ -442,7 +442,7 @@ status: new
heat_map_annotator = sv.HeatMapAnnotator()
video_info = sv.VideoInfo.from_video_path(video_path='...')
frames_generator = get_video_frames_generator(source_path='...')
frames_generator = sv.get_video_frames_generator(source_path='...')
with sv.VideoSink(target_path='...', video_info=video_info) as sink:
for frame in frames_generator:

2
examples/count_people_in_zone/inference_example.py
View File

@ -75,7 +75,7 @@ def detect(
frame (np.ndarray): The frame to process, expected to be a NumPy array.
model (RoboflowInferenceModel): The Inference model used for processing the
frame.
confidence_threshold (float, optional): The confidence threshold for filtering
confidence_threshold (float): The confidence threshold for filtering
detections. Default is 0.5.
Returns:

2
examples/count_people_in_zone/ultralytics_example.py
View File

@ -72,7 +72,7 @@ def detect(
Args:
frame (np.ndarray): The frame to process, expected to be a NumPy array.
model (YOLO): The YOLO model used for processing the frame.
confidence_threshold (float, optional): The confidence threshold for filtering
confidence_threshold (float): The confidence threshold for filtering
detections. Default is 0.5.
Returns:

2
examples/time_in_zone/utils/timers.py
View File

@ -22,7 +22,7 @@ class FPSBasedTimer:
"""Initializes the FPSBasedTimer with the specified frames per second rate.
Args:
fps (int, optional): The frame rate of the video stream. Defaults to 30.
fps (int): The frame rate of the video stream. Defaults to 30.
"""
self.fps = fps
self.frame_id = 0

97
supervision/annotators/core.py
View File

@ -4,7 +4,8 @@ from typing import List, Optional, Tuple, Union
import cv2
import numpy as np
from PIL import ImageDraw, ImageFont
import numpy.typing as npt
from PIL import Image, ImageDraw, ImageFont
from supervision.annotators.base import BaseAnnotator, ImageType
from supervision.annotators.utils import (
@ -93,6 +94,7 @@ class BoxAnnotator(BaseAnnotator):
![bounding-box-annotator-example](https://media.roboflow.com/
supervision-annotator-examples/bounding-box-annotator-example-purple.png)
"""
assert isinstance(scene, np.ndarray)
for detection_idx in range(len(detections)):
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
color = resolve_color(
@ -178,6 +180,7 @@ class BoundingBoxAnnotator(BaseAnnotator):
![bounding-box-annotator-example](https://media.roboflow.com/
supervision-annotator-examples/bounding-box-annotator-example-purple.png)
"""
assert isinstance(scene, np.ndarray)
for detection_idx in range(len(detections)):
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
color = resolve_color(
@ -262,12 +265,13 @@ class OrientedBoxAnnotator(BaseAnnotator):
)
```
""" # noqa E501 // docs
assert isinstance(scene, np.ndarray)
if detections.data is None or ORIENTED_BOX_COORDINATES not in detections.data:
return scene
obb_boxes = np.array(detections.data[ORIENTED_BOX_COORDINATES]).astype(int)
for detection_idx in range(len(detections)):
bbox = np.intp(detections.data.get(ORIENTED_BOX_COORDINATES)[detection_idx])
obb = obb_boxes[detection_idx]
color = resolve_color(
color=self.color,
detections=detections,
@ -277,7 +281,7 @@ class OrientedBoxAnnotator(BaseAnnotator):
else custom_color_lookup,
)
cv2.drawContours(scene, [bbox], 0, color.as_bgr(), self.thickness)
cv2.drawContours(scene, [obb], 0, color.as_bgr(), self.thickness)
return scene
@ -348,6 +352,7 @@ class MaskAnnotator(BaseAnnotator):
![mask-annotator-example](https://media.roboflow.com/
supervision-annotator-examples/mask-annotator-example-purple.png)
"""
assert isinstance(scene, np.ndarray)
if detections.mask is None:
return scene
@ -437,6 +442,7 @@ class PolygonAnnotator(BaseAnnotator):
![polygon-annotator-example](https://media.roboflow.com/
supervision-annotator-examples/polygon-annotator-example-purple.png)
"""
assert isinstance(scene, np.ndarray)
if detections.mask is None:
return scene
@ -523,6 +529,7 @@ class ColorAnnotator(BaseAnnotator):
![box-mask-annotator-example](https://media.roboflow.com/
supervision-annotator-examples/box-mask-annotator-example-purple.png)
"""
assert isinstance(scene, np.ndarray)
scene_with_boxes = scene.copy()
for detection_idx in range(len(detections)):
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
@ -618,6 +625,7 @@ class HaloAnnotator(BaseAnnotator):
![halo-annotator-example](https://media.roboflow.com/
supervision-annotator-examples/halo-annotator-example-purple.png)
"""
assert isinstance(scene, np.ndarray)
if detections.mask is None:
return scene
colored_mask = np.zeros_like(scene, dtype=np.uint8)
@ -717,6 +725,7 @@ class EllipseAnnotator(BaseAnnotator):
![ellipse-annotator-example](https://media.roboflow.com/
supervision-annotator-examples/ellipse-annotator-example-purple.png)
"""
assert isinstance(scene, np.ndarray)
for detection_idx in range(len(detections)):
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
color = resolve_color(
@ -808,6 +817,7 @@ class BoxCornerAnnotator(BaseAnnotator):
![box-corner-annotator-example](https://media.roboflow.com/
supervision-annotator-examples/box-corner-annotator-example-purple.png)
"""
assert isinstance(scene, np.ndarray)
for detection_idx in range(len(detections)):
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
color = resolve_color(
@ -897,6 +907,7 @@ class CircleAnnotator(BaseAnnotator):
![circle-annotator-example](https://media.roboflow.com/
supervision-annotator-examples/circle-annotator-example-purple.png)
"""
assert isinstance(scene, np.ndarray)
for detection_idx in range(len(detections)):
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
center = ((x1 + x2) // 2, (y1 + y2) // 2)
@ -994,6 +1005,7 @@ class DotAnnotator(BaseAnnotator):
![dot-annotator-example](https://media.roboflow.com/
supervision-annotator-examples/dot-annotator-example-purple.png)
"""
assert isinstance(scene, np.ndarray)
xy = detections.get_anchors_coordinates(anchor=self.position)
for detection_idx in range(len(detections)):
color = resolve_color(
@ -1115,6 +1127,7 @@ class LabelAnnotator(BaseAnnotator):
![label-annotator-example](https://media.roboflow.com/
supervision-annotator-examples/label-annotator-example-purple.png)
"""
assert isinstance(scene, np.ndarray)
font = cv2.FONT_HERSHEY_SIMPLEX
anchors_coordinates = detections.get_anchors_coordinates(
anchor=self.text_anchor
@ -1151,8 +1164,8 @@ class LabelAnnotator(BaseAnnotator):
if labels is not None:
text = labels[detection_idx]
elif detections[CLASS_NAME_DATA_FIELD] is not None:
text = detections[CLASS_NAME_DATA_FIELD][detection_idx]
elif CLASS_NAME_DATA_FIELD in detections.data:
text = detections.data[CLASS_NAME_DATA_FIELD][detection_idx]
elif detections.class_id is not None:
text = str(detections.class_id[detection_idx])
else:
@ -1331,6 +1344,7 @@ class RichLabelAnnotator(BaseAnnotator):
```
"""
assert isinstance(scene, Image.Image)
draw = ImageDraw.Draw(scene)
anchors_coordinates = detections.get_anchors_coordinates(
anchor=self.text_anchor
@ -1366,8 +1380,8 @@ class RichLabelAnnotator(BaseAnnotator):
if labels is not None:
text = labels[detection_idx]
elif detections[CLASS_NAME_DATA_FIELD] is not None:
text = detections[CLASS_NAME_DATA_FIELD][detection_idx]
elif CLASS_NAME_DATA_FIELD in detections.data:
text = detections.data[CLASS_NAME_DATA_FIELD][detection_idx]
elif detections.class_id is not None:
text = str(detections.class_id[detection_idx])
else:
@ -1563,6 +1577,7 @@ class BlurAnnotator(BaseAnnotator):
![blur-annotator-example](https://media.roboflow.com/
supervision-annotator-examples/blur-annotator-example-purple.png)
"""
assert isinstance(scene, np.ndarray)
image_height, image_width = scene.shape[:2]
clipped_xyxy = clip_boxes(
xyxy=detections.xyxy, resolution_wh=(image_width, image_height)
@ -1661,8 +1676,14 @@ class TraceAnnotator(BaseAnnotator):
![trace-annotator-example](https://media.roboflow.com/
supervision-annotator-examples/trace-annotator-example-purple.png)
"""
self.trace.put(detections)
assert isinstance(scene, np.ndarray)
if detections.tracker_id is None:
raise ValueError(
"The `tracker_id` field is missing in the provided detections."
" See more: https://supervision.roboflow.com/latest/how_to/track_objects"
)
self.trace.put(detections)
for detection_idx in range(len(detections)):
tracker_id = int(detections.tracker_id[detection_idx])
color = resolve_color(
@ -1715,9 +1736,9 @@ class HeatMapAnnotator(BaseAnnotator):
self.opacity = opacity
self.radius = radius
self.kernel_size = kernel_size
self.heat_mask = None
self.top_hue = top_hue
self.low_hue = low_hue
self.heat_mask: Optional[npt.NDArray[np.float32]] = None
@ensure_cv2_image_for_annotation
def annotate(self, scene: ImageType, detections: Detections) -> ImageType:
@ -1744,7 +1765,7 @@ class HeatMapAnnotator(BaseAnnotator):
heat_map_annotator = sv.HeatMapAnnotator()
video_info = sv.VideoInfo.from_video_path(video_path='...')
frames_generator = get_video_frames_generator(source_path='...')
frames_generator = sv.get_video_frames_generator(source_path='...')
with sv.VideoSink(target_path='...', video_info=video_info) as sink:
for frame in frames_generator:
@ -1759,12 +1780,20 @@ class HeatMapAnnotator(BaseAnnotator):
![heatmap-annotator-example](https://media.roboflow.com/
supervision-annotator-examples/heat-map-annotator-example-purple.png)
"""
assert isinstance(scene, np.ndarray)
if self.heat_mask is None:
self.heat_mask = np.zeros(scene.shape[:2])
self.heat_mask = np.zeros(scene.shape[:2], dtype=np.float32)
mask = np.zeros(scene.shape[:2])
for xy in detections.get_anchors_coordinates(self.position):
cv2.circle(mask, (int(xy[0]), int(xy[1])), self.radius, 1, -1)
x, y = int(xy[0]), int(xy[1])
cv2.circle(
img=mask,
center=(x, y),
radius=self.radius,
color=(1,),
thickness=-1, # fill
)
self.heat_mask = mask + self.heat_mask
temp = self.heat_mask.copy()
temp = self.low_hue - temp / temp.max() * (self.low_hue - self.top_hue)
@ -1832,6 +1861,7 @@ class PixelateAnnotator(BaseAnnotator):
![pixelate-annotator-example](https://media.roboflow.com/
supervision-annotator-examples/pixelate-annotator-example-10.png)
"""
assert isinstance(scene, np.ndarray)
image_height, image_width = scene.shape[:2]
clipped_xyxy = clip_boxes(
xyxy=detections.xyxy, resolution_wh=(image_width, image_height)
@ -1930,6 +1960,7 @@ class TriangleAnnotator(BaseAnnotator):
![triangle-annotator-example](https://media.roboflow.com/
supervision-annotator-examples/triangle-annotator-example.png)
"""
assert isinstance(scene, np.ndarray)
xy = detections.get_anchors_coordinates(anchor=self.position)
for detection_idx in range(len(detections)):
color = resolve_color(
@ -2042,7 +2073,7 @@ class RoundBoxAnnotator(BaseAnnotator):
![round-box-annotator-example](https://media.roboflow.com/
supervision-annotator-examples/round-box-annotator-example-purple.png)
"""
assert isinstance(scene, np.ndarray)
for detection_idx in range(len(detections)):
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
color = resolve_color(
@ -2115,7 +2146,7 @@ class PercentageBarAnnotator(BaseAnnotator):
border_color: Color = Color.BLACK,
position: Position = Position.TOP_CENTER,
color_lookup: ColorLookup = ColorLookup.CLASS,
border_thickness: int = None,
border_thickness: Optional[int] = None,
):
"""
Args:
@ -2127,7 +2158,7 @@ class PercentageBarAnnotator(BaseAnnotator):
position (Position): The anchor position of drawing the percentage bar.
color_lookup (ColorLookup): Strategy for mapping colors to annotations.
Options are `INDEX`, `CLASS`, `TRACK`.
border_thickness (int): The thickness of the border lines.
border_thickness (Optional[int]): The thickness of the border lines.
"""
self.height: int = height
self.width: int = width
@ -2186,9 +2217,9 @@ class PercentageBarAnnotator(BaseAnnotator):
![percentage-bar-example](https://media.roboflow.com/
supervision-annotator-examples/percentage-bar-annotator-example-purple.png)
"""
self.validate_custom_values(
custom_values=custom_values, detections_count=len(detections)
)
assert isinstance(scene, np.ndarray)
self.validate_custom_values(custom_values=custom_values, detections=detections)
anchors = detections.get_anchors_coordinates(anchor=self.position)
for detection_idx in range(len(detections)):
anchor = anchors[detection_idx]
@ -2199,11 +2230,11 @@ class PercentageBarAnnotator(BaseAnnotator):
)
border_width = border_coordinates[1][0] - border_coordinates[0][0]
value = (
custom_values[detection_idx]
if custom_values is not None
else detections.confidence[detection_idx]
)
if custom_values is not None:
value = custom_values[detection_idx]
else:
assert detections.confidence is not None # MyPy type hint
value = detections.confidence[detection_idx]
color = resolve_color(
color=self.color,
@ -2263,15 +2294,23 @@ class PercentageBarAnnotator(BaseAnnotator):
@staticmethod
def validate_custom_values(
custom_values: Optional[Union[np.ndarray, List[float]]], detections_count: int
custom_values: Optional[Union[np.ndarray, List[float]]], detections: Detections
) -> None:
if custom_values is not None:
if custom_values is None:
if detections.confidence is None:
raise ValueError(
"The provided detections do not contain confidence values. "
"Please provide `custom_values` or ensure that the detections "
"contain confidence values (e.g. by using a different model)."
)
else:
if not isinstance(custom_values, (np.ndarray, list)):
raise TypeError(
"custom_values must be either a numpy array or a list of floats."
)
if len(custom_values) != detections_count:
if len(custom_values) != len(detections):
raise ValueError(
"The length of custom_values must match the number of detections."
)
@ -2351,6 +2390,7 @@ class CropAnnotator(BaseAnnotator):
)
```
"""
assert isinstance(scene, np.ndarray)
crops = [
crop_image(image=scene, xyxy=xyxy) for xyxy in detections.xyxy.astype(int)
]
@ -2489,6 +2529,7 @@ class BackgroundOverlayAnnotator(BaseAnnotator):
![background-overlay-annotator-example](https://media.roboflow.com/
supervision-annotator-examples/background-color-annotator-example-purple.png)
"""
assert isinstance(scene, np.ndarray)
colored_mask = np.full_like(scene, self.color.as_bgr(), dtype=np.uint8)
cv2.addWeighted(

35
supervision/dataset/core.py
View File

@ -43,7 +43,10 @@ class BaseDataset(ABC):
@abstractmethod
def split(
self, split_ratio=0.8, random_state=None, shuffle: bool = True
self,
split_ratio: float = 0.8,
random_state: Optional[int] = None,
shuffle: bool = True,
) -> Tuple[BaseDataset, BaseDataset]:
pass
@ -174,18 +177,21 @@ class DetectionDataset(BaseDataset):
return True
def split(
self, split_ratio=0.8, random_state=None, shuffle: bool = True
self,
split_ratio: float = 0.8,
random_state: Optional[int] = None,
shuffle: bool = True,
) -> Tuple[DetectionDataset, DetectionDataset]:
"""
Splits the dataset into two parts (training and testing)
using the provided split_ratio.
Args:
split_ratio (float, optional): The ratio of the training
split_ratio (float): The ratio of the training
set to the entire dataset.
random_state (int, optional): The seed for the random number generator.
random_state (Optional[int]): The seed for the random number generator.
This is used for reproducibility.
shuffle (bool, optional): Whether to shuffle the data before splitting.
shuffle (bool): Whether to shuffle the data before splitting.
Returns:
Tuple[DetectionDataset, DetectionDataset]: A tuple containing
@ -396,7 +402,7 @@ class DetectionDataset(BaseDataset):
images_directory_path (str): Path to the directory containing the images.
annotations_directory_path (str): Path to the directory
containing the PASCAL VOC XML annotations.
force_masks (bool, optional): If True, forces masks to
force_masks (bool): If True, forces masks to
be loaded for all annotations, regardless of whether they are present.
Returns:
@ -455,10 +461,10 @@ class DetectionDataset(BaseDataset):
containing the YOLO annotation files.
data_yaml_path (str): The path to the data
YAML file containing class information.
force_masks (bool, optional): If True, forces
force_masks (bool): If True, forces
masks to be loaded for all annotations,
regardless of whether they are present.
is_obb (bool, optional): If True, loads the annotations in OBB format.
is_obb (bool): If True, loads the annotations in OBB format.
OBB annotations are defined as `[class_id, x, y, x, y, x, y, x, y]`,
where pairs of [x, y] are box corners.
@ -565,7 +571,7 @@ class DetectionDataset(BaseDataset):
images_directory_path (str): The path to the
directory containing the images.
annotations_path (str): The path to the json annotation files.
force_masks (bool, optional): If True,
force_masks (bool): If True,
forces masks to be loaded for all annotations,
regardless of whether they are present.
@ -777,18 +783,21 @@ class ClassificationDataset(BaseDataset):
return True
def split(
self, split_ratio=0.8, random_state=None, shuffle: bool = True
self,
split_ratio: float = 0.8,
random_state: Optional[int] = None,
shuffle: bool = True,
) -> Tuple[ClassificationDataset, ClassificationDataset]:
"""
Splits the dataset into two parts (training and testing)
using the provided split_ratio.
Args:
split_ratio (float, optional): The ratio of the training
split_ratio (float): The ratio of the training
set to the entire dataset.
random_state (int, optional): The seed for the
random_state (Optional[int]): The seed for the
random number generator. This is used for reproducibility.
shuffle (bool, optional): Whether to shuffle the data before splitting.
shuffle (bool): Whether to shuffle the data before splitting.
Returns:
Tuple[ClassificationDataset, ClassificationDataset]: A tuple containing

2
supervision/dataset/formats/pascal_voc.py
View File

@ -147,7 +147,7 @@ def load_pascal_voc_annotations(
images_directory_path (str): The path to the directory containing the images.
annotations_directory_path (str): The path to the directory containing the
PASCAL VOC annotation files.
force_masks (bool, optional): If True, forces masks to be loaded for all
force_masks (bool): If True, forces masks to be loaded for all
annotations, regardless of whether they are present.
Returns:

4
supervision/dataset/formats/yolo.py
View File

@ -138,9 +138,9 @@ def load_yolo_annotations(
containing the YOLO annotation files.
data_yaml_path (str): The path to the data
YAML file containing class information.
force_masks (bool, optional): If True, forces masks to be loaded
force_masks (bool): If True, forces masks to be loaded
for all annotations, regardless of whether they are present.
is_obb (bool, optional): If True, loads the annotations in OBB format.
is_obb (bool): If True, loads the annotations in OBB format.
OBB annotations are defined as `[class_id, x, y, x, y, x, y, x, y]`,
where pairs of [x, y] are box corners.

8
supervision/detection/core.py
View File

@ -1158,10 +1158,10 @@ class Detections:
from a segmentation model, the IoU mask is applied. Otherwise, box IoU is used.
Args:
threshold (float, optional): The intersection-over-union threshold
threshold (float): The intersection-over-union threshold
to use for non-maximum suppression. I'm the lower the value the more
restrictive the NMS becomes. Defaults to 0.5.
class_agnostic (bool, optional): Whether to perform class-agnostic
class_agnostic (bool): Whether to perform class-agnostic
non-maximum suppression. If True, the class_id of each detection
will be ignored. Defaults to False.
@ -1213,9 +1213,9 @@ class Detections:
Perform non-maximum merging on the current set of object detections.
Args:
threshold (float, optional): The intersection-over-union threshold
threshold (float): The intersection-over-union threshold
to use for non-maximum merging. Defaults to 0.5.
class_agnostic (bool, optional): Whether to perform class-agnostic
class_agnostic (bool): Whether to perform class-agnostic
non-maximum merging. If True, the class_id of each detection
will be ignored. Defaults to False.

10
supervision/detection/overlap_filter.py
View File

@ -55,9 +55,9 @@ def mask_non_max_suppression(
masks (np.ndarray): A 3D array of binary masks corresponding to the predictions.
Shape: `(N, H, W)`, where N is the number of predictions, and H, W are the
dimensions of each mask.
iou_threshold (float, optional): The intersection-over-union threshold
iou_threshold (float): The intersection-over-union threshold
to use for non-maximum suppression.
mask_dimension (int, optional): The dimension to which the masks should be
mask_dimension (int): The dimension to which the masks should be
resized before computing IOU values. Defaults to 640.
Returns:
@ -103,7 +103,7 @@ def box_non_max_suppression(
predictions (np.ndarray): An array of object detection predictions in
the format of `(x_min, y_min, x_max, y_max, score)`
or `(x_min, y_min, x_max, y_max, score, class)`.
iou_threshold (float, optional): The intersection-over-union threshold
iou_threshold (float): The intersection-over-union threshold
to use for non-maximum suppression.
Returns:
@ -158,7 +158,7 @@ def group_overlapping_boxes(
predictions (npt.NDArray[np.float64]): An array of shape `(n, 5)` containing
the bounding boxes coordinates in format `[x1, y1, x2, y2]`
and the confidence scores.
iou_threshold (float, optional): The intersection-over-union threshold
iou_threshold (float): The intersection-over-union threshold
to use for non-maximum suppression. Defaults to 0.5.
Returns:
@ -202,7 +202,7 @@ def box_non_max_merge(
containing the bounding boxes coordinates in format `[x1, y1, x2, y2]`,
the confidence scores and class_ids. Omit class_id column to allow
detections of different classes to be merged.
iou_threshold (float, optional): The intersection-over-union threshold
iou_threshold (float): The intersection-over-union threshold
to use for non-maximum suppression. Defaults to 0.5.
Returns:

2
supervision/detection/tools/polygon_zone.py
View File

@ -147,7 +147,7 @@ class PolygonZoneAnnotator:
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
label (Optional[str]): A label for the count of detected objects
within the polygon zone (default: None)
Returns:

2
supervision/detection/utils.py
View File

@ -106,7 +106,7 @@ def mask_iou_batch(
Args:
masks_true (np.ndarray): 3D `np.ndarray` representing ground-truth masks.
masks_detection (np.ndarray): 3D `np.ndarray` representing detection masks.
memory_limit (int, optional): memory limit in MB, default is 1024 * 5 MB (5GB).
memory_limit (int): memory limit in MB, default is 1024 * 5 MB (5GB).
Returns:
np.ndarray: Pairwise IoU of masks from `masks_true` and `masks_detection`.

14
supervision/draw/utils.py
View File

@ -142,7 +142,7 @@ def draw_polygon(
scene (np.ndarray): The scene to draw the polygon on.
polygon (np.ndarray): The polygon to be drawn, given as a list of vertices.
color (Color): The color of the polygon.
thickness (int, optional): The thickness of the polygon lines, by default 2.
thickness (int): The thickness of the polygon lines, by default 2.
Returns:
np.ndarray: The scene with the polygon drawn on it.
@ -172,14 +172,14 @@ def draw_text(
text (str): The text to be drawn.
text_anchor (Point): The anchor point for the text, represented as a
Point object with x and y attributes.
text_color (Color, optional): The color of the text. Defaults to black.
text_scale (float, optional): The scale of the text. Defaults to 0.5.
text_thickness (int, optional): The thickness of the text. Defaults to 1.
text_padding (int, optional): The amount of padding to add around the text
text_color (Color): The color of the text. Defaults to black.
text_scale (float): The scale of the text. Defaults to 0.5.
text_thickness (int): The thickness of the text. Defaults to 1.
text_padding (int): The amount of padding to add around the text
when drawing a rectangle in the background. Defaults to 10.
text_font (int, optional): The font to use for the text.
text_font (int): The font to use for the text.
Defaults to cv2.FONT_HERSHEY_SIMPLEX.
background_color (Color, optional): The color of the background rectangle,
background_color (Optional[Color]): The color of the background rectangle,
if one is to be drawn. Defaults to None.
Returns:

34
supervision/keypoint/annotators.py
View File

@ -34,8 +34,8 @@ class VertexAnnotator(BaseKeyPointAnnotator):
) -> None:
"""
Args:
color (Color, optional): The color to use for annotating key points.
radius (int, optional): The radius of the circles used to represent the key
color (Color): The color to use for annotating key points.
radius (int): The radius of the circles used to represent the key
points.
"""
self.color = color
@ -78,6 +78,7 @@ class VertexAnnotator(BaseKeyPointAnnotator):
![vertex-annotator-example](https://media.roboflow.com/
supervision-annotator-examples/vertex-annotator-example.png)
"""
assert isinstance(scene, np.ndarray)
if len(key_points) == 0:
return scene
@ -108,8 +109,8 @@ class EdgeAnnotator(BaseKeyPointAnnotator):
) -> None:
"""
Args:
color (Color, optional): The color to use for the edges.
thickness (int, optional): The thickness of the edges.
color (Color): The color to use for the edges.
thickness (int): The thickness of the edges.
edges (Optional[List[Tuple[int, int]]]): The edges to draw.
If set to `None`, will attempt to select automatically.
"""
@ -155,6 +156,7 @@ class EdgeAnnotator(BaseKeyPointAnnotator):
![edge-annotator-example](https://media.roboflow.com/
supervision-annotator-examples/edge-annotator-example.png)
"""
assert isinstance(scene, np.ndarray)
if len(key_points) == 0:
return scene
@ -202,16 +204,16 @@ class VertexLabelAnnotator:
):
"""
Args:
color (Union[Color, List[Color]], optional): The color to use for each
color (Union[Color, List[Color]]): The color to use for each
keypoint label. If a list is provided, the colors will be used in order
for each keypoint.
text_color (Union[Color, List[Color]], optional): The color to use
text_color (Union[Color, List[Color]]): The color to use
for the labels. If a list is provided, the colors will be used in order
for each keypoint.
text_scale (float, optional): The scale of the text.
text_thickness (int, optional): The thickness of the text.
text_padding (int, optional): The padding around the text.
border_radius (int, optional): The radius of the rounded corners of the
text_scale (float): The scale of the text.
text_thickness (int): The thickness of the text.
text_padding (int): The padding around the text.
border_radius (int): The radius of the rounded corners of the
boxes. Set to a high value to produce circles.
"""
self.border_radius: int = border_radius
@ -222,7 +224,10 @@ class VertexLabelAnnotator:
self.text_padding: int = text_padding
def annotate(
self, scene: ImageType, key_points: KeyPoints, labels: List[str] = None
self,
scene: ImageType,
key_points: KeyPoints,
labels: Optional[List[str]] = None,
) -> ImageType:
"""
A class that draws labels of skeleton vertices on images. It uses specified key
@ -234,7 +239,7 @@ class VertexLabelAnnotator:
`PIL.Image.Image`.
key_points (KeyPoints): A collection of key points where each key point
consists of x and y coordinates.
labels (List[str], optional): A list of labels to be displayed on the
labels (Optional[List[str]]): A list of labels to be displayed on the
annotated image. If not provided, keypoint indices will be used.
Returns:
@ -305,6 +310,7 @@ class VertexLabelAnnotator:
![vertex-label-annotator-custom-example](https://media.roboflow.com/
supervision-annotator-examples/vertex-label-annotator-custom-example.png)
"""
assert isinstance(scene, np.ndarray)
font = cv2.FONT_HERSHEY_SIMPLEX
skeletons_count, points_count, _ = key_points.xy.shape
@ -400,7 +406,7 @@ class VertexLabelAnnotator:
@staticmethod
def preprocess_and_validate_labels(
labels: Optional[List[str]], points_count: int, skeletons_count: int
) -> np.array:
) -> np.ndarray:
if labels and len(labels) != points_count:
raise ValueError(
f"Number of labels ({len(labels)}) must match number of key points "
@ -416,7 +422,7 @@ class VertexLabelAnnotator:
colors: Optional[Union[Color, List[Color]]],
points_count: int,
skeletons_count: int,
) -> np.array:
) -> np.ndarray:
if isinstance(colors, list) and len(colors) != points_count:
raise ValueError(
f"Number of colors ({len(colors)}) must match number of key points "

14
supervision/keypoint/core.py
View File

@ -88,15 +88,17 @@ class KeyPoints:
```
Attributes:
xy (np.ndarray): An array of shape `(n, 2)` containing
the bounding boxes coordinates in format `[x1, y1]`
xy (np.ndarray): An array of shape `(n, m, 2)` containing
`n` detected objects, each composed of `m` equally-sized
sets of keypoints, where each point is `[x, y]`.
confidence (Optional[np.ndarray]): An array of shape
`(n,)` containing the confidence scores of the keypoint keypoints.
`(n, m)` containing the confidence scores of each keypoint.
class_id (Optional[np.ndarray]): An array of shape
`(n,)` containing the class ids of the keypoint keypoints.
`(n,)` containing the class ids of the detected objects.
data (Dict[str, Union[np.ndarray, List]]): A dictionary containing additional
data where each key is a string representing the data type, and the value
is either a NumPy array or a list of corresponding data.
is either a NumPy array or a list of corresponding data of length `n`
(one entry per detected object).
""" # noqa: E501 // docs
xy: npt.NDArray[np.float32]
@ -132,7 +134,7 @@ class KeyPoints:
]:
"""
Iterates over the Keypoint object and yield a tuple of
`(xy, confidence, class_id, data)` for each keypoint detection.
`(xy, confidence, class_id, data)` for each object detection.
"""
for i in range(len(self.xy)):
yield (

2
supervision/metrics/detection.py
View File

@ -806,7 +806,7 @@ class MeanAveragePrecision:
prediction_confidence (np.ndarray): Objectness value from 0-1.
prediction_class_ids (np.ndarray): Predicted object classes.
true_class_ids (np.ndarray): True object classes.
eps (float, optional): Small value to prevent division by zero.
eps (float): Small value to prevent division by zero.
Returns:
np.ndarray: Average precision for different IoU levels.

10
supervision/tracker/byte_tracker/core.py
View File

@ -197,19 +197,19 @@ class ByteTrack:
</video>
Parameters:
track_activation_threshold (float, optional): Detection confidence threshold
track_activation_threshold (float): Detection confidence threshold
for track activation. Increasing track_activation_threshold improves accuracy
and stability but might miss true detections. Decreasing it increases
completeness but risks introducing noise and instability.
lost_track_buffer (int, optional): Number of frames to buffer when a track is lost.
lost_track_buffer (int): Number of frames to buffer when a track is lost.
Increasing lost_track_buffer enhances occlusion handling, significantly
reducing the likelihood of track fragmentation or disappearance caused
by brief detection gaps.
minimum_matching_threshold (float, optional): Threshold for matching tracks with detections.
minimum_matching_threshold (float): Threshold for matching tracks with detections.
Increasing minimum_matching_threshold improves accuracy but risks fragmentation.
Decreasing it improves completeness but risks false positives and drift.
frame_rate (int, optional): The frame rate of the video.
minimum_consecutive_frames (int, optional): Number of consecutive frames that an object must
frame_rate (int): The frame rate of the video.
minimum_consecutive_frames (int): Number of consecutive frames that an object must
be tracked before it is considered a 'valid' track.
Increasing minimum_consecutive_frames prevents the creation of accidental tracks from
false detection or double detection, but risks missing shorter tracks.

8
supervision/utils/image.py
View File

@ -158,7 +158,7 @@ def resize_image(
accepting either `numpy.ndarray` or `PIL.Image.Image`.
resolution_wh (Tuple[int, int]): The target resolution as
`(width, height)`.
keep_aspect_ratio (bool, optional): Flag to maintain the image's original
keep_aspect_ratio (bool): Flag to maintain the image's original
aspect ratio. Defaults to `False`.
Returns:
@ -380,9 +380,9 @@ class ImageSink:
Args:
target_dir_path (str): The target directory where images will be saved.
overwrite (bool, optional): Whether to overwrite the existing directory.
overwrite (bool): Whether to overwrite the existing directory.
Defaults to False.
image_name_pattern (str, optional): The image file name pattern.
image_name_pattern (str): The image file name pattern.
Defaults to "image_{:05d}.png".
Examples:
@ -419,7 +419,7 @@ class ImageSink:
Args:
image (np.ndarray): The image to be saved. The image must be in BGR color
format.
image_name (str, optional): The name to use for the saved image.
image_name (Optional[str]): The name to use for the saved image.
If not provided, a name will be
generated using the `image_name_pattern`.
"""

29
supervision/utils/internal.py
View File

@ -2,7 +2,7 @@ import functools
import inspect
import os
import warnings
from typing import Any, Callable, Set
from typing import Any, Callable, Generic, Optional, Set, TypeVar
class SupervisionWarnings(Warning):
@ -56,9 +56,9 @@ def deprecated_parameter(
Parameters:
old_parameter (str): The name of the deprecated parameter.
new_parameter (str): The name of the parameter that should be used instead.
map_function (Callable, optional): A function used to map the value of the old
map_function (Callable): A function used to map the value of the old
parameter to the new parameter. Defaults to the identity function.
warning_message (str, optional): The warning message to be displayed when the
warning_message (str): The warning message to be displayed when the
deprecated parameter is used. Defaults to a generic warning message with
placeholders for the old parameter, new parameter, and function name.
**message_kwargs: Additional keyword arguments that can be used to customize
@ -122,7 +122,10 @@ def deprecated(reason: str):
return decorator
class classproperty(property):
T = TypeVar("T")
class classproperty(Generic[T]):
"""
A decorator that combines @classmethod and @property.
It allows a method to be accessed as a property of the class,
@ -134,17 +137,27 @@ class classproperty(property):
...
"""
def __get__(self, owner_self: object, owner_cls: type) -> object:
def __init__(self, fget: Callable[..., T]):
"""
Args:
The function that is called when the property is accessed.
"""
self.fget = fget
def __get__(self, owner_self: Any, owner_cls: Optional[type] = None) -> T:
"""
Override the __get__ method to return the result of the function call.
Args:
owner_self: The instance through which the attribute was accessed, or None.
owner_cls: The class through which the attribute was accessed.
owner_self: The instance through which the attribute was accessed, or None.
Irrelevant for class properties.
owner_cls: The class through which the attribute was accessed.
Returns:
The result of calling the function stored in 'fget' with 'owner_cls'.
The result of calling the function stored in 'fget' with 'owner_cls'.
"""
if self.fget is None:
raise AttributeError("unreadable attribute")
return self.fget(owner_cls)

2
supervision/utils/video.py
View File

@ -19,7 +19,7 @@ class VideoInfo:
width (int): width of the video in pixels
height (int): height of the video in pixels
fps (int): frames per second of the video
total_frames (int, optional): total number of frames in the video,
total_frames (Optional[int]): total number of frames in the video,
default is None
Examples:

2
test/dataset/formats/test_yolo.py
View File

@ -13,7 +13,7 @@ from supervision.dataset.formats.yolo import (
from supervision.detection.core import Detections
def _mock_simple_mask(resolution_wh: Tuple[int, int], box: List[int]) -> np.array:
def _mock_simple_mask(resolution_wh: Tuple[int, int], box: List[int]) -> np.ndarray:
x_min, y_min, x_max, y_max = box
mask = np.full(resolution_wh, False, dtype=bool)
mask[y_min:y_max, x_min:x_max] = True