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

refactor docstrings in `draw`, `classification`, and `key_points` (#2161) 

* refactor docstrings in draw, classification, and key_points
* Refactor type annotations, logging, and empty output handling across key modules
* Refactor type annotations in `core.py` to include conditional `TYPE_CHECKING` for `torch` imports
* Apply suggestions from code review

---------

Co-authored-by: Jirka Borovec <6035284+Borda@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Omkar Kabde 2026-03-10 22:37:39 +05:30 committed by GitHub
parent 7d8b2b3fc5
commit a673640aad
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 199 additions and 151 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
pyproject.toml
View File

@ -204,7 +204,6 @@ module = [
"tests.*",
"examples.*",
# TODO: fix type errors in the following modules
"supervision.classification.core",
"supervision.detection.core",
"supervision.detection.line_zone",
"supervision.detection.tools.csv_sink",
@ -214,8 +213,6 @@ module = [
"supervision.detection.tools.smoother",
"supervision.detection.tools.transformers",
"supervision.detection.vlm",
"supervision.key_points.annotators",
"supervision.key_points.core",
"supervision.key_points.skeletons",
"supervision.metrics.utils.utils",
]

48
src/supervision/classification/core.py
View File

@ -1,9 +1,13 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import Any
from typing import TYPE_CHECKING, Any
import numpy as np
import numpy.typing as npt
if TYPE_CHECKING:
import torch
def _validate_class_ids(class_id: Any, n: int) -> None:
@ -27,8 +31,8 @@ def _validate_confidence(confidence: Any, n: int) -> None:
@dataclass
class Classifications:
class_id: np.ndarray
confidence: np.ndarray | None = None
class_id: npt.NDArray[np.int_]
confidence: npt.NDArray[np.floating] | None = None
def __post_init__(self) -> None:
"""
@ -46,16 +50,16 @@ class Classifications:
return len(self.class_id)
@classmethod
def from_clip(cls, clip_results) -> Classifications:
def from_clip(cls, clip_results: torch.Tensor) -> Classifications:
"""
Creates a Classifications instance from a
[clip](https://github.com/openai/clip) inference result.
Args:
clip_results (np.ndarray): The inference result from clip model.
clip_results: The inference result from clip model.
Returns:
Classifications: A new Classifications object.
A new Classifications object.
Example:
```python
@ -77,23 +81,25 @@ class Classifications:
confidence = clip_results.softmax(dim=-1).cpu().detach().numpy()[0]
if len(confidence) == 0:
return cls(class_id=np.array([]), confidence=np.array([]))
return cls(
class_id=np.array([], dtype=np.int_),
confidence=np.array([], dtype=np.float32),
)
class_ids = np.arange(len(confidence))
return cls(class_id=class_ids, confidence=confidence)
@classmethod
def from_ultralytics(cls, ultralytics_results) -> Classifications:
def from_ultralytics(cls, ultralytics_results: Any) -> Classifications:
"""
Creates a Classifications instance from a
[ultralytics](https://github.com/ultralytics/ultralytics) inference result.
Args:
ultralytics_results (ultralytics.engine.results.Results):
The inference result from ultralytics model.
ultralytics_results: The inference result from ultralytics model.
Returns:
Classifications: A new Classifications object.
A new Classifications object.
Example:
```python
@ -112,16 +118,16 @@ class Classifications:
return cls(class_id=np.arange(confidence.shape[0]), confidence=confidence)
@classmethod
def from_timm(cls, timm_results) -> Classifications:
def from_timm(cls, timm_results: Any) -> Classifications:
"""
Creates a Classifications instance from a
[timm](https://huggingface.co/docs/hub/timm) inference result.
Args:
timm_results (torch.Tensor): The inference result from timm model.
timm_results: The inference result from timm model.
Returns:
Classifications: A new Classifications object.
A new Classifications object.
Example:
```python
@ -149,22 +155,26 @@ class Classifications:
confidence = timm_results.cpu().detach().numpy()[0]
if len(confidence) == 0:
return cls(class_id=np.array([]), confidence=np.array([]))
return cls(
class_id=np.array([], dtype=np.int_),
confidence=np.array([], dtype=np.float32),
)
class_id = np.arange(len(confidence))
return cls(class_id=class_id, confidence=confidence)
def get_top_k(self, k: int) -> tuple[np.ndarray, np.ndarray]:
def get_top_k(
self, k: int
) -> tuple[npt.NDArray[np.int_], npt.NDArray[np.floating]]:
"""
Retrieve the top k class IDs and confidences,
ordered in descending order by confidence.
Args:
k (int): The number of top class IDs and confidences to retrieve.
k: The number of top class IDs and confidences to retrieve.
Returns:
Tuple[np.ndarray, np.ndarray]: A tuple containing
the top k class IDs and confidences.
A tuple containing the top k class IDs and confidences.
Example:
```pycon

48
src/supervision/draw/color.py
View File

@ -71,9 +71,9 @@ class Color:
codes, converting colors to hex strings, RGB tuples, and BGR tuples.
Attributes:
r (int): Red channel value (0-255).
g (int): Green channel value (0-255).
b (int): Blue channel value (0-255).
r: Red channel value (0-255).
g: Green channel value (0-255).
b: Blue channel value (0-255).
Example:
```pycon
@ -105,13 +105,13 @@ class Color:
Create a Color instance from a hex string.
Args:
color_hex (str): The hex string representing the color. This string can
color_hex: The hex string representing the color. This string can
start with '#' followed by either 3 or 6 hexadecimal characters. In
case of 3 characters, each character is repeated to form the full
6-character hex code.
Returns:
Color: An instance representing the color.
An instance representing the color.
Example:
```pycon
@ -136,11 +136,11 @@ class Color:
Create a Color instance from an RGB tuple.
Args:
color_tuple (Tuple[int, int, int]): A tuple representing the color in RGB
format, where each element is an integer in the range 0-255.
color_tuple: A tuple representing the color in RGB format, where each
element is an integer in the range 0-255.
Returns:
Color: An instance representing the color.
An instance representing the color.
Raises:
ValueError: If any RGB value is outside the range 0-255.
@ -164,11 +164,11 @@ class Color:
Create a Color instance from a BGR tuple.
Args:
color_tuple (Tuple[int, int, int]): A tuple representing the color in BGR
format, where each element is an integer in the range 0-255.
color_tuple: A tuple representing the color in BGR format, where each
element is an integer in the range 0-255.
Returns:
Color: An instance representing the color.
An instance representing the color.
Raises:
ValueError: If any BGR value is outside the range 0-255.
@ -191,7 +191,7 @@ class Color:
Converts the Color instance to a hex string.
Returns:
str: The hexadecimal color string.
The hexadecimal color string.
Example:
```pycon
@ -208,7 +208,7 @@ class Color:
Returns the color as an RGB tuple.
Returns:
Tuple[int, int, int]: RGB tuple.
RGB tuple.
Example:
```pycon
@ -225,7 +225,7 @@ class Color:
Returns the color as a BGR tuple.
Returns:
Tuple[int, int, int]: BGR tuple.
BGR tuple.
Example:
```pycon
@ -291,7 +291,7 @@ class ColorPalette:
Returns a default color palette.
Returns:
ColorPalette: A ColorPalette instance with default colors.
A ColorPalette instance with default colors.
Example:
```pycon
@ -312,7 +312,7 @@ class ColorPalette:
Returns a Roboflow color palette.
Returns:
ColorPalette: A ColorPalette instance with Roboflow colors.
A ColorPalette instance with Roboflow colors.
Example:
```pycon
@ -337,10 +337,10 @@ class ColorPalette:
Create a ColorPalette instance from a list of hex strings.
Args:
color_hex_list (List[str]): List of color hex strings.
color_hex_list: List of color hex strings.
Returns:
ColorPalette: A ColorPalette instance.
A ColorPalette instance.
Example:
```pycon
@ -360,11 +360,11 @@ class ColorPalette:
Create a ColorPalette instance from a Matplotlib color palette.
Args:
palette_name (str): Name of the Matplotlib palette.
color_count (int): Number of colors to sample from the palette.
palette_name: Name of the Matplotlib palette.
color_count: Number of colors to sample from the palette.
Returns:
ColorPalette: A ColorPalette instance.
A ColorPalette instance.
Example:
```pycon
@ -393,10 +393,10 @@ class ColorPalette:
Return the color at a given index in the palette.
Args:
idx (int): Index of the color in the palette.
idx: Index of the color in the palette.
Returns:
Color: Color at the given index.
Color at the given index.
Example:
```pycon
@ -418,7 +418,7 @@ class ColorPalette:
Returns the number of colors in the palette.
Returns:
int: The number of colors.
The number of colors.
"""
return len(self.colors)

42
src/supervision/draw/utils.py
View File

@ -229,23 +229,25 @@ def draw_text(
"""
Draw text with background on a scene.
Parameters:
scene (np.ndarray): A 2-dimensional numpy ndarray representing an image or scene
text (str): The text to be drawn.
text_anchor (Point): The anchor point for the text, represented as a
Args:
scene: A numpy ndarray representing the image, typically of shape
(H, W, 3) for a color BGR image or (H, W) for grayscale,
with dtype uint8.
text: The text to be drawn.
text_anchor: The anchor point for the text, represented as a
Point object with x and y attributes.
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
text_color: The color of the text. Defaults to black.
text_scale: The scale of the text. Defaults to 0.5.
text_thickness: The thickness of the text. Defaults to 1.
text_padding: The amount of padding to add around the text
when drawing a rectangle in the background. Defaults to 10.
text_font (int): The font to use for the text.
text_font: The font to use for the text.
Defaults to cv2.FONT_HERSHEY_SIMPLEX.
background_color (Optional[Color]): The color of the background rectangle,
background_color: The color of the background rectangle,
if one is to be drawn. Defaults to None.
Returns:
np.ndarray: The input scene with the text drawn on it.
The input scene with the text drawn on it.
Examples:
```pycon
@ -306,13 +308,13 @@ def draw_image(
Draws an image onto a given scene with specified opacity and dimensions.
Args:
scene (np.ndarray): Background image where the new image will be drawn.
image (Union[str, np.ndarray]): Image to draw.
opacity (float): Opacity of the image to be drawn.
rect (Rect): Rectangle specifying where to draw the image.
scene: Background image where the new image will be drawn.
image: Image to draw, either a file path or an already-loaded image array.
opacity: Opacity of the image to be drawn.
rect: Rectangle specifying where to draw the image.
Returns:
np.ndarray: The updated scene.
The updated scene.
Raises:
FileNotFoundError: If the image path does not exist.
@ -373,10 +375,10 @@ def calculate_optimal_text_scale(resolution_wh: tuple[int, int]) -> float:
consistent readability.
Args:
resolution_wh (tuple[int, int]): (width, height) of the image in pixels
resolution_wh: A tuple of `(width, height)` of the image in pixels.
Returns:
float: recommended font scale factor
Recommended font scale factor.
Examples:
```pycon
@ -398,10 +400,10 @@ def calculate_optimal_line_thickness(resolution_wh: tuple[int, int]) -> int:
image resolution.
Args:
resolution_wh (tuple[int, int]): (width, height) of the image in pixels
resolution_wh: A tuple of `(width, height)` of the image in pixels.
Returns:
int: recommended line thickness in pixels
Recommended line thickness in pixels.
Examples:
```pycon

94
src/supervision/key_points/annotators.py
View File

@ -1,10 +1,12 @@
from __future__ import annotations
from abc import ABC, abstractmethod
from logging import warn
from collections.abc import Sequence
from typing import Any
import cv2
import numpy as np
import numpy.typing as npt
from supervision.detection.utils.boxes import pad_boxes, spread_out_boxes
from supervision.draw.base import ImageType
@ -14,6 +16,9 @@ from supervision.geometry.core import Rect
from supervision.key_points.core import KeyPoints
from supervision.key_points.skeletons import SKELETONS_BY_VERTEX_COUNT
from supervision.utils.conversion import ensure_cv2_image_for_class_method
from supervision.utils.logger import _get_logger
logger = _get_logger(__name__)
class BaseKeyPointAnnotator(ABC):
@ -36,9 +41,8 @@ class VertexAnnotator(BaseKeyPointAnnotator):
) -> None:
"""
Args:
color (Color): The color to use for annotating key points.
radius (int): The radius of the circles used to represent the key
points.
color: The color to use for annotating key points.
radius: The radius of the circles used to represent the key points.
"""
self.color = color
self.radius = radius
@ -50,11 +54,10 @@ class VertexAnnotator(BaseKeyPointAnnotator):
points. It draws circles at each key point location.
Args:
scene (ImageType): The image where skeleton vertices will be drawn.
`ImageType` is a flexible type, accepting either `numpy.ndarray` or
`PIL.Image.Image`.
key_points (KeyPoints): A collection of key points where each key point
consists of x and y coordinates.
scene: The image where skeleton vertices will be drawn. `ImageType` is a
flexible type, accepting either `numpy.ndarray` or `PIL.Image.Image`.
key_points: A collection of key points where each key point consists of x
and y coordinates.
Returns:
The annotated image, matching the type of `scene` (`numpy.ndarray`
@ -108,14 +111,14 @@ class EdgeAnnotator(BaseKeyPointAnnotator):
self,
color: Color = Color.ROBOFLOW,
thickness: int = 2,
edges: list[tuple[int, int]] | None = None,
edges: Sequence[tuple[int, int]] | None = None,
) -> None:
"""
Args:
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.
color: The color to use for the edges.
thickness: The thickness of the edges.
edges: The edges to draw. If set to `None`, will attempt to select
automatically.
"""
self.color = color
self.thickness = thickness
@ -128,16 +131,14 @@ class EdgeAnnotator(BaseKeyPointAnnotator):
edges.
Args:
scene (ImageType): The image where skeleton edges will be drawn. `ImageType`
is a flexible type, accepting either `numpy.ndarray` or
`PIL.Image.Image`.
key_points (KeyPoints): A collection of key points where each key point
consists of x and y coordinates.
scene: The image where skeleton edges will be drawn. `ImageType` is a
flexible type, accepting either `numpy.ndarray` or `PIL.Image.Image`.
key_points: A collection of key points where each key point consists of x
and y coordinates.
Returns:
Returns:
The annotated image, matching the type of `scene` (`numpy.ndarray`
or `PIL.Image.Image`)
The annotated image, matching the type of `scene` (`numpy.ndarray`
or `PIL.Image.Image`)
Example:
```pycon
@ -169,7 +170,7 @@ class EdgeAnnotator(BaseKeyPointAnnotator):
if not edges:
edges = SKELETONS_BY_VERTEX_COUNT.get(len(xy))
if not edges:
warn(f"No skeleton found with {len(xy)} vertices")
logger.warning("No skeleton found with %d vertices", len(xy))
return scene
for class_a, class_b in edges:
@ -209,18 +210,16 @@ class VertexLabelAnnotator:
):
"""
Args:
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]]): 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): 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.
smart_position (bool): Spread out the labels to avoid overlap.
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: The color to use for the labels. If a list is provided, the
colors will be used in order for each keypoint.
text_scale: The scale of the text.
text_thickness: The thickness of the text.
text_padding: The padding around the text.
border_radius: The radius of the rounded corners of the boxes. Set to a
high value to produce circles.
smart_position: Spread out the labels to avoid overlap.
"""
self.border_radius: int = border_radius
self.color: Color | list[Color] = color
@ -241,13 +240,12 @@ class VertexLabelAnnotator:
points to determine the locations where the vertices should be drawn.
Args:
scene (ImageType): The image where vertex labels will be drawn. `ImageType`
is a flexible type, accepting either `numpy.ndarray` or
`PIL.Image.Image`.
key_points (KeyPoints): A collection of key points where each key point
consists of x and y coordinates.
labels (Optional[List[str]]): A list of labels to be displayed on the
annotated image. If not provided, keypoint indices will be used.
scene: The image where vertex labels will be drawn. `ImageType` is a
flexible type, accepting either `numpy.ndarray` or `PIL.Image.Image`.
key_points: A collection of key points where each key point consists of x
and y coordinates.
labels: A list of labels to be displayed on the annotated image. If not
provided, keypoint indices will be used.
Returns:
The annotated image, matching the type of `scene` (`numpy.ndarray`
@ -351,14 +349,14 @@ class VertexLabelAnnotator:
skeletons_count=skeletons_count,
)
labels = self.preprocess_and_validate_labels(
processed_labels = self.preprocess_and_validate_labels(
labels=labels, points_count=points_count, skeletons_count=skeletons_count
)
anchors = anchors[mask]
colors = colors[mask]
text_colors = text_colors[mask]
labels = labels[mask]
filtered_labels = processed_labels[mask]
xyxy = np.array(
[
@ -369,7 +367,7 @@ class VertexLabelAnnotator:
text_thickness=self.text_thickness,
center_coordinates=tuple(anchor),
)
for anchor, label in zip(anchors, labels)
for anchor, label in zip(anchors, filtered_labels)
]
)
xyxy_padded = pad_boxes(xyxy=xyxy, px=self.text_padding)
@ -379,7 +377,7 @@ class VertexLabelAnnotator:
xyxy = pad_boxes(xyxy=xyxy_padded, px=-self.text_padding)
for text, color, text_color, box, box_padded in zip(
labels, colors, text_colors, xyxy, xyxy_padded
filtered_labels, colors, text_colors, xyxy, xyxy_padded
):
draw_rounded_rectangle(
scene=scene,
@ -425,7 +423,7 @@ class VertexLabelAnnotator:
@staticmethod
def preprocess_and_validate_labels(
labels: list[str] | None, points_count: int, skeletons_count: int
) -> np.ndarray:
) -> npt.NDArray[np.str_]:
if labels and len(labels) != points_count:
raise ValueError(
f"Number of labels ({len(labels)}) must match number of key points "
@ -441,7 +439,7 @@ class VertexLabelAnnotator:
colors: Color | list[Color] | None,
points_count: int,
skeletons_count: int,
) -> np.ndarray:
) -> npt.NDArray[Any]:
if isinstance(colors, list) and len(colors) != points_count:
raise ValueError(
f"Number of colors ({len(colors)}) must match number of key points "

82
src/supervision/key_points/core.py
View File

@ -2,7 +2,7 @@ from __future__ import annotations
from collections.abc import Iterable, Iterator
from dataclasses import dataclass, field
from typing import Any
from typing import Any, Union, cast
import numpy as np
import numpy.typing as npt
@ -12,6 +12,16 @@ from supervision.detection.core import Detections
from supervision.detection.utils.internal import get_data_item, is_data_equal
from supervision.validators import validate_key_points_fields
Index1D = Union[
int,
slice,
list[int],
list[bool],
npt.NDArray[np.int_],
npt.NDArray[np.bool_],
]
Index2D = tuple[Index1D, Index1D]
@dataclass
class KeyPoints:
@ -157,9 +167,9 @@ class KeyPoints:
xy: npt.NDArray[np.float32]
class_id: npt.NDArray[np.int_] | None = None
confidence: npt.NDArray[np.float32] | None = None
data: dict[str, npt.NDArray[np.generic] | list] = field(default_factory=dict)
data: dict[str, npt.NDArray[np.generic] | list[Any]] = field(default_factory=dict)
def __post_init__(self):
def __post_init__(self) -> None:
validate_key_points_fields(
xy=self.xy,
confidence=self.confidence,
@ -172,7 +182,7 @@ class KeyPoints:
Returns the number of objects in the `sv.KeyPoints` object.
Returns:
int: The number of objects.
The number of objects.
Example:
```pycon
@ -191,12 +201,10 @@ class KeyPoints:
self,
) -> Iterator[
tuple[
np.ndarray,
np.ndarray | None,
float | None,
int | None,
int | None,
dict[str, np.ndarray | list],
npt.NDArray[np.float32],
npt.NDArray[np.float32] | None,
npt.NDArray[np.int_] | None,
dict[str, npt.NDArray[np.generic] | list[Any]],
]
]:
"""
@ -211,7 +219,9 @@ class KeyPoints:
get_data_item(self.data, i),
)
def __eq__(self, other: KeyPoints) -> bool:
def __eq__(self, other: object) -> bool:
if not isinstance(other, KeyPoints):
return NotImplemented
return all(
[
np.array_equal(self.xy, other.xy),
@ -222,14 +232,14 @@ class KeyPoints:
)
@classmethod
def from_inference(cls, inference_result: dict | Any) -> KeyPoints:
def from_inference(cls, inference_result: Any) -> KeyPoints:
"""
Create a `sv.KeyPoints` object from the [Roboflow](https://roboflow.com/)
API inference result or the [Inference](https://inference.roboflow.com/)
package results.
Args:
inference_result (dict, any): The result from the
inference_result: The result from the
Roboflow API or Inference package containing predictions with keypoints.
Returns:
@ -305,7 +315,7 @@ class KeyPoints:
@classmethod
def from_mediapipe(
cls, mediapipe_results, resolution_wh: tuple[int, int]
cls, mediapipe_results: Any, resolution_wh: tuple[int, int]
) -> KeyPoints:
"""
Creates a `sv.KeyPoints` instance from a
@ -313,12 +323,11 @@ class KeyPoints:
pose landmark detection inference result.
Args:
mediapipe_results (Union[PoseLandmarkerResult, FaceLandmarkerResult, SolutionOutputs]):
The output results from Mediapipe. It support pose and face landmarks
from `PoseLandmaker`, `FaceLandmarker` and the legacy ones
from `Pose` and `FaceMesh`.
resolution_wh (Tuple[int, int]): A tuple of the form `(width, height)`
representing the resolution of the frame.
mediapipe_results: The output results from Mediapipe. It supports pose
and face landmarks from `PoseLandmarker`, `FaceLandmarker` and the
legacy ones from `Pose` and `FaceMesh`.
resolution_wh: A tuple of the form `(width, height)` representing the
resolution of the frame.
Returns:
A `sv.KeyPoints` object containing the keypoint coordinates and
@ -382,7 +391,7 @@ class KeyPoints:
face_landmarker_result, (image_width, image_height))
```
""" # noqa: E501 // docs
"""
if hasattr(mediapipe_results, "pose_landmarks"):
results = mediapipe_results.pose_landmarks
if not isinstance(mediapipe_results.pose_landmarks, list):
@ -431,14 +440,13 @@ class KeyPoints:
)
@classmethod
def from_ultralytics(cls, ultralytics_results) -> KeyPoints:
def from_ultralytics(cls, ultralytics_results: Any) -> KeyPoints:
"""
Creates a `sv.KeyPoints` instance from a
[YOLOv8](https://github.com/ultralytics/ultralytics) pose inference result.
Args:
ultralytics_results (ultralytics.engine.results.Keypoints):
The output Results instance from YOLOv8
ultralytics_results: The output Results instance from YOLOv8.
Returns:
A `sv.KeyPoints` object containing the keypoint coordinates, class IDs,
@ -469,14 +477,13 @@ class KeyPoints:
return cls(xy, class_id, confidence, data)
@classmethod
def from_yolo_nas(cls, yolo_nas_results) -> KeyPoints:
def from_yolo_nas(cls, yolo_nas_results: Any) -> KeyPoints:
"""
Create a `sv.KeyPoints` instance from a [YOLO-NAS](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS-POSE.md)
pose inference results.
Args:
yolo_nas_results (ImagePoseEstimationPrediction): The output object from
YOLO NAS.
yolo_nas_results: The output object from YOLO NAS.
Returns:
A `sv.KeyPoints` object containing the keypoint coordinates, class IDs,
@ -534,7 +541,7 @@ class KeyPoints:
[Detectron2](https://github.com/facebookresearch/detectron2) inference result.
Args:
detectron2_results (Any): The output of a
detectron2_results: The output of a
Detectron2 model containing instances with prediction data.
Returns:
@ -585,7 +592,7 @@ class KeyPoints:
[Transformers](https://github.com/huggingface/transformers) inference result.
Args:
transformers_results (Any): The output of a
transformers_results: The output of a
Transformers model containing instances with prediction data.
Returns:
@ -664,8 +671,9 @@ class KeyPoints:
return cls.empty()
def __getitem__(
self, index: int | slice | list[int] | np.ndarray | tuple | str
) -> KeyPoints | np.ndarray | list | None:
self,
index: Index1D | Index2D | str,
) -> KeyPoints | npt.NDArray[np.generic] | list[Any] | None:
if isinstance(index, str):
return self.data.get(index)
@ -728,13 +736,13 @@ class KeyPoints:
data=data_selected,
)
def __setitem__(self, key: str, value: np.ndarray | list):
def __setitem__(self, key: str, value: npt.NDArray[np.generic] | list[Any]) -> None:
"""
Set a value in the data dictionary of the `sv.KeyPoints` object.
Args:
key (str): The key in the data dictionary to set.
value (Union[np.ndarray, List]): The value to set for the key.
key: The key in the data dictionary to set.
value: The value to set for the key.
Examples:
```python
@ -787,7 +795,7 @@ class KeyPoints:
Returns `True` if the `KeyPoints` object is considered empty.
Returns:
bool: `True` if the object is empty, `False` otherwise.
`True` if the object is empty, `False` otherwise.
Example:
```pycon
@ -810,7 +818,7 @@ class KeyPoints:
approximates the bounding box of the detected object by
taking the bounding box that fits all key points.
Arguments:
Args:
selected_keypoint_indices: The
indices of the key points to include in the bounding box
calculation. This helps focus on a subset of key points,
@ -869,6 +877,6 @@ class KeyPoints:
detections = Detections.merge(detections_list)
detections.class_id = self.class_id
detections.data = self.data
detections = detections[detections.area > 0]
detections = cast(Detections, detections[detections.area > 0])
return detections

33
tests/classification/test_core.py
View File

@ -8,6 +8,23 @@ import pytest
from supervision.classification.core import Classifications
class _MockTensor:
def __init__(self, value: np.ndarray) -> None:
self.value = value
def softmax(self, dim: int) -> _MockTensor:
return self
def cpu(self) -> _MockTensor:
return self
def detach(self) -> _MockTensor:
return self
def numpy(self) -> np.ndarray:
return self.value
@pytest.mark.parametrize(
("class_id", "confidence", "k", "expected_result", "exception"),
[
@ -62,3 +79,19 @@ def test_top_k(
assert np.array_equal(result[0], expected_result[0])
assert np.array_equal(result[1], expected_result[1])
def test_from_clip_empty_output_dtypes() -> None:
result = Classifications.from_clip(_MockTensor(np.empty((1, 0), dtype=np.float32)))
assert result.class_id.dtype == np.int_
assert result.confidence is not None
assert result.confidence.dtype == np.float32
def test_from_timm_empty_output_dtypes() -> None:
result = Classifications.from_timm(_MockTensor(np.empty((1, 0), dtype=np.float32)))
assert result.class_id.dtype == np.int_
assert result.confidence is not None
assert result.confidence.dtype == np.float32