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

fix(pre_commit): 🎨 auto format pre-commit hooks

This commit is contained in:
pre-commit-ci[bot] 2023-10-12 17:11:58 +00:00
parent e1e455f99c
commit 476a31fa79
3 changed files with 86 additions and 82 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

69
supervision/annotators/core.py
View File

@ -5,11 +5,7 @@ import cv2
import numpy as np
from supervision.annotators.base import BaseAnnotator
from supervision.annotators.utils import (
ColorLookup,
Trace,
resolve_color,
)
from supervision.annotators.utils import ColorLookup, Trace, resolve_color
from supervision.detection.core import Detections
from supervision.draw.color import Color, ColorPalette
from supervision.geometry.core import Position
@ -42,7 +38,7 @@ class BoundingBoxAnnotator(BaseAnnotator):
self,
scene: np.ndarray,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None
custom_color_lookup: Optional[np.ndarray] = None,
) -> np.ndarray:
"""
Annotates the given scene with bounding boxes based on the provided detections.
@ -79,8 +75,9 @@ class BoundingBoxAnnotator(BaseAnnotator):
color=self.color,
detections=detections,
detection_idx=detection_idx,
color_lookup=self.color_lookup if
custom_color_lookup is None else custom_color_lookup
color_lookup=self.color_lookup
if custom_color_lookup is None
else custom_color_lookup,
)
cv2.rectangle(
img=scene,
@ -119,7 +116,7 @@ class MaskAnnotator(BaseAnnotator):
self,
scene: np.ndarray,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None
custom_color_lookup: Optional[np.ndarray] = None,
) -> np.ndarray:
"""
Annotates the given scene with masks based on the provided detections.
@ -158,8 +155,9 @@ class MaskAnnotator(BaseAnnotator):
color=self.color,
detections=detections,
detection_idx=detection_idx,
color_lookup=self.color_lookup if
custom_color_lookup is None else custom_color_lookup
color_lookup=self.color_lookup
if custom_color_lookup is None
else custom_color_lookup,
)
mask = detections.mask[detection_idx]
colored_mask = np.zeros_like(scene, dtype=np.uint8)
@ -198,7 +196,7 @@ class BoxMaskAnnotator(BaseAnnotator):
self,
scene: np.ndarray,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None
custom_color_lookup: Optional[np.ndarray] = None,
) -> np.ndarray:
"""
Annotates the given scene with box masks based on the provided detections.
@ -236,8 +234,9 @@ class BoxMaskAnnotator(BaseAnnotator):
color=self.color,
detections=detections,
detection_idx=detection_idx,
color_lookup=self.color_lookup if
custom_color_lookup is None else custom_color_lookup
color_lookup=self.color_lookup
if custom_color_lookup is None
else custom_color_lookup,
)
cv2.rectangle(
img=scene,
@ -283,7 +282,7 @@ class HaloAnnotator(BaseAnnotator):
self,
scene: np.ndarray,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None
custom_color_lookup: Optional[np.ndarray] = None,
) -> np.ndarray:
"""
Annotates the given scene with halos based on the provided detections.
@ -326,8 +325,9 @@ class HaloAnnotator(BaseAnnotator):
color=self.color,
detections=detections,
detection_idx=detection_idx,
color_lookup=self.color_lookup if
custom_color_lookup is None else custom_color_lookup
color_lookup=self.color_lookup
if custom_color_lookup is None
else custom_color_lookup,
)
mask = detections.mask[detection_idx]
fmask = np.logical_or(fmask, mask)
@ -376,7 +376,7 @@ class EllipseAnnotator(BaseAnnotator):
self,
scene: np.ndarray,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None
custom_color_lookup: Optional[np.ndarray] = None,
) -> np.ndarray:
"""
Annotates the given scene with ellipses based on the provided detections.
@ -413,8 +413,9 @@ class EllipseAnnotator(BaseAnnotator):
color=self.color,
detections=detections,
detection_idx=detection_idx,
color_lookup=self.color_lookup if
custom_color_lookup is None else custom_color_lookup
color_lookup=self.color_lookup
if custom_color_lookup is None
else custom_color_lookup,
)
center = (int((x1 + x2) / 2), y2)
width = x2 - x1
@ -462,7 +463,7 @@ class BoxCornerAnnotator(BaseAnnotator):
self,
scene: np.ndarray,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None
custom_color_lookup: Optional[np.ndarray] = None,
) -> np.ndarray:
"""
Annotates the given scene with box corners based on the provided detections.
@ -499,8 +500,9 @@ class BoxCornerAnnotator(BaseAnnotator):
color=self.color,
detections=detections,
detection_idx=detection_idx,
color_lookup=self.color_lookup if
custom_color_lookup is None else custom_color_lookup
color_lookup=self.color_lookup
if custom_color_lookup is None
else custom_color_lookup,
)
corners = [(x1, y1), (x2, y1), (x1, y2), (x2, y2)]
@ -545,7 +547,7 @@ class CircleAnnotator(BaseAnnotator):
self,
scene: np.ndarray,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None
custom_color_lookup: Optional[np.ndarray] = None,
) -> np.ndarray:
"""
Annotates the given scene with circles based on the provided detections.
@ -585,8 +587,9 @@ class CircleAnnotator(BaseAnnotator):
color=self.color,
detections=detections,
detection_idx=detection_idx,
color_lookup=self.color_lookup if
custom_color_lookup is None else custom_color_lookup
color_lookup=self.color_lookup
if custom_color_lookup is None
else custom_color_lookup,
)
cv2.circle(
img=scene,
@ -682,7 +685,7 @@ class LabelAnnotator:
scene: np.ndarray,
detections: Detections,
labels: List[str] = None,
custom_color_lookup: Optional[np.ndarray] = None
custom_color_lookup: Optional[np.ndarray] = None,
) -> np.ndarray:
"""
Annotates the given scene with labels based on the provided detections.
@ -721,8 +724,9 @@ class LabelAnnotator:
color=self.color,
detections=detections,
detection_idx=detection_idx,
color_lookup=self.color_lookup if
custom_color_lookup is None else custom_color_lookup
color_lookup=self.color_lookup
if custom_color_lookup is None
else custom_color_lookup,
)
text = (
f"{detections.class_id[detection_idx]}"
@ -861,7 +865,7 @@ class TraceAnnotator:
self,
scene: np.ndarray,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None
custom_color_lookup: Optional[np.ndarray] = None,
) -> np.ndarray:
"""
Draws trace paths on the frame based on the detection coordinates provided.
@ -901,8 +905,9 @@ class TraceAnnotator:
color=self.color,
detections=detections,
detection_idx=detection_idx,
color_lookup=self.color_lookup if
custom_color_lookup is None else custom_color_lookup
color_lookup=self.color_lookup
if custom_color_lookup is None
else custom_color_lookup,
)
xy = self.trace.get(tracker_id=tracker_id)
if len(xy) > 1:

5
supervision/annotators/utils.py
View File

@ -12,6 +12,7 @@ class ColorLookup(Enum):
"""
Enum for annotator color lookup.
"""
INDEX = "index"
CLASS = "class"
TRACK = "track"
@ -20,7 +21,7 @@ class ColorLookup(Enum):
def resolve_color_idx(
detections: Detections,
detection_idx: int,
color_lookup: Union[ColorLookup, np.ndarray] = ColorLookup.CLASS
color_lookup: Union[ColorLookup, np.ndarray] = ColorLookup.CLASS,
) -> int:
if detection_idx >= len(detections):
raise ValueError(
@ -63,7 +64,7 @@ def resolve_color(
color: Union[Color, ColorPalette],
detections: Detections,
detection_idx: int,
color_lookup: Union[ColorLookup, np.ndarray] = ColorLookup.CLASS
color_lookup: Union[ColorLookup, np.ndarray] = ColorLookup.CLASS,
) -> Color:
idx = resolve_color_idx(
detections=detections,

94
test/annotators/test_utils.py
View File

@ -1,10 +1,8 @@
from contextlib import ExitStack as DoesNotRaise
import numpy as np
from test.utils import mock_detections
from typing import Optional
import numpy as np
import pytest
from supervision.annotators.utils import ColorLookup, resolve_color_idx
@ -15,83 +13,83 @@ from supervision.detection.core import Detections
"detections, detection_idx, color_lookup, expected_result, exception",
[
(
mock_detections(
mock_detections(
xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
class_id=[5, 3],
tracker_id=[2, 6],
),
0,
ColorLookup.INDEX,
0,
DoesNotRaise(),
0,
ColorLookup.INDEX,
0,
DoesNotRaise(),
), # multiple detections; index lookup
(
mock_detections(
mock_detections(
xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
class_id=[5, 3],
tracker_id=[2, 6],
),
0,
ColorLookup.CLASS,
5,
DoesNotRaise(),
0,
ColorLookup.CLASS,
5,
DoesNotRaise(),
), # multiple detections; class lookup
(
mock_detections(
mock_detections(
xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
class_id=[5, 3],
tracker_id=[2, 6],
),
0,
ColorLookup.TRACK,
2,
DoesNotRaise(),
0,
ColorLookup.TRACK,
2,
DoesNotRaise(),
), # multiple detections; track lookup
(
Detections.empty(),
0,
ColorLookup.INDEX,
None,
pytest.raises(ValueError),
Detections.empty(),
0,
ColorLookup.INDEX,
None,
pytest.raises(ValueError),
), # no detections; index lookup; out of bounds
(
mock_detections(
mock_detections(
xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
class_id=[5, 3],
tracker_id=[2, 6],
),
2,
ColorLookup.INDEX,
None,
pytest.raises(ValueError),
2,
ColorLookup.INDEX,
None,
pytest.raises(ValueError),
), # multiple detections; index lookup; out of bounds
(
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
0,
ColorLookup.CLASS,
None,
pytest.raises(ValueError),
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
0,
ColorLookup.CLASS,
None,
pytest.raises(ValueError),
), # multiple detections; class lookup; no class_id
(
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
0,
ColorLookup.TRACK,
None,
pytest.raises(ValueError),
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
0,
ColorLookup.TRACK,
None,
pytest.raises(ValueError),
), # multiple detections; class lookup; no track_id
(
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
0,
np.array([1, 0]),
1,
DoesNotRaise(),
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
0,
np.array([1, 0]),
1,
DoesNotRaise(),
), # multiple detections; custom lookup; correct length
(
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
0,
np.array([1]),
None,
pytest.raises(ValueError),
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
0,
np.array([1]),
None,
pytest.raises(ValueError),
), # multiple detections; custom lookup; wrong length
],
)