261 lines
8.9 KiB
Python
261 lines
8.9 KiB
Python
from typing import List, Optional, Union
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from supervision.detection.core import Detections
|
|
from supervision.draw.color import Color, ColorPalette
|
|
from supervision.geometry.core import Position
|
|
|
|
|
|
class BoxAnnotator:
|
|
"""
|
|
A class for drawing bounding boxes on an image using detections provided.
|
|
|
|
Attributes:
|
|
color (Union[Color, ColorPalette]): The color to draw the bounding box,
|
|
can be a single color or a color palette
|
|
thickness (int): The thickness of the bounding box lines, default is 2
|
|
text_color (Color): The color of the text on the bounding box, default is white
|
|
text_scale (float): The scale of the text on the bounding box, default is 0.5
|
|
text_thickness (int): The thickness of the text on the bounding box,
|
|
default is 1
|
|
text_padding (int): The padding around the text on the bounding box,
|
|
default is 5
|
|
|
|
"""
|
|
|
|
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,
|
|
):
|
|
self.color: Union[Color, ColorPalette] = color
|
|
self.thickness: int = thickness
|
|
self.text_color: Color = text_color
|
|
self.text_scale: float = text_scale
|
|
self.text_thickness: int = text_thickness
|
|
self.text_padding: int = text_padding
|
|
|
|
def annotate(
|
|
self,
|
|
scene: np.ndarray,
|
|
detections: Detections,
|
|
labels: Optional[List[str]] = None,
|
|
skip_label: bool = False,
|
|
) -> 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
|
|
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.
|
|
skip_label (bool): Is set to `True`, skips bounding box label annotation.
|
|
Returns:
|
|
np.ndarray: The image with the bounding boxes drawn on it
|
|
|
|
Example:
|
|
```python
|
|
>>> import supervision as sv
|
|
|
|
>>> classes = ['person', ...]
|
|
>>> image = ...
|
|
>>> detections = sv.Detections(...)
|
|
|
|
>>> box_annotator = sv.BoxAnnotator()
|
|
>>> labels = [
|
|
... f"{classes[class_id]} {confidence:0.2f}"
|
|
... for _, _, confidence, class_id, _
|
|
... in detections
|
|
... ]
|
|
>>> annotated_frame = box_annotator.annotate(
|
|
... scene=image.copy(),
|
|
... detections=detections,
|
|
... labels=labels
|
|
... )
|
|
```
|
|
"""
|
|
font = cv2.FONT_HERSHEY_SIMPLEX
|
|
for i in range(len(detections)):
|
|
x1, y1, x2, y2 = detections.xyxy[i].astype(int)
|
|
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,
|
|
)
|
|
if skip_label:
|
|
continue
|
|
|
|
text = (
|
|
f"{class_id}"
|
|
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 Trace:
|
|
def __init__(
|
|
self,
|
|
max_size: Optional[int] = None,
|
|
start_frame_id: int = 0,
|
|
anchor: Position = Position.CENTER,
|
|
) -> None:
|
|
self.current_frame_id = start_frame_id
|
|
self.max_size = max_size
|
|
self.anchor = anchor
|
|
|
|
self.frame_id = np.array([], dtype=int)
|
|
self.xy = np.empty((0, 2), dtype=np.float32)
|
|
self.tracker_id = np.array([], dtype=int)
|
|
|
|
def put(self, detections: Detections) -> None:
|
|
frame_id = np.full(len(detections), self.current_frame_id, dtype=int)
|
|
self.frame_id = np.concatenate([self.frame_id, frame_id])
|
|
self.xy = np.concatenate(
|
|
[self.xy, detections.get_anchor_coordinates(self.anchor)]
|
|
)
|
|
self.tracker_id = np.concatenate([self.tracker_id, detections.tracker_id])
|
|
|
|
unique_frame_id = np.unique(self.frame_id)
|
|
|
|
if 0 < self.max_size < len(unique_frame_id):
|
|
max_allowed_frame_id = self.current_frame_id - self.max_size + 1
|
|
filtering_mask = self.frame_id >= max_allowed_frame_id
|
|
self.frame_id = self.frame_id[filtering_mask]
|
|
self.xy = self.xy[filtering_mask]
|
|
self.tracker_id = self.tracker_id[filtering_mask]
|
|
|
|
self.current_frame_id += 1
|
|
|
|
def get(self, tracker_id: int) -> np.ndarray:
|
|
return self.xy[self.tracker_id == tracker_id]
|
|
|
|
|
|
class TraceAnnotator:
|
|
"""
|
|
A class for drawing trace paths on an image based on detection coordinates.
|
|
|
|
Attributes:
|
|
color (Union[Color, ColorPalette]): The color to draw the trace, can be
|
|
a single color or a color palette.
|
|
position (Optional[Position]): The position of the trace. Defaults to `CENTER`.
|
|
trace_length (int): The maximum length of the trace in terms of historical
|
|
points. Defaults to `30`.
|
|
thickness (int): The thickness of the trace lines. Defaults to `2`.
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
|
position: Optional[Position] = Position.CENTER,
|
|
trace_length: int = 30,
|
|
thickness: int = 2,
|
|
):
|
|
self.color: Union[Color, ColorPalette] = color
|
|
self.position = position
|
|
self.trace = Trace(max_size=trace_length)
|
|
self.thickness = thickness
|
|
|
|
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
|
|
"""
|
|
Draws trace paths on the frame based on the detection coordinates provided.
|
|
|
|
Args:
|
|
scene (np.ndarray): The image on which the traces will be drawn.
|
|
detections (Detections): The detections which include coordinates for
|
|
which the traces will be drawn.
|
|
|
|
Returns:
|
|
np.ndarray: The image with the trace paths drawn on it.
|
|
|
|
Example:
|
|
```python
|
|
>>> import supervision as sv
|
|
|
|
>>> image = ...
|
|
>>> detections = sv.Detections(...)
|
|
|
|
>>> trace_annotator = sv.TraceAnnotator()
|
|
>>> annotated_frame = trace_annotator.annotate(
|
|
... scene=image.copy(),
|
|
... detections=detections
|
|
... )
|
|
```
|
|
"""
|
|
self.trace.put(detections)
|
|
|
|
for i, (xyxy, mask, confidence, class_id, tracker_id) in enumerate(detections):
|
|
class_id = detections.class_id[i] if 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
|
|
)
|
|
|
|
xy = self.trace.get(tracker_id=tracker_id)
|
|
if len(xy) > 1:
|
|
scene = cv2.polylines(
|
|
scene,
|
|
[xy.astype(np.int32)],
|
|
False,
|
|
color=color.as_bgr(),
|
|
thickness=self.thickness,
|
|
)
|
|
return scene
|