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

Merge pull request #171 from hardikdava/advanced_box_annotator

Browse Source 
cornered box annotator
This commit is contained in:
Piotr Skalski authored and GitHub committed 2023-07-23 20:44:29 +02:00
2 parents 8066125f03 cc2e08ac72
commit 1d28b8ae26
8 files changed
+871 -30
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

setup.py
+2 -1
View File
@@ -27,7 +27,8 @@ setuptools.setup(
'numpy>=1.20.0',
'opencv-python',
'matplotlib',
'pyyaml'
'pyyaml',
'pillow'
],
packages=find_packages(exclude=("tests",)),
extras_require={
supervision/__init__.py
+11
View File
@@ -1,5 +1,16 @@
__version__ = "0.11.1"
from supervision.annotators.composable import DetectionAnnotator, SegmentationAnnotator
from supervision.annotators.core import (
BoxCornerAnnotator,
BoxLineAnnotator,
BoxMaskAnnotator,
EllipseAnnotator,
LabelAdvancedAnnotator,
LabelAnnotator,
MaskAnnotator,
build_label_formatter,
)
from supervision.classification.core import Classifications
from supervision.dataset.core import (
BaseDataset,
supervision/annotators/composable.py
+84 -8
View File
@@ -1,18 +1,94 @@
from abc import ABC
from typing import List
from typing import List, Optional
import numpy as np
from supervision.annotators.core import BaseAnnotator
from supervision.annotators.core import BoxAnnotator, LabelAnnotator, MaskAnnotator
from supervision.detection.core import Detections
class ComposableAnnotator(ABC):
def __init__(self, annotators: List[BaseAnnotator]):
self.annotators = annotators
def __init__(
self,
):
self.annotators = []
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
annotated_image = scene
def annotate(
self,
scene: np.ndarray,
detections: Detections,
labels: Optional[List[str]] = None,
) -> np.ndarray:
for annotator in self.annotators:
annotated_image = annotator.annotate(scene=scene, detections=detections)
return annotated_image
if isinstance(annotator, LabelAnnotator):
scene = annotator.annotate(
scene=scene,
detections=detections,
labels=labels,
)
else:
scene = annotator.annotate(scene=scene, detections=detections)
return scene
class DetectionAnnotator(ComposableAnnotator):
"""
Highlevel API for drawing Object Detection output. This will use Box and Label Annotators
Example:
```python
>>> import supervision as sv
>>> classes = ['person', ...]
>>> image = ...
>>> detections = sv.Detections(...)
>>> detection_annotator = sv.DetectionAnnotator()
>>> annotated_frame = detection_annotator.annotate(
... scene=image.copy(),
... detections=detections
... )
```
"""
def __init__(
self,
color_by_track: bool = False,
skip_label: bool = False,
):
super().__init__()
self.annotators = [BoxAnnotator(color_by_track=color_by_track)]
if not skip_label:
self.annotators.append(LabelAnnotator(color_by_track=color_by_track))
class SegmentationAnnotator(ComposableAnnotator):
"""
High level API for drawing segmentation mask, bounding box and labels on an image using provided detections.
Example:
```python
>>> import supervision as sv
>>> classes = ['person', ...]
>>> image = ...
>>> detections = sv.Detections(...)
>>> segmentation_annotator = sv.SegmentationAnnotator()
>>> annotated_frame = segmentation_annotator.annotate(
... scene=image.copy(),
... detections=detections
... )
```
"""
def __init__(
self,
color_by_track: bool = False,
skip_label: bool = False,
):
super().__init__()
self.annotators = [
BoxAnnotator(color_by_track=color_by_track),
MaskAnnotator(color_by_track=color_by_track),
]
if not skip_label:
self.annotators.append(LabelAnnotator(color_by_track=color_by_track))
supervision/annotators/core.py
+657 -19
View File
@@ -1,7 +1,10 @@
import os.path
from abc import ABC, abstractmethod
from typing import Union
from typing import Callable, List, Optional, Union
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from supervision.detection.core import Detections
from supervision.draw.color import Color, ColorPalette
@@ -12,12 +15,85 @@ class BaseAnnotator(ABC):
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
pass
class BoxAnnotator(BaseAnnotator):
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
@staticmethod
def resolve_annotation_color(
color: Union[Color, ColorPalette], by_track: bool, detections: Detections
) -> Color:
pass
class BoxLineAnnotator(BaseAnnotator):
"""
Basic bounding box annotation class
"""
def __init__(
self,
color: Union[Color, ColorPalette] = ColorPalette.default(),
thickness: int = 2,
color_by_track: bool = False,
):
self.color: Union[Color, ColorPalette] = color
self.thickness: int = thickness
self.color_by_track = color_by_track
def annotate(
self, scene: np.ndarray, detections: Detections, **kwargs
) -> np.ndarray:
"""
Draws bounding boxes on the frame using the detections provided.
Args:
scene (np.ndarray): The image on which the bounding boxes will be drawn
detections (Detections): The detections for which the bounding boxes will be drawn
Returns:
np.ndarray: The image with the bounding boxes drawn on it.
Example:
```python
>>> import supervision as sv
>>> image = ...
>>> detections = sv.Detections(...)
>>> box_line_annotator = sv.BoxLineAnnotator()
>>> annotated_frame = box_line_annotator.annotate(
... scene=image.copy(),
... detections=detections
... )
```
"""
for i in range(len(detections)):
x1, y1, x2, y2 = detections.xyxy[i].astype(int)
if self.color_by_track:
tracker_id = (
detections.tracker_id[i]
if detections.tracker_id is not None
else None
)
idx = tracker_id if tracker_id is not None else i
else:
class_id = (
detections.class_id[i] if detections.class_id is not None else None
)
idx = class_id if class_id is not None else i
color = (
self.color.by_idx(idx)
if isinstance(self.color, ColorPalette)
else self.color
)
cv2.rectangle(
img=scene,
pt1=(x1, y1),
pt2=(x2, y2),
color=color.as_bgr(),
thickness=self.thickness,
)
return scene
class MaskAnnotator(BaseAnnotator):
"""
A class for overlaying masks on an image using detections provided.
@@ -31,11 +107,15 @@ class MaskAnnotator(BaseAnnotator):
self,
color: Union[Color, ColorPalette] = ColorPalette.default(),
opacity: float = 0.5,
color_by_track: bool = False,
):
self.color: Union[Color, ColorPalette] = color
self.opacity = opacity
self.color_by_track = color_by_track
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
def annotate(
self, scene: np.ndarray, detections: Detections, **kwargs
) -> np.ndarray:
"""
Overlays the masks on the given image based on the provided detections, with a specified opacity.
@@ -45,15 +125,38 @@ class MaskAnnotator(BaseAnnotator):
Returns:
np.ndarray: The image with the masks overlaid
Example:
```python
>>> import supervision as sv
>>> classes = ['person', ...]
>>> image = ...
>>> detections = sv.Detections(...)
>>> mask_annotator = sv.MaskAnnotator()
>>> annotated_frame = mask_annotator.annotate(
... scene=image.copy(),
... detections=detections
... )
```
"""
if detections.mask is None:
return scene
for i in np.flip(np.argsort(detections.area)):
class_id = (
detections.class_id[i] if detections.class_id is not None else None
)
idx = class_id if class_id is not None else i
if self.color_by_track:
tracker_id = (
detections.tracker_id[i]
if detections.tracker_id is not None
else None
)
idx = tracker_id if tracker_id is not None else i
else:
class_id = (
detections.class_id[i] if detections.class_id is not None else None
)
idx = class_id if class_id is not None else i
color = (
self.color.by_idx(idx)
if isinstance(self.color, ColorPalette)
@@ -73,21 +176,556 @@ class MaskAnnotator(BaseAnnotator):
return scene
def default_label_formatter(
detections: Detections,
) -> List[str]:
return [str(class_id) for class_id in detections.class_id]
def build_label_formatter(classes: List[str]) -> Callable[[Detections], List[str]]:
def default_label_formatter(detections: Detections) -> List[str]:
return [classes[class_id] for class_id in detections.class_id]
return default_label_formatter
class LabelAnnotator(BaseAnnotator):
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
pass
"""
A class for putting text on an image using provided detections.
Attributes:
color (Union[Color, ColorPalette]): The color to text on the image, can be a single color or a color palette
"""
class PillowLabelAnnotator(BaseAnnotator):
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
pass
def __init__(
self,
color: Union[Color, ColorPalette] = ColorPalette.default(),
thickness: int = 2,
text_color: Color = Color.black(),
text_scale: float = 0.5,
text_thickness: int = 1,
text_padding: int = 10,
color_by_track: bool = False,
classes: Optional[List[str]] = None,
label_formatter: Optional[
Callable[[Detections], List[str]]
] = default_label_formatter,
):
"""
Args:
color_by_track: pick color by tracker id
classes: Optional list of class name
label_formatter: Optional callback function for label generator, avoided if classes is provided
"""
self.color: Union[Color, ColorPalette] = color
self.text_color: Color = text_color
self.text_scale: float = text_scale
self.text_thickness: int = text_thickness
self.text_padding: int = text_padding
self.color_by_track = color_by_track
self.classes = classes
self.label_formatter = label_formatter
def annotate(
self,
scene: np.ndarray,
detections: Detections,
) -> np.ndarray:
"""
Draws text on the frame using the detections provided and label.
class TrackAnnotator(BaseAnnotator):
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
pass
Args:
scene (np.ndarray): The image on which the bounding boxes will be drawn
detections (Detections): The detections for which the bounding boxes will be drawn
Returns:
np.ndarray: The image with the bounding boxes drawn on it
Example:
```python
>>> import supervision as sv
>>> classes = ['person', ...]
>>> image = ...
>>> detections = sv.Detections(...)
>>> label_annotator = sv.LabelAnnotator()
>>> labels = [
... f"{classes[class_id]} {confidence:0.2f}"
... for _, _, confidence, class_id, _
... in detections
... ]
>>> annotated_frame = label_annotator.annotate(
... scene=image.copy(),
... detections=detections,
... )
```
"""
font = cv2.FONT_HERSHEY_SIMPLEX
labels = None
if self.classes:
label_formatter = build_label_formatter(self.classes)
labels = label_formatter(detections)
else:
labels = self.label_formatter(detections=detections)
for i in range(len(detections)):
x1, y1, x2, y2 = detections.xyxy[i].astype(int)
if self.color_by_track:
tracker_id = (
detections.tracker_id[i]
if detections.tracker_id is not None
else None
)
idx = tracker_id if tracker_id is not None else i
else:
class_id = (
detections.class_id[i] if detections.class_id is not None else None
)
idx = class_id if class_id is not None else i
color = (
self.color.by_idx(idx)
if isinstance(self.color, ColorPalette)
else self.color
)
text = (
f"{idx}"
if (labels is None or len(detections) != len(labels))
else labels[i]
)
text_width, text_height = cv2.getTextSize(
text=text,
fontFace=font,
fontScale=self.text_scale,
thickness=self.text_thickness,
)[0]
text_x = x1 + self.text_padding
text_y = y1 - self.text_padding
text_background_x1 = x1
text_background_y1 = y1 - 2 * self.text_padding - text_height
text_background_x2 = x1 + 2 * self.text_padding + text_width
text_background_y2 = y1
cv2.rectangle(
img=scene,
pt1=(text_background_x1, text_background_y1),
pt2=(text_background_x2, text_background_y2),
color=color.as_bgr(),
thickness=cv2.FILLED,
)
cv2.putText(
img=scene,
text=text,
org=(text_x, text_y),
fontFace=font,
fontScale=self.text_scale,
color=self.text_color.as_rgb(),
thickness=self.text_thickness,
lineType=cv2.LINE_AA,
)
return scene
class BoxMaskAnnotator(BaseAnnotator):
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
pass
def __init__(
self,
color: Union[Color, ColorPalette] = ColorPalette.default(),
opacity: float = 0.5,
color_by_track: bool = False,
):
self.color: Union[Color, ColorPalette] = color
self.opacity = opacity
self.color_by_track = color_by_track
def annotate(
self,
scene: np.ndarray,
detections: Detections,
) -> np.ndarray:
"""
Overlays the rectangle masks on the given image based on the provided detections, with a specified opacity.
Args:
scene (np.ndarray): The image on which the masks will be overlaid
detections (Detections): The detections for which the masks will be overlaid
Returns:
np.ndarray: The image with the masks overlaid
Example:
```python
>>> import supervision as sv
>>> classes = ['person', ...]
>>> image = ...
>>> detections = sv.Detections(...)
>>> box_mask_annotator = sv.MaskAnnotator()
>>> annotated_frame = box_mask_annotator.annotate(
... scene=image.copy(),
... detections=detections
... )
```
"""
overlay_img = np.zeros_like(scene, np.uint8)
for i in range(len(detections)):
x1, y1, x2, y2 = detections.xyxy[i].astype(int)
if self.color_by_track:
tracker_id = (
detections.tracker_id[i]
if detections.tracker_id is not None
else None
)
idx = tracker_id if tracker_id is not None else i
else:
class_id = (
detections.class_id[i] if detections.class_id is not None else None
)
idx = class_id if class_id is not None else i
color = (
self.color.by_idx(idx)
if isinstance(self.color, ColorPalette)
else self.color
)
cv2.rectangle(
img=overlay_img,
pt1=(x1, y1),
pt2=(x2, y2),
color=color.as_bgr(),
thickness=-1,
)
mask = overlay_img.astype(bool)
scene[mask] = cv2.addWeighted(
scene, self.opacity, overlay_img, 1 - self.opacity, 0
)[mask]
return scene
def default_label_formatter(detections: Detections) -> List[str]:
return [str(class_id) for class_id in detections.class_id]
class LabelAdvancedAnnotator(BaseAnnotator):
def __init__(
self,
color: Union[Color, ColorPalette] = ColorPalette.default(),
text_color: Color = Color.black(),
text_padding: int = 20,
color_by_track: bool = False,
font: Optional[str] = None,
font_size: Optional[int] = 15,
):
if font and os.path.exists(font):
self.font = ImageFont.truetype(font, font_size)
else:
self.font = ImageFont.load_default()
self.color: Union[Color, ColorPalette] = color
self.text_color: Color = text_color
self.text_padding: int = text_padding
self.color_by_track = color_by_track
def annotate(
self,
scene: np.ndarray,
detections: Detections,
labels: Optional[List[str]] = None,
) -> np.ndarray:
"""
Draws text on the frame using the detections provided and label.
Args:
scene (np.ndarray): The image on which the bounding boxes will be drawn
detections (Detections): The detections for which the bounding boxes will be drawn
labels (Optional[List[str]]): An optional list of labels corresponding to each detection. If `labels` are not provided, corresponding `class_id` will be used as label.
Returns:
np.ndarray: The image with the bounding boxes drawn on it
Example:
```python
>>> import supervision as sv
>>> classes = ['person', ...]
>>> image = ...
>>> detections = sv.Detections(...)
>>> pil_label_annotator = sv.LabelAdvancedAnnotator()
>>> labels = [
... f"{classes[class_id]} {confidence:0.2f}"
... for _, _, confidence, class_id, _
... in detections
... ]
>>> annotated_frame = pil_label_annotator.annotate(
... scene=image.copy(),
... detections=detections,
... labels=labels,
... )
```
"""
pil_image = Image.fromarray(scene)
draw = ImageDraw.Draw(pil_image)
text_color = "#fff"
for i in range(len(detections)):
x1, y1, x2, y2 = detections.xyxy[i].astype(int)
if self.color_by_track:
tracker_id = (
detections.tracker_id[i]
if detections.tracker_id is not None
else None
)
idx = tracker_id if tracker_id is not None else i
else:
class_id = (
detections.class_id[i] if detections.class_id is not None else None
)
idx = class_id if class_id is not None else i
color = (
self.color.by_idx(idx)
if isinstance(self.color, ColorPalette)
else self.color
)
text = (
f"{idx}"
if (labels is None or len(detections) != len(labels))
else labels[i]
)
text_bbox = draw.textbbox((x1, y1), text, font=self.font)
text_height = text_bbox[3] - text_bbox[1]
text_width = text_bbox[2] - text_bbox[0]
text_x = x1 + self.text_padding / 2
text_y = y1 - self.text_padding / 2 - text_height
text_background_x1 = x1
text_background_y1 = y1 - self.text_padding / 2 - text_height
text_background_x2 = x1 + 2 * self.text_padding / 2 + text_width
text_background_y2 = y1 # correct
draw.rectangle(
(
text_background_x1,
text_background_y1,
text_background_x2,
text_background_y2,
),
fill=color.as_bgr(),
)
draw.text((text_x, text_y), text, font=self.font, fill=text_color)
scene = np.asarray(pil_image)
return scene
class BoxCornerAnnotator(BaseAnnotator):
def __init__(
self,
color: Union[Color, ColorPalette] = ColorPalette.default(),
thickness: int = 2,
color_by_track: bool = False,
):
self.color: Union[Color, ColorPalette] = color
self.thickness: int = thickness
self.color_by_track = color_by_track
def annotate(
self,
scene: np.ndarray,
detections: Detections,
):
"""
Draws cornered bounding boxes on the frame using the detections provided.
Args:
scene (np.ndarray): The image on which the bounding boxes will be drawn
detections (Detections): The detections for which the bounding boxes will be drawn
Returns:
np.ndarray: The image with the bounding boxes drawn on it
Example:
```python
>>> import supervision as sv
>>> classes = ['person', ...]
>>> image = ...
>>> detections = sv.Detections(...)
>>> corner_box_annotator = sv.CorneredBoxAnotator()
>>> annotated_frame = corner_box_annotator.annotate(
... scene=image.copy(),
... detections=detections,
... labels=labels
... )
```
"""
line_thickness = self.thickness + 2
for i in range(len(detections)):
x1, y1, x2, y2 = detections.xyxy[i].astype(int)
if self.color_by_track:
tracker_id = (
detections.tracker_id[i]
if detections.tracker_id is not None
else None
)
idx = tracker_id if tracker_id is not None else i
else:
class_id = (
detections.class_id[i] if detections.class_id is not None else None
)
idx = class_id if class_id is not None else i
color = (
self.color.by_idx(idx)
if isinstance(self.color, ColorPalette)
else self.color
)
box_width = x2 - x1
box_height = y2 - y1
cv2.line(
scene,
(x1, y1),
(x1 + int(0.2 * box_width), y1),
color.as_bgr(),
thickness=line_thickness,
)
cv2.line(
scene,
(x2 - int(0.2 * box_width), y1),
(x2, y1),
color.as_bgr(),
thickness=line_thickness,
)
cv2.line(
scene,
(x1, y2),
(x1 + int(0.2 * box_width), y2),
color.as_bgr(),
thickness=line_thickness,
)
cv2.line(
scene,
(x2 - int(0.2 * box_width), y2),
(x2, y2),
color.as_bgr(),
thickness=line_thickness,
)
cv2.line(
scene,
(x1, y1),
(x1, y1 + int(0.2 * box_height)),
color.as_bgr(),
thickness=line_thickness,
)
cv2.line(
scene,
(x2, y1 + int(0.2 * box_height)),
(x2, y1),
color.as_bgr(),
thickness=line_thickness,
)
cv2.line(
scene,
(x1, y2 - int(0.2 * box_height)),
(x1, y2),
color.as_bgr(),
thickness=line_thickness,
)
cv2.line(
scene,
(x2, y2 - int(0.2 * box_height)),
(x2, y2),
color.as_bgr(),
thickness=line_thickness,
)
return scene
class EllipseAnnotator(BaseAnnotator):
def __init__(
self,
color: Union[Color, ColorPalette] = ColorPalette.default(),
thickness: int = 2,
color_by_track: bool = False,
start_angle: int = -45,
end_angle: int = 360,
):
self.color: Union[Color, ColorPalette] = color
self.thickness: int = thickness
self.color_by_track = color_by_track
self.start_angle = start_angle
self.end_angle = end_angle
def annotate(
self,
scene: np.ndarray,
detections: Detections,
):
"""
Draws ellipse at bottom of the objects on the frame using the detections provided.
Args:
scene (np.ndarray): The image on which the bounding boxes will be drawn
detections (Detections): The detections for which the bounding boxes will be drawn
Returns:
np.ndarray: The image with the bounding boxes drawn on it
Example:
```python
>>> import supervision as sv
>>> classes = ['person', ...]
>>> image = ...
>>> detections = sv.Detections(...)
>>> ellipse_annotator = sv.EllipseAnotator()
>>> annotated_frame = ellipse_annotator.annotate(
... scene=image.copy(),
... detections=detections,
... labels=labels
... )
```
"""
for i in range(len(detections)):
x1, y1, x2, y2 = detections.xyxy[i].astype(int)
if self.color_by_track:
tracker_id = (
detections.tracker_id[i]
if detections.tracker_id is not None
else None
)
idx = tracker_id if tracker_id is not None else i
else:
class_id = (
detections.class_id[i] if detections.class_id is not None else None
)
idx = class_id if class_id is not None else i
color = (
self.color.by_idx(idx)
if isinstance(self.color, ColorPalette)
else self.color
)
center = (int((x1 + x2) / 2), y2)
width = x2 - x1
cv2.ellipse(
scene,
center=center,
axes=(int(width), int(0.35 * width)),
angle=0.0,
startAngle=self.start_angle,
endAngle=self.end_angle,
color=color.as_bgr(),
thickness=self.thickness,
lineType=cv2.LINE_4,
)
return scene
supervision/detection/core.py
+26
View File
@@ -441,10 +441,36 @@ class Detections:
(self.xyxy[:, 1] + self.xyxy[:, 3]) / 2,
]
).transpose()
elif anchor == Position.CENTER_LEFT:
return np.array(
[
self.xyxy[:, 0],
(self.xyxy[:, 1] + self.xyxy[:, 3]) / 2,
]
).transpose()
elif anchor == Position.CENTER_RIGHT:
return np.array(
[
self.xyxy[:, 2],
(self.xyxy[:, 1] + self.xyxy[:, 3]) / 2,
]
).transpose()
elif anchor == Position.BOTTOM_CENTER:
return np.array(
[(self.xyxy[:, 0] + self.xyxy[:, 2]) / 2, self.xyxy[:, 3]]
).transpose()
elif anchor == Position.BOTTOM_LEFT:
return np.array([self.xyxy[:, 0], self.xyxy[:, 3]]).transpose()
elif anchor == Position.BOTTOM_RIGHT:
return np.array([self.xyxy[:, 2], self.xyxy[:, 3]]).transpose()
elif anchor == Position.TOP_CENTER:
return np.array(
[(self.xyxy[:, 0] + self.xyxy[:, 2]) / 2, self.xyxy[:, 1]]
).transpose()
elif anchor == Position.TOP_LEFT:
return np.array([self.xyxy[:, 0], self.xyxy[:, 1]]).transpose()
elif anchor == Position.TOP_RIGHT:
return np.array([self.xyxy[:, 2], self.xyxy[:, 1]]).transpose()
raise ValueError(f"{anchor} is not supported.")
supervision/geometry/core.py
+7
View File
@@ -7,7 +7,14 @@ from typing import Tuple
class Position(Enum):
CENTER = "CENTER"
CENTER_LEFT = "CENTER_LEFT"
CENTER_RIGHT = "CENTER_RIGHT"
TOP_CENTER = "TOP_CENTER"
TOP_LEFT = "TOP_LEFT"
TOP_RIGHT = "TOP_RIGHT"
BOTTOM_LEFT = "BOTTOM_LEFT"
BOTTOM_CENTER = "BOTTOM_CENTER"
BOTTOM_RIGHT = "BOTTOM_RIGHT"
@classmethod
def list(cls):
test/dataset/test_core.py
-1
View File
@@ -172,5 +172,4 @@ def test_dataset_merge(
) -> None:
with exception:
result = DetectionDataset.merge(dataset_list=dataset_list)
print(result.images.keys())
assert result == expected_result
test/detection/test_core.py
+84 -1
View File
@@ -2,7 +2,8 @@ from contextlib import ExitStack as DoesNotRaise
import pytest
from supervision import Detections
from supervision.detection.core import Detections
from supervision.geometry.core import Position
from typing import Optional, Union, List
@@ -243,3 +244,85 @@ def test_merge(
with exception:
result = Detections.merge(detections_list=detections_list)
assert result == expected_result
@pytest.mark.parametrize(
'detections, anchor, expected_result, exception',
[
(
Detections.empty(),
Position.CENTER,
np.empty((0, 2), dtype=np.float32),
DoesNotRaise()
), # empty detections
(
mock_detections(xyxy=[[10, 10, 20, 20]]),
Position.CENTER,
np.array([[15, 15]], dtype=np.float32),
DoesNotRaise()
), # single detection; center anchor
(
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
Position.CENTER,
np.array([[15, 15], [25, 25]], dtype=np.float32),
DoesNotRaise()
), # two detections; center anchor
(
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
Position.CENTER_LEFT,
np.array([[10, 15], [20, 25]], dtype=np.float32),
DoesNotRaise()
), # two detections; center left anchor
(
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
Position.CENTER_RIGHT,
np.array([[20, 15], [30, 25]], dtype=np.float32),
DoesNotRaise()
), # two detections; center right anchor
(
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
Position.TOP_CENTER,
np.array([[15, 10], [25, 20]], dtype=np.float32),
DoesNotRaise()
), # two detections; top center anchor
(
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
Position.TOP_LEFT,
np.array([[10, 10], [20, 20]], dtype=np.float32),
DoesNotRaise()
), # two detections; top left anchor
(
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
Position.TOP_RIGHT,
np.array([[20, 10], [30, 20]], dtype=np.float32),
DoesNotRaise()
), # two detections; top right anchor
(
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
Position.BOTTOM_CENTER,
np.array([[15, 20], [25, 30]], dtype=np.float32),
DoesNotRaise()
), # two detections; bottom center anchor
(
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
Position.BOTTOM_LEFT,
np.array([[10, 20], [20, 30]], dtype=np.float32),
DoesNotRaise()
), # two detections; bottom left anchor
(
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
Position.BOTTOM_RIGHT,
np.array([[20, 20], [30, 30]], dtype=np.float32),
DoesNotRaise()
), # two detections; bottom right anchor
]
)
def test_get_anchor_coordinates(
detections: Detections,
anchor: Position,
expected_result: np.ndarray,
exception: Exception
) -> None:
result = detections.get_anchor_coordinates(anchor)
with exception:
assert np.array_equal(result, expected_result)