From 664bf3cf796bc7b0d583fe56045cab80eff127d2 Mon Sep 17 00:00:00 2001 From: Omkar Kabde Date: Tue, 10 Mar 2026 23:10:05 +0530 Subject: [PATCH] refactor docstrings in `scr/supervision/detection` (#2162) * refactor docstrings in `scr/supervision/detection` * Enhance docstrings across multiple modules: clarify attributes/args, improve formatting, and update logic for handling sentinel values in metrics calculation. * Ensure consistent handling of `class_id` as integer across YOLO and Pascal VOC formats, fix NoneType handling in line zone logic, and add test coverage for multiclass annotator with None `class_id`. * Enforce `class_id` as integer in YOLO export, update line zone class count docstrings, and add test for non-integer `class_id`. * 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> --- pyproject.toml | 3 - src/supervision/dataset/formats/pascal_voc.py | 5 + src/supervision/dataset/formats/yolo.py | 10 +- src/supervision/detection/core.py | 362 +++++++++--------- src/supervision/detection/line_zone.py | 217 +++++------ src/supervision/detection/vlm.py | 177 ++++----- .../metrics/mean_average_precision.py | 2 +- tests/dataset/formats/test_yolo.py | 13 + tests/detection/test_line_counter.py | 19 +- 9 files changed, 428 insertions(+), 380 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 40aa2b78..c5fcb1dd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -204,15 +204,12 @@ module = [ "tests.*", "examples.*", # TODO: fix type errors in the following modules - "supervision.detection.core", - "supervision.detection.line_zone", "supervision.detection.tools.csv_sink", "supervision.detection.tools.inference_slicer", "supervision.detection.tools.json_sink", "supervision.detection.tools.polygon_zone", "supervision.detection.tools.smoother", "supervision.detection.tools.transformers", - "supervision.detection.vlm", "supervision.key_points.skeletons", "supervision.metrics.utils.utils", ] diff --git a/src/supervision/dataset/formats/pascal_voc.py b/src/supervision/dataset/formats/pascal_voc.py index 6832628c..91b6664c 100644 --- a/src/supervision/dataset/formats/pascal_voc.py +++ b/src/supervision/dataset/formats/pascal_voc.py @@ -117,6 +117,11 @@ def detections_to_pascal_voc( for xyxy, mask, _, class_id, _, _ in detections: if class_id is None: raise ValueError("Detections must include class_id for Pascal VOC export.") + if not isinstance(class_id, (int, np.integer)): + raise ValueError( + f"Detections class_id must be an integer for Pascal VOC export, " + f"got {type(class_id)!r}." + ) name = classes[class_id] if mask is not None: polygons = approximate_mask_with_polygons( diff --git a/src/supervision/dataset/formats/yolo.py b/src/supervision/dataset/formats/yolo.py index 16a14142..b9bb0233 100644 --- a/src/supervision/dataset/formats/yolo.py +++ b/src/supervision/dataset/formats/yolo.py @@ -246,6 +246,12 @@ def detections_to_yolo_annotations( for xyxy, mask, _, class_id, _, _ in detections: if class_id is None: raise ValueError("Class ID is required for YOLO annotations.") + if not isinstance(class_id, (int, np.integer)): + raise ValueError( + f"Detections class_id must be an integer for YOLO export, " + f"got {type(class_id)!r}." + ) + class_id_int = int(class_id) if mask is not None: polygons = approximate_mask_with_polygons( @@ -258,14 +264,14 @@ def detections_to_yolo_annotations( xyxy = polygon_to_xyxy(polygon=polygon) next_object = object_to_yolo( xyxy=xyxy, - class_id=class_id, + class_id=class_id_int, image_shape=image_shape, polygon=polygon, ) annotation.append(next_object) else: next_object = object_to_yolo( - xyxy=xyxy, class_id=class_id, image_shape=image_shape + xyxy=xyxy, class_id=class_id_int, image_shape=image_shape ) annotation.append(next_object) return annotation diff --git a/src/supervision/detection/core.py b/src/supervision/detection/core.py index 92388cbc..baabc568 100644 --- a/src/supervision/detection/core.py +++ b/src/supervision/detection/core.py @@ -4,9 +4,10 @@ from collections.abc import Iterator from dataclasses import dataclass, field from enum import Enum from functools import reduce -from typing import Any +from typing import Any, cast import numpy as np +import numpy.typing as npt from supervision.config import ( CLASS_NAME_DATA_FIELD, @@ -129,33 +130,33 @@ class Detections: ``` Attributes: - xyxy (np.ndarray): An array of shape `(n, 4)` containing + xyxy: An array of shape `(n, 4)` containing the bounding boxes coordinates in format `[x1, y1, x2, y2]` - mask: (Optional[np.ndarray]): An array of shape - `(n, H, W)` containing the segmentation masks (`bool` data type). - confidence (Optional[np.ndarray]): An array of shape - `(n,)` containing the confidence scores of the detections. - class_id (Optional[np.ndarray]): An array of shape - `(n,)` containing the class ids of the detections. - tracker_id (Optional[np.ndarray]): An array of shape - `(n,)` containing the tracker ids of the detections. - data (Dict[str, Union[np.ndarray, List]]): A dictionary containing additional + mask: An array of shape `(n, H, W)` containing the segmentation masks + (`bool` data type), or `None` when masks are not available. + confidence: An array of shape `(n,)` containing the confidence scores + of the detections, or `None` when confidence values are not available. + class_id: An array of shape `(n,)` containing the class ids of the + detections, or `None` when class ids are not available. + tracker_id: An array of shape `(n,)` containing the tracker ids of the + detections, or `None` when tracker ids are not available. + data: A dictionary containing additional data where each key is a string representing the data type, and the value is either a NumPy array or a list of corresponding data. - metadata (Dict[str, Any]): A dictionary containing collection-level metadata + metadata: A dictionary containing collection-level metadata that applies to the entire set of detections. This may include information such as the video name, camera parameters, timestamp, or other global metadata. """ # noqa: E501 // docs - xyxy: np.ndarray - mask: np.ndarray | None = None - confidence: np.ndarray | None = None - class_id: np.ndarray | None = None - tracker_id: np.ndarray | None = None - data: dict[str, np.ndarray | list] = field(default_factory=dict) + xyxy: npt.NDArray[np.generic] + mask: npt.NDArray[np.generic] | None = None + confidence: npt.NDArray[np.generic] | None = None + class_id: npt.NDArray[np.generic] | None = None + tracker_id: npt.NDArray[np.generic] | None = None + data: dict[str, npt.NDArray[np.generic] | list[Any]] = field(default_factory=dict) metadata: dict[str, Any] = field(default_factory=dict) - def __post_init__(self): + def __post_init__(self) -> None: validate_detections_fields( xyxy=self.xyxy, mask=self.mask, @@ -165,7 +166,7 @@ class Detections: data=self.data, ) - def __len__(self): + def __len__(self) -> int: """ Returns the number of detections in the Detections object. """ @@ -175,12 +176,12 @@ class Detections: self, ) -> Iterator[ tuple[ - np.ndarray, - np.ndarray | None, - float | None, - int | None, - int | None, - dict[str, np.ndarray | list], + npt.NDArray[np.generic], + npt.NDArray[np.generic] | None, + np.generic | None, + np.generic | None, + np.generic | None, + dict[str, npt.NDArray[np.generic] | list[Any]], ] ]: """ @@ -197,7 +198,9 @@ class Detections: get_data_item(self.data, i), ) - def __eq__(self, other: Detections): + def __eq__(self, other: object) -> bool: + if not isinstance(other, Detections): + return NotImplemented return all( [ np.array_equal(self.xyxy, other.xyxy), @@ -211,17 +214,16 @@ class Detections: ) @classmethod - def from_yolov5(cls, yolov5_results) -> Detections: + def from_yolov5(cls, yolov5_results: Any) -> Detections: """ Creates a Detections instance from a [YOLOv5](https://github.com/ultralytics/yolov5) inference result. Args: - yolov5_results (yolov5.models.common.Detections): - The output Detections instance from YOLOv5 + yolov5_results: The output Detections instance from YOLOv5. Returns: - Detections: A new Detections object. + A new Detections object. Example: ```python @@ -244,7 +246,7 @@ class Detections: ) @classmethod - def from_ultralytics(cls, ultralytics_results) -> Detections: + def from_ultralytics(cls, ultralytics_results: Any) -> Detections: """ Creates a `sv.Detections` instance from a [YOLOv8](https://github.com/ultralytics/ultralytics) inference result. @@ -257,11 +259,10 @@ class Detections: [OBB](https://docs.ultralytics.com/tasks/obb/) models. Args: - ultralytics_results (ultralytics.yolo.engine.results.Results): - The output Results instance from Ultralytics + ultralytics_results: The output Results instance from Ultralytics. Returns: - Detections: A new Detections object. + A new Detections object. Example: ```python @@ -325,20 +326,19 @@ class Detections: return cls.empty() @classmethod - def from_yolo_nas(cls, yolo_nas_results) -> Detections: + def from_yolo_nas(cls, yolo_nas_results: Any) -> Detections: """ Creates a Detections instance from a [YOLO-NAS](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS.md) inference result. Args: - yolo_nas_results (ImageDetectionPrediction): - The output Results instance from YOLO-NAS + yolo_nas_results: The output Results instance from YOLO-NAS. ImageDetectionPrediction is coming from - 'super_gradients.training.models.prediction_results' + 'super_gradients.training.models.prediction_results'. Returns: - Detections: A new Detections object. + A new Detections object. Example: ```python @@ -364,7 +364,7 @@ class Detections: @classmethod def from_tensorflow( - cls, tensorflow_results: dict, resolution_wh: tuple + cls, tensorflow_results: dict[str, Any], resolution_wh: tuple[int, int] ) -> Detections: """ Creates a Detections instance from a @@ -372,11 +372,13 @@ class Detections: inference result. Args: - tensorflow_results (dict): - The output results from Tensorflow Hub. + tensorflow_results: The output results from Tensorflow Hub. + resolution_wh: The input image resolution as `(width, height)`. + Bounding boxes from Tensorflow are normalized and are scaled + to absolute coordinates using this resolution. Returns: - Detections: A new Detections object. + A new Detections object. Example: ```python @@ -389,7 +391,9 @@ class Detections: model = hub.load(module_handle) img = np.array(cv2.imread(SOURCE_IMAGE_PATH)) result = model(img) - detections = sv.Detections.from_tensorflow(result) + detections = sv.Detections.from_tensorflow( + result, resolution_wh=(img.shape[1], img.shape[0]) + ) ``` """ @@ -404,18 +408,17 @@ class Detections: ) @classmethod - def from_deepsparse(cls, deepsparse_results) -> Detections: + def from_deepsparse(cls, deepsparse_results: Any) -> Detections: """ Creates a Detections instance from a [DeepSparse](https://github.com/neuralmagic/deepsparse) inference result. Args: - deepsparse_results (deepsparse.yolo.schemas.YOLOOutput): - The output Results instance from DeepSparse. + deepsparse_results: The output Results instance from DeepSparse. Returns: - Detections: A new Detections object. + A new Detections object. Example: ```python @@ -441,18 +444,17 @@ class Detections: ) @classmethod - def from_mmdetection(cls, mmdet_results) -> Detections: + def from_mmdetection(cls, mmdet_results: Any) -> Detections: """ Creates a Detections instance from a [mmdetection](https://github.com/open-mmlab/mmdetection) and [mmyolo](https://github.com/open-mmlab/mmyolo) inference result. Args: - mmdet_results (mmdet.structures.DetDataSample): - The output Results instance from MMDetection. + mmdet_results: The output Results instance from MMDetection. Returns: - Detections: A new Detections object. + A new Detections object. Example: ```python @@ -481,7 +483,9 @@ class Detections: @classmethod def from_transformers( - cls, transformers_results: dict, id2label: dict[int, str] | None = None + cls, + transformers_results: dict[str, Any], + id2label: dict[int, str] | None = None, ) -> Detections: """ Creates a Detections instance from object detection or panoptic, semantic @@ -489,17 +493,17 @@ class Detections: [Transformer](https://github.com/huggingface/transformers) inference result. Args: - transformers_results (Union[dict, torch.Tensor]): Inference results from - your Transformers model. This can be either a dictionary containing - valuable outputs like `scores`, `labels`, `boxes`, `masks`, - `segments_info`, and `segmentation`, or a `torch.Tensor` holding a - segmentation map where values represent class IDs. - id2label (Optional[Dict[int, str]]): A dictionary mapping class IDs to - labels, typically part of the `transformers` model configuration. If - provided, the resulting dictionary will include class names. + transformers_results: Inference results from your Transformers model. + This can be either a dictionary containing valuable outputs like + `scores`, `labels`, `boxes`, `masks`, `segments_info`, and + `segmentation`, or a `torch.Tensor` holding a segmentation map + where values represent class IDs. + id2label: A dictionary mapping class IDs to labels, typically part of + the `transformers` model configuration. If provided, the resulting + dictionary will include class names. Returns: - Detections: A new Detections object. + A new Detections object. Example: ```python @@ -565,11 +569,11 @@ class Detections: [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: - (Detections): A Detections object containing the bounding boxes, + A Detections object containing the bounding boxes, class IDs, and confidences of the predictions. Example: @@ -606,7 +610,7 @@ class Detections: ) @classmethod - def from_inference(cls, roboflow_result: dict | Any) -> Detections: + def from_inference(cls, roboflow_result: dict[str, Any] | Any) -> Detections: """ Create a `sv.Detections` object from the [Roboflow](https://roboflow.com/) API inference result or the [Inference](https://inference.roboflow.com/) @@ -615,11 +619,11 @@ class Detections: them into a Detections object. Args: - roboflow_result (dict, any): The result from the + roboflow_result: The result from the Roboflow API or Inference package containing predictions. Returns: - (Detections): A Detections object containing the bounding boxes, class IDs, + A Detections object containing the bounding boxes, class IDs, and confidences of the predictions. Example: @@ -658,17 +662,17 @@ class Detections: ) @classmethod - def from_sam(cls, sam_result: list[dict]) -> Detections: + def from_sam(cls, sam_result: list[dict[str, Any]]) -> Detections: """ Creates a Detections instance from [Segment Anything Model](https://github.com/facebookresearch/segment-anything) inference result. Args: - sam_result (List[dict]): The output Results instance from SAM + sam_result: The output Results instance from SAM. Returns: - Detections: A new Detections object. + A new Detections object. Example: ```python @@ -701,7 +705,7 @@ class Detections: @classmethod def from_sam3( - cls, sam3_result: dict | Any, resolution_wh: tuple[int, int] + cls, sam3_result: dict[str, Any] | Any, resolution_wh: tuple[int, int] ) -> Detections: """ Creates a Detections instance from @@ -709,15 +713,15 @@ class Detections: Supports both PVS and PCS SAM3 segmentation formats. Args: - sam3_result (dict | Any): The output result from SAM 3 inference, - either Sam3PromptResult from inference package or dict containing + sam3_result: The output result from SAM 3 inference, either + Sam3PromptResult from inference package or dict containing prompt_results with polygon predictions. - resolution_wh (Tuple[int, int]): The width and height of the image - used for mask generation. + resolution_wh: The width and height of the image used for mask + generation. Returns: - Detections: A new Detections object. - The `class_id` field contains the prompt index for each polygon. + A new Detections object. The `class_id` field contains the prompt + index for each polygon. Example: ```python @@ -798,7 +802,7 @@ class Detections: if not pred_masks: continue - full_mask = np.zeros((height, width), dtype=bool) + full_mask: npt.NDArray[np.bool_] = np.zeros((height, width), dtype=bool) for poly in pred_masks: polygon = np.array(poly, dtype=np.int32) mask = polygon_to_mask( @@ -826,7 +830,7 @@ class Detections: @classmethod def from_azure_analyze_image( - cls, azure_result: dict, class_map: dict[int, str] | None = None + cls, azure_result: dict[str, Any], class_map: dict[int, str] | None = None ) -> Detections: """ Creates a Detections instance from [Azure Image Analysis 4.0]( @@ -834,13 +838,13 @@ class Detections: concept-object-detection-40). Args: - azure_result (dict): The result from Azure Image Analysis. It should + azure_result: The result from Azure Image Analysis. It should contain detected objects and their bounding box coordinates. - class_map (Optional[Dict[int, str]]): A mapping ofclass IDs (int) to class - names (str). If None, a new mapping is created dynamically. + class_map: A mapping of class IDs to class names. If None, a new + mapping is created dynamically. Returns: - Detections: A new Detections object. + A new Detections object. Example: ```python @@ -873,10 +877,10 @@ class Detections: xyxy, confidences, class_ids = [], [], [] is_dynamic_mapping = class_map is None - if is_dynamic_mapping: + if class_map is None: class_map = {} - class_map = {value: key for key, value in class_map.items()} + inverted_map: dict[str, int] = {value: key for key, value in class_map.items()} for detection in azure_result["objectsResult"]["values"]: bbox = detection["boundingBox"] @@ -890,17 +894,17 @@ class Detections: for tag in tags: confidence = tag["confidence"] - class_name = tag["name"] - class_id = class_map.get(class_name, None) + class_name: str = tag["name"] + class_id_val: int | None = inverted_map.get(class_name, None) - if is_dynamic_mapping and class_id is None: - class_id = len(class_map) - class_map[class_name] = class_id + if is_dynamic_mapping and class_id_val is None: + class_id_val = len(inverted_map) + inverted_map[class_name] = class_id_val - if class_id is not None: + if class_id_val is not None: xyxy.append([x0, y0, x1, y1]) confidences.append(confidence) - class_ids.append(class_id) + class_ids.append(class_id_val) if len(xyxy) == 0: return Detections.empty() @@ -912,17 +916,17 @@ class Detections: ) @classmethod - def from_paddledet(cls, paddledet_result) -> Detections: + def from_paddledet(cls, paddledet_result: Any) -> Detections: """ Creates a Detections instance from [PaddleDetection](https://github.com/PaddlePaddle/PaddleDetection) inference result. Args: - paddledet_result (List[dict]): The output Results instance from PaddleDet + paddledet_result: The output Results instance from PaddleDet. Returns: - Detections: A new Detections object. + A new Detections object. Example: ```python @@ -958,7 +962,9 @@ class Detections: "`Detections.from_lmm` property is deprecated and will be removed in " "`supervision-0.31.0`. Use Detections.from_vlm instead." ) - def from_lmm(cls, lmm: LMM | str, result: str | dict, **kwargs: Any) -> Detections: + def from_lmm( + cls, lmm: LMM | str, result: str | dict[str, Any], **kwargs: Any + ) -> Detections: """ !!! deprecated "Deprecated" `Detections.from_lmm` is **deprecated** and will be removed in `supervision-0.31.0`. @@ -978,12 +984,12 @@ class Detections: | DeepSeek-VL2 | `DEEPSEEK_VL_2` | detection | `resolution_wh` | `classes` | Args: - lmm (Union[LMM, str]): The type of LMM (Large Multimodal Model) to use. - result (str): The result string containing the detection data. - **kwargs (Any): Additional keyword arguments required by the specified LMM. + lmm: The type of LMM (Large Multimodal Model) to use. + result: The result string containing the detection data. + **kwargs: Additional keyword arguments required by the specified LMM. Returns: - Detections: A new Detections object. + A new Detections object. Raises: ValueError: If the LMM is invalid, required arguments are missing, or @@ -1440,7 +1446,9 @@ class Detections: return cls.from_vlm(vlm=vlm, result=result, **kwargs) @classmethod - def from_vlm(cls, vlm: VLM | str, result: str | dict, **kwargs: Any) -> Detections: + def from_vlm( + cls, vlm: VLM | str, result: str | dict[str, Any], **kwargs: Any + ) -> Detections: """ Creates a Detections object from the given result string based on the specified @@ -1458,12 +1466,12 @@ class Detections: | DeepSeek-VL2 | `DEEPSEEK_VL_2` | detection | `resolution_wh` | `classes` | Args: - vlm (Union[VLM, str]): The type of VLM (Vision Language Model) to use. - result (str): The result string containing the detection data. - **kwargs (Any): Additional keyword arguments required by the specified VLM. + vlm: The type of VLM (Vision Language Model) to use. + result: The result string containing the detection data. + **kwargs: Additional keyword arguments required by the specified VLM. Returns: - Detections: A new Detections object. + A new Detections object. Raises: ValueError: If the VLM is invalid, required arguments are missing, or @@ -1863,28 +1871,41 @@ class Detections: vlm = validate_vlm_parameters(vlm, result, kwargs) if vlm == VLM.PALIGEMMA: + assert isinstance(result, str) xyxy, class_id, class_name = from_paligemma(result, **kwargs) - data = {CLASS_NAME_DATA_FIELD: class_name} + data: dict[str, npt.NDArray[np.generic] | list[Any]] = { + CLASS_NAME_DATA_FIELD: class_name, + } return cls(xyxy=xyxy, class_id=class_id, data=data) if vlm == VLM.QWEN_2_5_VL: + assert isinstance(result, str) xyxy, class_id, class_name = from_qwen_2_5_vl(result, **kwargs) data = {CLASS_NAME_DATA_FIELD: class_name} - confidence = np.ones(len(xyxy), dtype=float) - return cls(xyxy=xyxy, class_id=class_id, confidence=confidence, data=data) + confidence_arr: npt.NDArray[np.floating[Any]] = np.ones( + len(xyxy), dtype=float + ) + return cls( + xyxy=xyxy, class_id=class_id, confidence=confidence_arr, data=data + ) if vlm == VLM.QWEN_3_VL: + assert isinstance(result, str) xyxy, class_id, class_name = from_qwen_3_vl(result, **kwargs) data = {CLASS_NAME_DATA_FIELD: class_name} - confidence = np.ones(len(xyxy), dtype=float) - return cls(xyxy=xyxy, class_id=class_id, confidence=confidence, data=data) + confidence_arr = np.ones(len(xyxy), dtype=float) + return cls( + xyxy=xyxy, class_id=class_id, confidence=confidence_arr, data=data + ) if vlm == VLM.DEEPSEEK_VL_2: + assert isinstance(result, str) xyxy, class_id, class_name = from_deepseek_vl_2(result, **kwargs) data = {CLASS_NAME_DATA_FIELD: class_name} return cls(xyxy=xyxy, class_id=class_id, data=data) if vlm == VLM.FLORENCE_2: + assert isinstance(result, dict) xyxy, labels, mask, xyxyxyxy = from_florence_2(result, **kwargs) if len(xyxy) == 0: return cls.empty() @@ -1898,31 +1919,32 @@ class Detections: return cls(xyxy=xyxy, mask=mask, data=data) if vlm == VLM.GOOGLE_GEMINI_2_0: + assert isinstance(result, str) xyxy, class_id, class_name = from_google_gemini_2_0(result, **kwargs) data = {CLASS_NAME_DATA_FIELD: class_name} return cls(xyxy=xyxy, class_id=class_id, data=data) if vlm == VLM.MOONDREAM: + assert isinstance(result, dict) xyxy = from_moondream(result, **kwargs) return cls(xyxy=xyxy) if vlm == VLM.GOOGLE_GEMINI_2_5: - xyxy, class_id, class_name, confidence, mask = from_google_gemini_2_5( - result, **kwargs - ) - data = {CLASS_NAME_DATA_FIELD: class_name} + assert isinstance(result, str) + gemini_result = from_google_gemini_2_5(result, **kwargs) + data = {CLASS_NAME_DATA_FIELD: gemini_result[2]} return cls( - xyxy=xyxy, - class_id=class_id, - mask=mask, - confidence=confidence, + xyxy=gemini_result[0], + class_id=gemini_result[1], + mask=gemini_result[4], + confidence=gemini_result[3], data=data, ) return cls.empty() @classmethod - def from_easyocr(cls, easyocr_results: list) -> Detections: + def from_easyocr(cls, easyocr_results: list[Any]) -> Detections: """ Create a Detections object from the [EasyOCR](https://github.com/JaidedAI/EasyOCR) result. @@ -1930,10 +1952,10 @@ class Detections: Results are placed in the `data` field with the key `"class_name"`. Args: - easyocr_results (List): The output Results instance from EasyOCR + easyocr_results: The output Results instance from EasyOCR. Returns: - Detections: A new Detections object. + A new Detections object. Example: ```python @@ -1968,7 +1990,7 @@ class Detections: ) @classmethod - def from_ncnn(cls, ncnn_results) -> Detections: + def from_ncnn(cls, ncnn_results: Any) -> Detections: """ Creates a Detections instance from the [ncnn](https://github.com/Tencent/ncnn) inference result. @@ -2032,7 +2054,7 @@ class Detections: confidences, or class IDs. Returns: - (Detections): An empty Detections object. + An empty Detections object. Example: ```python @@ -2054,7 +2076,7 @@ class Detections: empty_detections = Detections.empty() empty_detections.data = self.data empty_detections.metadata = self.metadata - return self == empty_detections + return bool(self == empty_detections) @classmethod def merge(cls, detections_list: list[Detections]) -> Detections: @@ -2128,7 +2150,7 @@ class Detections: xyxy = np.vstack([d.xyxy for d in detections_list]) - def stack_or_none(name: str): + def stack_or_none(name: str) -> npt.NDArray[np.generic] | None: if all(d.__getattribute__(name) is None for d in detections_list): return None if any(d.__getattribute__(name) is None for d in detections_list): @@ -2159,7 +2181,7 @@ class Detections: metadata=metadata, ) - def get_anchors_coordinates(self, anchor: Position) -> np.ndarray: + def get_anchors_coordinates(self, anchor: Position) -> npt.NDArray[np.generic]: """ Calculates and returns the coordinates of a specific anchor point within the bounding boxes defined by the `xyxy` attribute. The anchor @@ -2167,12 +2189,11 @@ class Detections: such as `CENTER`, `CENTER_LEFT`, `BOTTOM_RIGHT`, etc. Args: - anchor (Position): An enum specifying the position of the anchor point - within the bounding box. Supported positions are defined in the - `Position` enum. + anchor: An enum specifying the position of the anchor point within the + bounding box. Supported positions are defined in the `Position` enum. Returns: - np.ndarray: An array of shape `(n, 2)`, where `n` is the number of bounding + An array of shape `(n, 2)`, where `n` is the number of bounding boxes. Each row contains the `[x, y]` coordinates of the specified anchor point for the corresponding bounding box. @@ -2226,8 +2247,8 @@ class Detections: raise ValueError(f"{anchor} is not supported.") def __getitem__( - self, index: int | slice | list[int] | np.ndarray | str - ) -> Detections | list | np.ndarray | None: + self, index: int | slice | list[int] | npt.NDArray[np.generic] | str + ) -> Detections | list[Any] | npt.NDArray[np.generic] | None: """ Get a subset of the Detections object or access an item from its data field. @@ -2237,12 +2258,11 @@ class Detections: the data dictionary. Args: - index (Union[int, slice, List[int], np.ndarray, str]): The index, indices, - or key to access a subset of the Detections or an item from the data. + index: The index, indices, or key to access a subset of the Detections + or an item from the data. Returns: - Union[Detections, Any]: A subset of the Detections object or an item from - the data field. + A subset of the Detections object or an item from the data field. Example: ```python @@ -2275,13 +2295,13 @@ class Detections: metadata=self.metadata, ) - 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 Detections 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. Example: ```python @@ -2311,16 +2331,16 @@ class Detections: self.data[key] = value @property - def area(self) -> np.ndarray: + def area(self) -> npt.NDArray[np.generic]: """ Calculate the area of each detection in the set of object detections. If masks field is defined property returns are of each mask. If only box is given property return area of each box. Returns: - np.ndarray: An array of floats containing the area of each detection - in the format of `(area_1, area_2, , area_n)`, - where n is the number of detections. + An array of floats containing the area of each detection + in the format of `(area_1, area_2, ..., area_n)`, + where n is the number of detections. """ if self.mask is not None: return np.array([np.sum(mask) for mask in self.mask]) @@ -2328,25 +2348,25 @@ class Detections: return self.box_area @property - def box_area(self) -> np.ndarray: + def box_area(self) -> npt.NDArray[np.generic]: """ Calculate the area of each bounding box in the set of object detections. Returns: - np.ndarray: An array of floats containing the area of each bounding - box in the format of `(area_1, area_2, , area_n)`, + An array of floats containing the area of each bounding + box in the format of `(area_1, area_2, ..., area_n)`, where n is the number of detections. """ return (self.xyxy[:, 3] - self.xyxy[:, 1]) * (self.xyxy[:, 2] - self.xyxy[:, 0]) @property - def box_aspect_ratio(self) -> np.ndarray: + def box_aspect_ratio(self) -> npt.NDArray[np.generic]: """ Compute the aspect ratio (width divided by height) for each bounding box. Returns: - np.ndarray: Array of shape `(N,)` containing aspect ratios, where `N` is the - number of boxes (width / height for each box). + Array of shape `(N,)` containing aspect ratios, where `N` is the + number of boxes (width / height for each box). Examples: ```python @@ -2387,17 +2407,17 @@ class Detections: from a segmentation model, the IoU mask is applied. Otherwise, box IoU is used. Args: - threshold (float): The intersection-over-union threshold - to use for non-maximum suppression. I'm the lower the value the more + threshold: The intersection-over-union threshold + to use for non-maximum suppression. The lower the value the more restrictive the NMS becomes. Defaults to 0.5. - class_agnostic (bool): Whether to perform class-agnostic + class_agnostic: Whether to perform class-agnostic non-maximum suppression. If True, the class_id of each detection will be ignored. Defaults to False. - overlap_metric (OverlapMetric): Metric used to compute the degree of + overlap_metric: Metric used to compute the degree of overlap between pairs of masks or boxes (e.g., IoU, IoS). Returns: - Detections: A new Detections object containing the subset of detections + A new Detections object containing the subset of detections after non-maximum suppression. Raises: @@ -2440,7 +2460,7 @@ class Detections: overlap_metric=overlap_metric, ) - return self[indices] + return cast(Detections, self[indices]) def with_nmm( self, @@ -2452,16 +2472,16 @@ class Detections: Perform non-maximum merging on the current set of object detections. Args: - threshold (float): The intersection-over-union threshold + threshold: The intersection-over-union threshold to use for non-maximum merging. Defaults to 0.5. - class_agnostic (bool): Whether to perform class-agnostic + class_agnostic: Whether to perform class-agnostic non-maximum merging. If True, the class_id of each detection will be ignored. Defaults to False. - overlap_metric (OverlapMetric): Metric used to compute the degree of + overlap_metric: Metric used to compute the degree of overlap between pairs of masks or boxes (e.g., IoU, IoS). Returns: - Detections: A new Detections object containing the subset of detections + A new Detections object containing the subset of detections after non-maximum merging. Raises: @@ -2506,9 +2526,9 @@ class Detections: overlap_metric=overlap_metric, ) - result = [] + result: list[Detections] = [] for merge_group in merge_groups: - unmerged_detections = [self[i] for i in merge_group] + unmerged_detections = [cast(Detections, self[i]) for i in merge_group] merged_detections = merge_inner_detections_objects_without_iou( unmerged_detections ) @@ -2521,7 +2541,7 @@ def merge_inner_detection_object_pair( detections_1: Detections, detections_2: Detections ) -> Detections: """ - Merges two Detections object into a single Detections object. + Merges two Detections objects into a single Detections object. Assumes each Detections contains exactly one object. A `winning` detection is determined based on the confidence score of the two @@ -2529,18 +2549,16 @@ def merge_inner_detection_object_pair( `class_id`, `tracker_id`, and `data` to include in the merged Detections object. The resulting `confidence` of the merged object is calculated by the weighted - contribution of ea detection to the merged object. + contribution of each detection to the merged object. The bounding boxes and masks of the two input detections are merged into a single bounding box and mask, respectively. Args: - detections_1 (Detections): - The first Detections object - detections_2 (Detections): - The second Detections object + detections_1: The first Detections object. + detections_2: The second Detections object. Returns: - Detections: A new Detections object, with merged attributes. + A new Detections object, with merged attributes. Raises: ValueError: If the input Detections objects do not have exactly 1 detected @@ -2572,6 +2590,8 @@ def merge_inner_detection_object_pair( if detections_1.confidence is None and detections_2.confidence is None: merged_confidence = None else: + assert detections_1.confidence is not None + assert detections_2.confidence is not None detection_1_area = (xyxy_1[2] - xyxy_1[0]) * (xyxy_1[3] - xyxy_1[1]) detections_2_area = (xyxy_2[2] - xyxy_2[0]) * (xyxy_2[3] - xyxy_2[1]) merged_confidence = ( @@ -2589,7 +2609,7 @@ def merge_inner_detection_object_pair( else: merged_mask = np.logical_or(detections_1.mask, detections_2.mask) - if detections_1.confidence is None and detections_2.confidence is None: + if detections_1.confidence is None or detections_2.confidence is None: winning_detection = detections_1 elif detections_1.confidence[0] >= detections_2.confidence[0]: winning_detection = detections_1 @@ -2611,7 +2631,7 @@ def merge_inner_detection_object_pair( def merge_inner_detections_objects( detections: list[Detections], - threshold=0.5, + threshold: float = 0.5, overlap_metric: OverlapMetric = OverlapMetric.IOU, ) -> Detections: """ diff --git a/src/supervision/detection/line_zone.py b/src/supervision/detection/line_zone.py index 574860f3..040f2188 100644 --- a/src/supervision/detection/line_zone.py +++ b/src/supervision/detection/line_zone.py @@ -5,7 +5,7 @@ import warnings from collections import Counter, defaultdict, deque from collections.abc import Iterable from functools import lru_cache -from typing import Any, Literal +from typing import Any, Literal, cast import cv2 import numpy as np @@ -41,14 +41,16 @@ class LineZone: tracking into your inference pipeline. Attributes: - in_count (int): The number of objects that have crossed the line from outside + in_count: The number of objects that have crossed the line from outside to inside. - out_count (int): The number of objects that have crossed the line from inside + out_count: The number of objects that have crossed the line from inside to outside. - in_count_per_class (Dict[int, int]): Number of objects of each class that have - crossed the line from outside to inside. - out_count_per_class (Dict[int, int]): Number of objects of each class that have - crossed the line from inside to outside. + in_count_per_class: Number of objects of each class that have + crossed the line from outside to inside, keyed by `class_id` + (`int` for classified detections, `None` for unclassified ones). + out_count_per_class: Number of objects of each class that have + crossed the line from inside to outside, keyed by `class_id` + (`int` for classified detections, `None` for unclassified ones). Example: ```python @@ -86,27 +88,25 @@ class LineZone: ): """ Args: - start (Point): The starting point of the line. - end (Point): The ending point of the line. - triggering_anchors (List[sv.Position]): A list of positions - specifying which anchors of the detections bounding box - to consider when deciding on whether the detection - has passed the line counter or not. By default, this - contains the four corners of the detection's bounding box - minimum_crossing_threshold (int): Detection needs to be seen - on the other side of the line for this many frames to be - considered as having crossed the line. This is useful when - dealing with unstable bounding boxes or when detections - may linger on the line. + start: The starting point of the line. + end: The ending point of the line. + triggering_anchors: A list of positions specifying which anchors of + the detections bounding box to consider when deciding on whether + the detection has passed the line counter or not. By default, + this contains the four corners of the detection's bounding box. + minimum_crossing_threshold: Detection needs to be seen on the other + side of the line for this many frames to be considered as having + crossed the line. This is useful when dealing with unstable + bounding boxes or when detections may linger on the line. """ self.vector = Vector(start=start, end=end) self.limits = self._calculate_region_of_interest_limits(vector=self.vector) self.crossing_history_length = max(2, minimum_crossing_threshold + 1) - self.crossing_state_history: dict[tuple[int, int], deque[bool]] = defaultdict( - lambda: deque(maxlen=self.crossing_history_length) + self.crossing_state_history: dict[tuple[int, int | None], deque[bool]] = ( + defaultdict(lambda: deque(maxlen=self.crossing_history_length)) ) - self._in_count_per_class: Counter = Counter() - self._out_count_per_class: Counter = Counter() + self._in_count_per_class: Counter[int | None] = Counter() + self._out_count_per_class: Counter[int | None] = Counter() self.triggering_anchors = triggering_anchors if not list(self.triggering_anchors): raise ValueError("Triggering anchors cannot be empty.") @@ -121,20 +121,21 @@ class LineZone: return sum(self._out_count_per_class.values()) @property - def in_count_per_class(self) -> dict[int, int]: + def in_count_per_class(self) -> dict[int | None, int]: return dict(self._in_count_per_class) @property - def out_count_per_class(self) -> dict[int, int]: + def out_count_per_class(self) -> dict[int | None, int]: return dict(self._out_count_per_class) - def trigger(self, detections: Detections) -> tuple[np.ndarray, np.ndarray]: + def trigger( + self, detections: Detections + ) -> tuple[npt.NDArray[np.bool_], npt.NDArray[np.bool_]]: """ Update the `in_count` and `out_count` based on the objects that cross the line. Args: - detections (Detections): A list of detections for which to update the - counts. + detections: A Detections object for which to update the counts. Returns: A tuple of two boolean NumPy arrays. The first array indicates which @@ -255,11 +256,10 @@ class LineZone: ``` Args: - detections (Detections): The detections to check. + detections: The detections to check. Returns: - result (Tuple[np.ndarray, np.ndarray, np.ndarray]): - All 3 arrays are boolean arrays of shape (N, ) where N is the + All 3 arrays are boolean arrays of shape (N, ) where N is the number of detections. The first array, `in_limits`, indicates if the detection's anchor is within the line zone limits. The second array, `has_any_left_trigger`, indicates if the @@ -335,22 +335,22 @@ class LineZoneAnnotator: A class for drawing the `LineZone` and its detected object count on an image. - Attributes: - thickness (int): Line thickness. - color (Color): Line color. - text_thickness (int): Text thickness. - text_color (Color): Text color. - text_scale (float): Text scale. - text_offset (float): How far the text will be from the line. - text_padding (int): The empty space in the text box, surrounding the text. - custom_in_text (Optional[str]): Write something else instead of "in". - custom_out_text (Optional[str]): Write something else instead of "out". - display_in_count (bool): Pass `False` to hide the "in" count. - display_out_count (bool): Pass `False` to hide the "out" count. - display_text_box (bool): Pass `False` to hide the text background box. - text_orient_to_line (bool): ⭐ Match text orientation to the line. + Args: + thickness: Line thickness. + color: Line color. + text_thickness: Text thickness. + text_color: Text color. + text_scale: Text scale. + text_offset: How far the text will be from the line. + text_padding: The empty space in the text box, surrounding the text. + custom_in_text: Write something else instead of "in". + custom_out_text: Write something else instead of "out". + display_in_count: Pass `False` to hide the "in" count. + display_out_count: Pass `False` to hide the "out" count. + display_text_box: Pass `False` to hide the text background box. + text_orient_to_line: Match text orientation to the line. Recommended to set to `True`. - text_centered (bool): Pass `False` to disable text centering. Useful + text_centered: Pass `False` to disable text centering. Useful when the label overlaps something important. """ @@ -369,17 +369,18 @@ class LineZoneAnnotator: self.text_orient_to_line: bool = text_orient_to_line self.text_centered: bool = text_centered - def annotate(self, frame: np.ndarray, line_counter: LineZone) -> np.ndarray: + def annotate( + self, frame: npt.NDArray[np.uint8], line_counter: LineZone + ) -> npt.NDArray[np.uint8]: """ Draws the line on the frame using the line zone provided. - Attributes: - frame (np.ndarray): The image on which the line will be drawn. - line_counter (LineZone): The line zone - that will be used to draw the line. + Args: + frame: The image on which the line will be drawn. + line_counter: The line zone that will be used to draw the line. Returns: - (np.ndarray): The image with the line drawn on it. + The image with the line drawn on it. """ line_start = line_counter.vector.start.as_xy_int_tuple() @@ -443,10 +444,10 @@ class LineZoneAnnotator: Calculate the line counter angle (in degrees). Args: - line_zone (LineZone): The line zone object. + line_zone: The line zone object. Returns: - (float): Line counter angle, in degrees. + Line counter angle, in degrees. """ start_point = line_zone.vector.start.as_xy_int_tuple() end_point = line_zone.vector.end.as_xy_int_tuple() @@ -475,15 +476,15 @@ class LineZoneAnnotator: Calculate insertion anchor in frame to position the center of the count image. Args: - line_zone (LineZone): The line counter object used for counting. - text_width (int): Text width. - text_height (int): Text height. - is_in_count (bool): Whether the count should be placed over or below line. - label_dimension (int): Size of the label image. Assumes the + line_zone: The line counter object used for counting. + text_width: Text width. + text_height: Text height. + is_in_count: Whether the count should be placed over or below line. + label_dimension: Size of the label image. Assumes the label is rectangular. Returns: - (Tuple[int, int]): xy, point in an image where the label will be placed. + xy, point in an image where the label will be placed. """ line_angle = self._get_line_angle(line_zone) @@ -529,24 +530,24 @@ class LineZoneAnnotator: def _draw_basic_label( self, - frame: np.ndarray, + frame: npt.NDArray[np.uint8], line_center: Point, text: str, is_in_count: bool, - ) -> np.ndarray: + ) -> npt.NDArray[np.uint8]: """ Draw the count label on the frame. For example: "out: 7". The label contains horizontal text and is not rotated. Args: - frame (np.ndarray): The entire scene, on which the label will be placed. - line_center (Point): The center of the line zone. - text (str): The text that will be drawn. - is_in_count (bool): Whether to display the in count (above line) + frame: The entire scene, on which the label will be placed. + line_center: The center of the line zone. + text: The text that will be drawn. + is_in_count: Whether to display the in count (above line) or out count (below line). Returns: - (np.ndarray): The scene with the label drawn on it. + The scene with the label drawn on it. """ _, text_height = cv2.getTextSize( text, cv2.FONT_HERSHEY_SIMPLEX, self.text_scale, self.text_thickness @@ -572,25 +573,24 @@ class LineZoneAnnotator: def _draw_oriented_label( self, - frame: np.ndarray, + frame: npt.NDArray[np.uint8], line_zone: LineZone, text: str, is_in_count: bool, - ) -> np.ndarray: + ) -> npt.NDArray[np.uint8]: """ Draw the count label on the frame. For example: "out: 7". The label is oriented to match the line angle. Args: - frame (np.ndarray): The entire scene, on which the label will be placed. - line_zone (LineZone): The line zone responsible for counting - objects crossing it. - text (str): The text that will be drawn. - is_in_count (bool): Whether to display the in count (above line) + frame: The entire scene, on which the label will be placed. + line_zone: The line zone responsible for counting objects crossing it. + text: The text that will be drawn. + is_in_count: Whether to display the in count (above line) or out count (below line). Returns: - (np.ndarray): The scene with the label drawn on it. + The scene with the label drawn on it. """ line_angle_degrees = self._get_line_angle(line_zone) @@ -634,22 +634,22 @@ class LineZoneAnnotator: text_box_show: bool, text_box_color: Color, line_angle_degrees: float, - ) -> np.ndarray: + ) -> npt.NDArray[np.uint8]: """ Create the small text box displaying line zone count. E.g. "out: 7". Args: - text (str): The text to display. - text_scale (float): The scale of the text. - text_thickness (int): The thickness of the text. - text_padding (int): The padding around the text. - text_color (Color): The color of the text. - text_box_show (bool): Whether to display the text box. - text_box_color (Color): The color of the text box. - line_angle_degrees (float): The angle of the line in degrees. + text: The text to display. + text_scale: The scale of the text. + text_thickness: The thickness of the text. + text_padding: The padding around the text. + text_color: The color of the text. + text_box_show: Whether to display the text box. + text_box_color: The color of the text box. + line_angle_degrees: The angle of the line in degrees. Returns: - (np.ndarray): The label of shape (H, W, 4), in BGRA format. + The label of shape (H, W, 4), in BGRA format. """ text_width, text_height = cv2.getTextSize( text, cv2.FONT_HERSHEY_SIMPLEX, text_scale, text_thickness @@ -693,7 +693,7 @@ class LineZoneAnnotator: ) annotation = cv2.warpAffine(annotation, rotation_matrix, annotation_shape) - return annotation + return cast(npt.NDArray[np.uint8], annotation) class LineZoneAnnotatorMulticlass: @@ -719,15 +719,15 @@ class LineZoneAnnotatorMulticlass: Draw a table showing how many items of each class crossed each line. Args: - table_position (Position): The position of the table. - table_color (Color): The color of the table. - table_margin (int): The margin of the table from the image border. - table_padding (int): The padding of the table. - table_max_width (int): The maximum width of the table. - text_color (Color): The color of the text. - text_scale (float): The scale of the text. - text_thickness (int): The thickness of the text. - force_draw_class_ids (bool): Instead of writing the class names, + table_position: The position of the table. + table_color: The color of the table. + table_margin: The margin of the table from the image border. + table_padding: The padding of the table. + table_max_width: The maximum width of the table. + text_color: The color of the text. + text_scale: The scale of the text. + text_thickness: The thickness of the text. + force_draw_class_ids: Instead of writing the class names, on the table, write the class IDs. E.g. instead of `person: 6`, write `0: 6`. """ @@ -754,21 +754,21 @@ class LineZoneAnnotatorMulticlass: def annotate( self, - frame: np.ndarray, + frame: npt.NDArray[np.uint8], line_zones: list[LineZone], line_zone_labels: list[str] | None = None, - ) -> np.ndarray: + ) -> npt.NDArray[np.uint8]: """ Draws a table with the number of objects of each class that crossed each line. - Attributes: - frame (np.ndarray): The image on which the table will be drawn. - line_zones (List[LineZone]): The line zones to be annotated. - line_zone_labels (Optional[List[str]]): The labels, one for each - line zone. If not provided, the default labels will be used. + Args: + frame: The image on which the table will be drawn. + line_zones: The line zones to be annotated. + line_zone_labels: The labels, one for each line zone. If not + provided, the default labels will be used. Returns: - (np.ndarray): The image with the table drawn on it. + The image with the table drawn on it. """ if line_zone_labels is None: @@ -790,11 +790,12 @@ class LineZoneAnnotatorMulticlass: text_lines.append(f" {direction}:") for class_id, count in count_per_class.items(): - class_name = ( - class_id_to_name.get(class_id, str(class_id)) - if not self.force_draw_class_ids - else str(class_id) - ) + if self.force_draw_class_ids: + class_name = str(class_id) + elif class_id is None: + class_name = "None" + else: + class_name = class_id_to_name.get(class_id, str(class_id)) text_lines.append(f" {class_name}: {count}") table_width, table_height = 0, 0 diff --git a/src/supervision/detection/vlm.py b/src/supervision/detection/vlm.py index 814e417c..0576659b 100644 --- a/src/supervision/detection/vlm.py +++ b/src/supervision/detection/vlm.py @@ -6,9 +6,10 @@ import io import json import re from enum import Enum -from typing import Any +from typing import Any, cast import numpy as np +import numpy.typing as npt from PIL import Image from supervision.detection.utils.boxes import denormalize_boxes @@ -45,8 +46,8 @@ class LMM(Enum): MOONDREAM = "moondream" @classmethod - def list(cls): - return list(map(lambda c: c.value, cls)) + def list(cls) -> list[str]: + return [c.value for c in cls] @classmethod def from_value(cls, value: LMM | str) -> LMM: @@ -88,8 +89,8 @@ class VLM(Enum): MOONDREAM = "moondream" @classmethod - def list(cls): - return list(map(lambda c: c.value, cls)) + def list(cls) -> list[str]: + return [c.value for c in cls] @classmethod def from_value(cls, value: VLM | str) -> VLM: @@ -164,7 +165,7 @@ def validate_vlm_parameters(vlm: VLM | str, result: Any, kwargs: dict[str, Any]) kwargs: Dictionary of arguments to validate against required/allowed lists. Returns: - VLM: The validated VLM enum value. + The validated VLM enum value. Raises: ValueError: If the VLM, result type, or arguments are invalid. @@ -197,7 +198,7 @@ def validate_vlm_parameters(vlm: VLM | str, result: Any, kwargs: dict[str, Any]) def from_paligemma( result: str, resolution_wh: tuple[int, int], classes: list[str] | None = None -) -> tuple[np.ndarray, np.ndarray | None, np.ndarray]: +) -> tuple[npt.NDArray[Any], npt.NDArray[Any] | None, npt.NDArray[Any]]: """ Parse bounding boxes from paligemma-formatted text, scale them to the specified resolution, and optionally filter by classes. @@ -209,13 +210,10 @@ def from_paligemma( in this list are filtered out. Returns: - xyxy (np.ndarray): An array of shape `(n, 4)` containing - the bounding boxes coordinates in format `[x1, y1, x2, y2]`. - class_id (Optional[np.ndarray]): An array of shape `(n,)` containing - the class indices for each bounding box (or `None` if classes is not - provided). - class_name (np.ndarray): An array of shape `(n,)` containing - the class labels for each bounding box. + A tuple of `(xyxy, class_id, class_name)` where `xyxy` is an array of + shape `(n, 4)` in format `[x1, y1, x2, y2]`, `class_id` is an + optional array of shape `(n,)` with class indices, and `class_name` + is an array of shape `(n,)` with class labels. """ w, h = validate_resolution(resolution_wh) @@ -252,7 +250,7 @@ def recover_truncated_qwen_2_5_vl_response(text: str) -> Any | None: malformed, cleans trailing commas, and attempts to parse it into a Python object. Args: - text (str): Raw text containing the JSON snippet possibly truncated or + text: Raw text containing the JSON snippet possibly truncated or incomplete. Returns: @@ -293,7 +291,7 @@ def from_qwen_2_5_vl( input_wh: tuple[int, int], resolution_wh: tuple[int, int], classes: list[str] | None = None, -) -> tuple[np.ndarray, np.ndarray | None, np.ndarray]: +) -> tuple[npt.NDArray[Any], npt.NDArray[Any] | None, npt.NDArray[Any]]: """ Parse and rescale bounding boxes and class labels from Qwen-2.5-VL JSON output. @@ -306,20 +304,18 @@ def from_qwen_2_5_vl( ``` Args: - result (str): String containing Qwen-2.5-VL JSON bounding box and label data. - input_wh (tuple[int, int]): Width and height of the coordinate space where boxes + result: String containing Qwen-2.5-VL JSON bounding box and label data. + input_wh: Width and height of the coordinate space where boxes are normalized. - resolution_wh (tuple[int, int]): Target width and height to scale bounding - boxes. - classes (list[str] or None): Optional list of valid class names to filter - results. If provided, only boxes with labels in this list are returned. + resolution_wh: Target width and height to scale bounding boxes. + classes: Optional list of valid class names to filter results. If provided, + only boxes with labels in this list are returned. Returns: - xyxy (np.ndarray): Array of shape `(N, 4)` with rescaled bounding boxes in - `(x_min, y_min, x_max, y_max)` format. - class_id (np.ndarray or None): Array of shape `(N,)` with indices of classes, - or `None` if no filtering applied. - class_name (np.ndarray): Array of shape `(N,)` with class names as strings. + A tuple of `(xyxy, class_id, class_name)` where `xyxy` is an array of + shape `(N, 4)` in `(x_min, y_min, x_max, y_max)` format, `class_id` + is an optional array of shape `(N,)` with class indices, and + `class_name` is an array of shape `(N,)` with class names. """ in_w, in_h = validate_resolution(input_wh) @@ -386,23 +382,20 @@ def from_qwen_3_vl( result: str, resolution_wh: tuple[int, int], classes: list[str] | None = None, -) -> tuple[np.ndarray, np.ndarray | None, np.ndarray]: +) -> tuple[npt.NDArray[Any], npt.NDArray[Any] | None, npt.NDArray[Any]]: """ Parse and scale bounding boxes from Qwen-3-VL style JSON output. Args: - result (str): String containing the Qwen-3-VL JSON output. - resolution_wh (tuple[int, int]): Target resolution `(width, height)` to - scale bounding boxes. - classes (list[str] or None): Optional list of valid classes to filter - results. + result: String containing the Qwen-3-VL JSON output. + resolution_wh: Target resolution `(width, height)` to scale bounding boxes. + classes: Optional list of valid classes to filter results. Returns: - xyxy (np.ndarray): Array of bounding boxes with shape `(N, 4)` in - `(x_min, y_min, x_max, y_max)` format scaled to `resolution_wh`. - class_id (np.ndarray or None): Array of class indices for each box, or - None if no filtering by classes. - class_name (np.ndarray): Array of class names as strings. + A tuple of `(xyxy, class_id, class_name)` where `xyxy` is an array of + shape `(N, 4)` in `(x_min, y_min, x_max, y_max)` format scaled to + `resolution_wh`, `class_id` is an optional array of class indices, + and `class_name` is an array of class names. """ return from_qwen_2_5_vl( result=result, @@ -414,7 +407,7 @@ def from_qwen_3_vl( def from_deepseek_vl_2( result: str, resolution_wh: tuple[int, int], classes: list[str] | None = None -) -> tuple[np.ndarray, np.ndarray | None, np.ndarray]: +) -> tuple[npt.NDArray[Any], npt.NDArray[Any] | None, npt.NDArray[Any]]: """ Parse bounding boxes from deepseek-vl2-formatted text, scale them to the specified resolution, and optionally filter by classes. @@ -435,13 +428,10 @@ def from_deepseek_vl_2( in this list are filtered out. Returns: - xyxy (np.ndarray): An array of shape `(n, 4)` containing - the bounding boxes coordinates in format `[x1, y1, x2, y2]`. - class_id (Optional[np.ndarray]): An array of shape `(n,)` containing - the class indices for each bounding box (or `None` if classes is not - provided). - class_name (np.ndarray): An array of shape `(n,)` containing - the class labels for each bounding box. + A tuple of `(xyxy, class_id, class_name)` where `xyxy` is an array of + shape `(n, 4)` in format `[x1, y1, x2, y2]`, `class_id` is an + optional array of shape `(n,)` with class indices, and `class_name` + is an array of shape `(n,)` with class labels. """ # noqa: E501 width, height = resolution_wh @@ -486,8 +476,13 @@ def from_deepseek_vl_2( def from_florence_2( - result: dict, resolution_wh: tuple[int, int] -) -> tuple[np.ndarray, np.ndarray | None, np.ndarray | None, np.ndarray | None]: + result: dict[str, Any], resolution_wh: tuple[int, int] +) -> tuple[ + npt.NDArray[Any], + npt.NDArray[Any] | None, + npt.NDArray[Any] | None, + npt.NDArray[Any] | None, +]: """ Parse results from the Florence 2 multi-model model. https://huggingface.co/microsoft/Florence-2-large @@ -497,14 +492,12 @@ def from_florence_2( resolution_wh: (output_width, output_height) to which we rescale the boxes. Returns: - xyxy (np.ndarray): An array of shape `(n, 4)` containing - the bounding boxes coordinates in format `[x1, y1, x2, y2]` - labels: (Optional[np.ndarray]): An array of shape `(n,)` containing - the class labels for each bounding box - masks: (Optional[np.ndarray]): An array of shape `(n, h, w)` containing - the segmentation masks for each bounding box - obb_boxes: (Optional[np.ndarray]): An array of shape `(n, 4, 2)` containing - oriented bounding boxes. + A tuple of `(xyxy, labels, masks, obb_boxes)` where `xyxy` is an array + of shape `(n, 4)` in format `[x1, y1, x2, y2]`, `labels` is an + optional array of shape `(n,)` with class labels, `masks` is an + optional array of shape `(n, h, w)` with segmentation masks, and + `obb_boxes` is an optional array of shape `(n, 4, 2)` with oriented + bounding boxes. """ assert len(result) == 1, f"Expected result with a single element. Got: {result}" task = next(iter(result.keys())) @@ -582,7 +575,7 @@ def from_google_gemini_2_0( result: str, resolution_wh: tuple[int, int], classes: list[str] | None = None, -) -> tuple[np.ndarray, np.ndarray | None, np.ndarray]: +) -> tuple[npt.NDArray[Any], npt.NDArray[Any] | None, npt.NDArray[Any]]: """ Parse and scale bounding boxes from Google Gemini style [JSON output](https://ai.google.dev/gemini-api/docs/vision?lang=python). @@ -610,13 +603,10 @@ def from_google_gemini_2_0( are filtered to only those classes found here. Returns: - xyxy (np.ndarray): An array of shape `(n, 4)` containing - the bounding boxes coordinates in format `[x1, y1, x2, y2]` - class_id (Optional[np.ndarray]): An array of shape `(n,)` containing - the class indices for each bounding box (or None if `classes` is not - provided) - class_name (np.ndarray): An array of shape `(n,)` containing - the class labels for each bounding box + A tuple of `(xyxy, class_id, class_name)` where `xyxy` is an array of + shape `(n, 4)` in format `[x1, y1, x2, y2]`, `class_id` is an + optional array of shape `(n,)` with class indices, and `class_name` + is an array of shape `(n,)` with class labels. """ @@ -670,11 +660,11 @@ def from_google_gemini_2_5( resolution_wh: tuple[int, int], classes: list[str] | None = None, ) -> tuple[ - np.ndarray, - np.ndarray | None, - np.ndarray, - np.ndarray | None, - np.ndarray | None, + npt.NDArray[Any], + npt.NDArray[Any] | None, + npt.NDArray[Any], + npt.NDArray[Any] | None, + npt.NDArray[Any] | None, ]: """ Parse and scale bounding boxes and masks from Google Gemini 2.5 style @@ -700,17 +690,13 @@ def from_google_gemini_2_5( are filtered to only those classes found here. Returns: - xyxy (np.ndarray): An array of shape `(n, 4)` containing - the bounding boxes coordinates in format `[x1, y1, x2, y2]` - class_id (np.ndarray): An array of shape `(n,)` containing - the class indices for each bounding box - class_name (np.ndarray): An array of shape `(n,)` containing - the class labels for each bounding box - confidence: Optional[np.ndarray]: An array of shape `(n,)` containing - the confidence scores for each bounding box. If not provided, - it defaults to 0.0 for each box. - masks (Optional[np.ndarray]): An array of shape `(n, h, w)` containing - the segmentation masks for each bounding box + A tuple of `(xyxy, class_id, class_name, confidence, masks)` where + `xyxy` is an array of shape `(n, 4)` in format `[x1, y1, x2, y2]`, + `class_id` is an array of shape `(n,)` with class indices, + `class_name` is an array of shape `(n,)` with class labels, + `confidence` is an optional array of shape `(n,)` with confidence + scores, and `masks` is an optional array of shape `(n, h, w)` with + segmentation masks. """ w, h = validate_resolution(resolution_wh) @@ -732,10 +718,10 @@ def from_google_gemini_2_5( None, ) - boxes_list: list = [] - labels_list: list = [] - confidence_list: list | None = [] - masks_list: list | None = [] + boxes_list: list[Any] = [] + labels_list: list[str] = [] + confidence_list: list[float] | None = [] + masks_list: list[npt.NDArray[Any]] | None = [] for item in data: if "box_2d" not in item or "label" not in item: @@ -771,7 +757,7 @@ def from_google_gemini_2_5( mask_img = mask_img.resize( (bbox_width, bbox_height), resample=Image.Resampling.BILINEAR ) - np_mask = np.zeros((h, w), dtype=bool) + np_mask: npt.NDArray[np.bool_] = np.zeros((h, w), dtype=bool) np_mask[y_min:y_max, x_min:x_max] = np.array(mask_img) > 0 masks_list.append(np_mask) else: @@ -796,7 +782,7 @@ def from_google_gemini_2_5( xyxy = np.array(boxes_list, dtype=float) class_name = np.array(labels_list) - class_id: np.ndarray + class_id: npt.NDArray[Any] if classes is not None: mask = np.array([name in classes for name in class_name], dtype=bool) @@ -828,9 +814,9 @@ def from_google_gemini_2_5( def from_moondream( - result: dict, + result: dict[str, Any], resolution_wh: tuple[int, int], -) -> np.ndarray: +) -> npt.NDArray[Any]: """ Parse and scale bounding boxes from moondream JSON output. @@ -855,8 +841,8 @@ def from_moondream( resolution_wh: (output_width, output_height) to which we rescale the boxes. Returns: - xyxy (np.ndarray): An array of shape `(n, 4)` containing - the bounding boxes coordinates in format `[x1, y1, x2, y2]` + An array of shape `(n, 4)` containing the bounding boxes coordinates + in format `[x1, y1, x2, y2]`. """ w, h = resolution_wh @@ -882,9 +868,12 @@ def from_moondream( xyxy.append([x_min, y_min, x_max, y_max]) if len(xyxy) == 0: - return np.empty((0, 4)) + return cast(npt.NDArray[Any], np.empty((0, 4))) - return denormalize_boxes( - np.array(xyxy).astype(np.float64), - resolution_wh=(w, h), + return cast( + npt.NDArray[Any], + denormalize_boxes( + np.array(xyxy).astype(np.float64), + resolution_wh=(w, h), + ), ) diff --git a/src/supervision/metrics/mean_average_precision.py b/src/supervision/metrics/mean_average_precision.py index f95656bb..6d063033 100644 --- a/src/supervision/metrics/mean_average_precision.py +++ b/src/supervision/metrics/mean_average_precision.py @@ -1025,7 +1025,7 @@ class COCOEvaluator: :, :, :, area_range_idx, max_100_dets_idx ] # mAP over thresholds (dimension=num_thresholds) - # Use masked array to exclude -1 values when computing mean + # Exclude -1 sentinel values when computing mean mAP_scores_all_sizes, ap_per_class_all_sizes = compute_average_precision( average_precision_all_sizes ) diff --git a/tests/dataset/formats/test_yolo.py b/tests/dataset/formats/test_yolo.py index 1e2a10eb..419bbac7 100644 --- a/tests/dataset/formats/test_yolo.py +++ b/tests/dataset/formats/test_yolo.py @@ -8,6 +8,7 @@ import pytest from supervision.dataset.formats.yolo import ( _image_name_to_annotation_name, _with_seg_mask, + detections_to_yolo_annotations, object_to_yolo, yolo_annotations_to_detections, ) @@ -295,3 +296,15 @@ def test_object_to_yolo( xyxy=xyxy, class_id=class_id, image_shape=image_shape, polygon=polygon ) assert result == expected_result + + +def test_detections_to_yolo_annotations_raises_for_non_integer_class_id() -> None: + detections = Detections( + xyxy=np.array([[100, 100, 200, 200]], dtype=np.float32), + class_id=np.array([1.9], dtype=np.float32), + ) + + with pytest.raises(ValueError, match="must be an integer"): + detections_to_yolo_annotations( + detections=detections, image_shape=(1000, 1000, 3) + ) diff --git a/tests/detection/test_line_counter.py b/tests/detection/test_line_counter.py index be819060..75355a27 100644 --- a/tests/detection/test_line_counter.py +++ b/tests/detection/test_line_counter.py @@ -2,9 +2,10 @@ from __future__ import annotations from contextlib import ExitStack as DoesNotRaise +import numpy as np import pytest -from supervision import LineZone +from supervision import LineZone, LineZoneAnnotatorMulticlass from supervision.geometry.core import Point, Position, Vector from tests.helpers import _create_detections @@ -878,3 +879,19 @@ def test_line_zone_tracker_id_reuse_with_different_classes( assert line_zone.in_count_per_class == expected_in_count_per_class assert line_zone.out_count_per_class == expected_out_count_per_class + + +def test_line_zone_annotator_multiclass_supports_none_class_id() -> None: + line_zone = LineZone(start=Point(0, 0), end=Point(0, 10)) + for xyxy in [[4, 4, 6, 6], [-6, 4, -4, 6]]: + detections = _create_detections(xyxy=[xyxy], tracker_id=[0]) + line_zone.trigger(detections) + + assert line_zone.out_count_per_class == {None: 1} + + frame = np.zeros((100, 100, 3), dtype=np.uint8) + annotator = LineZoneAnnotatorMulticlass(force_draw_class_ids=False) + annotated_frame = annotator.annotate(frame=frame.copy(), line_zones=[line_zone]) + + assert annotated_frame.shape == frame.shape + assert not np.array_equal(annotated_frame, frame)