From a673640aada4cd37ff8ce377904376f7221dfcff Mon Sep 17 00:00:00 2001 From: Omkar Kabde Date: Tue, 10 Mar 2026 22:37:39 +0530 Subject: [PATCH] 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> --- pyproject.toml | 3 - src/supervision/classification/core.py | 48 +++++++----- src/supervision/draw/color.py | 48 ++++++------ src/supervision/draw/utils.py | 42 ++++++----- src/supervision/key_points/annotators.py | 94 ++++++++++++------------ src/supervision/key_points/core.py | 82 +++++++++++---------- tests/classification/test_core.py | 33 +++++++++ 7 files changed, 199 insertions(+), 151 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 05e33eb9..40aa2b78 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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", ] diff --git a/src/supervision/classification/core.py b/src/supervision/classification/core.py index 32f38722..e34c55c8 100644 --- a/src/supervision/classification/core.py +++ b/src/supervision/classification/core.py @@ -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 diff --git a/src/supervision/draw/color.py b/src/supervision/draw/color.py index d20a21a3..04526c54 100644 --- a/src/supervision/draw/color.py +++ b/src/supervision/draw/color.py @@ -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) diff --git a/src/supervision/draw/utils.py b/src/supervision/draw/utils.py index 3109d579..37ea37b6 100644 --- a/src/supervision/draw/utils.py +++ b/src/supervision/draw/utils.py @@ -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 diff --git a/src/supervision/key_points/annotators.py b/src/supervision/key_points/annotators.py index a13115b1..d4dfb730 100644 --- a/src/supervision/key_points/annotators.py +++ b/src/supervision/key_points/annotators.py @@ -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 " diff --git a/src/supervision/key_points/core.py b/src/supervision/key_points/core.py index 4099d0dd..a1636aa5 100644 --- a/src/supervision/key_points/core.py +++ b/src/supervision/key_points/core.py @@ -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 diff --git a/tests/classification/test_core.py b/tests/classification/test_core.py index 91784255..039b27c0 100644 --- a/tests/classification/test_core.py +++ b/tests/classification/test_core.py @@ -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