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

Merge pull request #406 from roboflow/feature/annotators-refactor-before-0.15.0-release

Browse Source 
Refactor TraceAnnotator to separate utility class
This commit is contained in:
Piotr Skalski authored and GitHub committed 2023-10-04 16:18:49 +02:00
2 parents 67fa66c7e7 8511bbb762
commit 10ad27d726
7 files changed
+206 -166
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/annotators.md
+25
View File
@@ -124,6 +124,27 @@
</div>
=== "Trace"
```python
>>> import supervision as sv
>>> image = ...
>>> detections = sv.Detections(...)
>>> trace_annotator = sv.TraceAnnotator()
>>> annotated_frame = trace_annotator.annotate(
... scene=image.copy(),
... detections=detections
... )
```
<div class="result" markdown>
![trace-annotator-example](https://media.roboflow.com/supervision-annotator-examples/trace-annotator-example.png){ align=center width="800" }
</div>
## BoundingBoxAnnotator
:::supervision.annotators.core.BoundingBoxAnnotator
@@ -147,3 +168,7 @@
## LabelAnnotator
:::supervision.annotators.core.LabelAnnotator
## TraceAnnotator
:::supervision.annotators.core.TraceAnnotator
supervision/__init__.py
+2 -1
View File
@@ -13,6 +13,7 @@ from supervision.annotators.core import (
EllipseAnnotator,
LabelAnnotator,
MaskAnnotator,
TraceAnnotator,
)
from supervision.classification.core import Classifications
from supervision.dataset.core import (
@@ -20,7 +21,7 @@ from supervision.dataset.core import (
ClassificationDataset,
DetectionDataset,
)
from supervision.detection.annotate import BoxAnnotator, TraceAnnotator
from supervision.detection.annotate import BoxAnnotator
from supervision.detection.core import Detections
from supervision.detection.line_counter import LineZone, LineZoneAnnotator
from supervision.detection.tools.inference_slicer import InferenceSlicer
supervision/annotators/base.py
-47
View File
@@ -1,55 +1,8 @@
from abc import ABC, abstractmethod
from enum import Enum
from typing import Union
import numpy as np
from supervision.detection.core import Detections
from supervision.draw.color import Color, ColorPalette
class ColorMap(Enum):
"""
Enum for annotator color mapping.
"""
INDEX = "index"
CLASS = "class"
TRACK = "track"
def resolve_color_idx(
detections: Detections, detection_idx: int, color_map: ColorMap = ColorMap.CLASS
) -> int:
if detection_idx >= len(detections):
raise ValueError(
f"Detection index {detection_idx}"
f"is out of bounds for detections of length {len(detections)}"
)
if color_map == ColorMap.INDEX:
return detection_idx
elif color_map == ColorMap.CLASS:
if detections.class_id is None:
raise ValueError(
"Could not resolve color by class because"
"Detections do not have class_id"
)
return detections.class_id[detection_idx]
elif color_map == ColorMap.TRACK:
if detections.tracker_id is None:
raise ValueError(
"Could not resolve color by track because"
"Detections do not have tracker_id"
)
return detections.tracker_id[detection_idx]
def resolve_color(color: Union[Color, ColorPalette], idx: int) -> Color:
if isinstance(color, ColorPalette):
return color.by_idx(idx)
else:
return color
class BaseAnnotator(ABC):
supervision/annotators/core.py
+89 -3
View File
@@ -1,12 +1,13 @@
from math import sqrt
from typing import List, Tuple, Union
from typing import List, Optional, Tuple, Union
import cv2
import numpy as np
from supervision.annotators.base import (
BaseAnnotator,
from supervision.annotators.base import BaseAnnotator
from supervision.annotators.utils import (
ColorMap,
Trace,
resolve_color,
resolve_color_idx,
)
@@ -562,3 +563,88 @@ class LabelAnnotator:
lineType=cv2.LINE_AA,
)
return scene
class TraceAnnotator:
"""
A class for drawing trace paths on an image based on detection coordinates.
!!! warning
This annotator utilizes the `tracker_id`. Read
[here](https://supervision.roboflow.com/trackers/) to learn how to plug
tracking into your inference pipeline.
"""
def __init__(
self,
color: Union[Color, ColorPalette] = ColorPalette.default(),
position: Optional[Position] = Position.CENTER,
trace_length: int = 30,
thickness: int = 2,
):
"""
Args:
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`.
"""
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
... )
```
![trace-annotator-example](https://media.roboflow.com/
supervision-annotator-examples/trace-annotator-example.png)
"""
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
supervision/annotators/utils.py
+89
View File
@@ -0,0 +1,89 @@
from enum import Enum
from typing import Optional, Union
import numpy as np
from supervision.detection.core import Detections
from supervision.draw.color import Color, ColorPalette
from supervision.geometry.core import Position
class ColorMap(Enum):
"""
Enum for annotator color mapping.
"""
INDEX = "index"
CLASS = "class"
TRACK = "track"
def resolve_color_idx(
detections: Detections, detection_idx: int, color_map: ColorMap = ColorMap.CLASS
) -> int:
if detection_idx >= len(detections):
raise ValueError(
f"Detection index {detection_idx}"
f"is out of bounds for detections of length {len(detections)}"
)
if color_map == ColorMap.INDEX:
return detection_idx
elif color_map == ColorMap.CLASS:
if detections.class_id is None:
raise ValueError(
"Could not resolve color by class because"
"Detections do not have class_id"
)
return detections.class_id[detection_idx]
elif color_map == ColorMap.TRACK:
if detections.tracker_id is None:
raise ValueError(
"Could not resolve color by track because"
"Detections do not have tracker_id"
)
return detections.tracker_id[detection_idx]
def resolve_color(color: Union[Color, ColorPalette], idx: int) -> Color:
if isinstance(color, ColorPalette):
return color.by_idx(idx)
return color
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]
supervision/detection/annotate.py
-114
View File
@@ -5,7 +5,6 @@ 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:
@@ -145,116 +144,3 @@ class BoxAnnotator:
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
test/annotators/test_base.py → test/annotators/test_utils.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import Optional
import pytest
from supervision.annotators.core import ColorMap, resolve_color_idx
from supervision.annotators.utils import ColorMap, resolve_color_idx
from supervision.detection.core import Detections