From e4cedb68617fe1178ec66c007e9e8142fc416771 Mon Sep 17 00:00:00 2001 From: Omkar Kabde Date: Tue, 10 Mar 2026 23:21:27 +0530 Subject: [PATCH] refactor docstrings in `src/supervision/detection/tools` (#2164) * refactor docstrings in `src/supervision/detection/tools` * Refine type annotations and docstrings in `transformers.py` to include `TensorLike` protocol and enhance segmentation handling * Enhance type safety in `transformers.py` by introducing `_is_tensor_like` with `TypeGuard` and refining segmentation result handling * Refine `CSVSink` type annotations by introducing `WriterProtocol` for improved type safety * Remove `_is_tensor_like`, `TensorLike` protocol, and `TypeGuard` usage from `transformers.py`, streamline segmentation handling, and fix typos in `CSVSink` and `JSONSink` docstrings --------- 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 | 5 - src/supervision/detection/tools/csv_sink.py | 49 ++++----- .../detection/tools/inference_slicer.py | 69 ++++++------ src/supervision/detection/tools/json_sink.py | 40 +++---- .../detection/tools/polygon_zone.py | 55 +++++----- .../detection/tools/transformers.py | 103 +++++++++--------- 6 files changed, 157 insertions(+), 164 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c5fcb1dd..cf291408 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -204,12 +204,7 @@ module = [ "tests.*", "examples.*", # TODO: fix type errors in the following modules - "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.key_points.skeletons", "supervision.metrics.utils.utils", ] diff --git a/src/supervision/detection/tools/csv_sink.py b/src/supervision/detection/tools/csv_sink.py index c994b4a3..d06d8966 100644 --- a/src/supervision/detection/tools/csv_sink.py +++ b/src/supervision/detection/tools/csv_sink.py @@ -1,8 +1,12 @@ from __future__ import annotations import csv +import io import os -from typing import Any +from collections.abc import Iterable +from typing import Any, Protocol + +import numpy as np from supervision.detection.core import Detections from supervision.utils.logger import _get_logger @@ -20,6 +24,10 @@ BASE_HEADER = [ ] +class WriterProtocol(Protocol): + def writerow(self, row: Iterable[Any]) -> Any: ... + + class CSVSink: """ A utility class for saving detection data to a CSV file. This class is designed to @@ -29,11 +37,11 @@ class CSVSink: !!! tip - CSVSink allow to pass custom data alongside the detection fields, providing + CSVSink allows passing custom data alongside detection fields, providing flexibility for logging various types of information. Args: - file_name (str): The name of the CSV file where the detections will be stored. + file_name: The name of the CSV file where the detections will be stored. Defaults to 'output.csv'. Example: @@ -66,16 +74,13 @@ class CSVSink: Initialize the CSVSink instance. Args: - file_name (str): The name of the CSV file. - - Returns: - None + file_name: The name of the CSV file. """ self.file_name = file_name - self.file: open | None = None - self.writer: csv.writer | None = None + self.file: io.TextIOWrapper | None = None + self.writer: WriterProtocol | None = None self.header_written = False - self.field_names = [] + self.field_names: list[str] = [] def __enter__(self) -> CSVSink: self.open() @@ -92,9 +97,6 @@ class CSVSink: def open(self) -> None: """ Open the CSV file for writing. - - Returns: - None """ parent_directory = os.path.dirname(self.file_name) if parent_directory and not os.path.exists(parent_directory): @@ -106,9 +108,6 @@ class CSVSink: def close(self) -> None: """ Close the CSV file. - - Returns: - None """ if self.file: self.file.close() @@ -137,10 +136,12 @@ class CSVSink: if hasattr(detections, "data"): for key, value in detections.data.items(): - if value.ndim == 0: + if isinstance(value, np.ndarray) and value.ndim == 0: row[key] = value - else: + elif isinstance(value, np.ndarray): row[key] = value[i] + else: + row[key] = value[i] if hasattr(value, "__getitem__") else value if custom_data: row.update(custom_data) @@ -154,11 +155,8 @@ class CSVSink: Append detection data to the CSV file. Args: - detections (Detections): The detection data. - custom_data (Dict[str, Any]): Custom data to include. - - Returns: - None + detections: The detection data. + custom_data: Custom data to include. """ if not self.writer: raise Exception( @@ -185,9 +183,10 @@ class CSVSink: @staticmethod def parse_field_names( - detections: Detections, custom_data: dict[str, Any] + detections: Detections, custom_data: dict[str, Any] | None = None ) -> list[str]: + custom_keys = set(custom_data.keys()) if custom_data else set() dynamic_header = sorted( - set(custom_data.keys()) | set(getattr(detections, "data", {}).keys()) + custom_keys | set(getattr(detections, "data", {}).keys()) ) return BASE_HEADER + dynamic_header diff --git a/src/supervision/detection/tools/inference_slicer.py b/src/supervision/detection/tools/inference_slicer.py index 4cc05f19..4e0fcbf8 100644 --- a/src/supervision/detection/tools/inference_slicer.py +++ b/src/supervision/detection/tools/inference_slicer.py @@ -3,8 +3,10 @@ from __future__ import annotations import warnings from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Any import numpy as np +import numpy.typing as npt from supervision.config import ORIENTED_BOX_COORDINATES from supervision.detection.core import Detections @@ -18,19 +20,19 @@ from supervision.utils.internal import SupervisionWarnings def move_detections( detections: Detections, - offset: np.ndarray, + offset: npt.NDArray[Any], resolution_wh: tuple[int, int] | None = None, ) -> Detections: """ Args: - detections (sv.Detections): Detections object to be moved. - offset (np.ndarray): An array of shape `(2,)` containing offset values in format - is `[dx, dy]`. - resolution_wh (Tuple[int, int]): The width and height of the desired mask + detections: Detections object to be moved. + offset: An array of shape `(2,)` containing offset values in the + format `[dx, dy]`. + resolution_wh: The width and height of the desired mask resolution. Required for segmentation detections. Returns: - (sv.Detections) repositioned Detections object. + Repositioned Detections object. """ detections.xyxy = move_boxes(xyxy=detections.xyxy, offset=offset) if ORIENTED_BOX_COORDINATES in detections.data: @@ -61,18 +63,17 @@ class InferenceSlicer: parallel slice inference. Args: - callback (Callable[[ImageType], Detections]): Inference function that takes - a sliced image and returns a `Detections` object. - slice_wh (int or tuple[int, int]): Size of each slice `(width, height)`. - If int, both width and height are set to this value. - overlap_wh (int or tuple[int, int]): Overlap size `(width, height)` between - slices. If int, both width and height are set to this value. - overlap_filter (OverlapFilter or str): Strategy to merge overlapping - detections (`NON_MAX_SUPPRESSION`, `NON_MAX_MERGE`, or `NONE`). - iou_threshold (float): IOU threshold used in merging overlap filtering. - overlap_metric (OverlapMetric or str): Metric to compute overlap - (`IOU` or `IOS`). - thread_workers (int): Number of threads for concurrent slice inference. + callback: Inference function that takes a sliced image and returns a + `Detections` object. + slice_wh: Size of each slice `(width, height)`. If int, both width and + height are set to this value. + overlap_wh: Overlap size `(width, height)` between slices. If int, both + width and height are set to this value. + overlap_filter: Strategy to merge overlapping detections + (`NON_MAX_SUPPRESSION`, `NON_MAX_MERGE`, or `NONE`). + iou_threshold: IOU threshold used in merging overlap filtering. + overlap_metric: Metric to compute overlap (`IOU` or `IOS`). + thread_workers: Number of threads for concurrent slice inference. Raises: ValueError: If `slice_wh` or `overlap_wh` are invalid or inconsistent. @@ -132,7 +133,7 @@ class InferenceSlicer: self.iou_threshold = iou_threshold self.overlap_metric = OverlapMetric.from_value(overlap_metric) self.overlap_filter = OverlapFilter.from_value(overlap_filter) - self.callback = callback + self.callback: Callable[[ImageType], Detections] = callback self.thread_workers = thread_workers def __call__(self, image: ImageType) -> Detections: @@ -140,10 +141,10 @@ class InferenceSlicer: Perform tiled inference on the full image and return merged detections. Args: - image (ImageType): The full image to run inference on. + image: The full image to run inference on. Returns: - Detections: Merged detections across all slices. + Merged detections across all slices. """ detections_list: list[Detections] = [] resolution_wh = get_image_resolution_wh(image) @@ -181,19 +182,19 @@ class InferenceSlicer: ) return merged - def _run_callback(self, image: ImageType, offset: np.ndarray) -> Detections: + def _run_callback(self, image: ImageType, offset: npt.NDArray[Any]) -> Detections: """ Run detection callback on a sliced portion of the image and adjust coordinates. Args: - image (ImageType): The full image. - offset (numpy.ndarray): Coordinates `(x_min, y_min, x_max, y_max)` defining + image: The full image. + offset: Coordinates `(x_min, y_min, x_max, y_max)` defining the slice region. Returns: - Detections: Detections adjusted to the full image coordinate system. + Detections adjusted to the full image coordinate system. """ - image_slice: ImageType = crop_image(image=image, xyxy=offset) + image_slice = crop_image(image=image, xyxy=offset) detections = self.callback(image_slice) resolution_wh = get_image_resolution_wh(image) @@ -260,17 +261,17 @@ class InferenceSlicer: resolution_wh: tuple[int, int], slice_wh: tuple[int, int], overlap_wh: tuple[int, int], - ) -> np.ndarray: + ) -> npt.NDArray[Any]: """ Generate bounding boxes defining the coordinates of image slices with overlap. Args: - resolution_wh (tuple[int, int]): Image resolution `(width, height)`. - slice_wh (tuple[int, int]): Size of each slice `(width, height)`. - overlap_wh (tuple[int, int]): Overlap size between slices `(width, height)`. + resolution_wh: Image resolution `(width, height)`. + slice_wh: Size of each slice `(width, height)`. + overlap_wh: Overlap size between slices `(width, height)`. Returns: - numpy.ndarray: Array of shape `(num_slices, 4)` with each row as + Array of shape `(num_slices, 4)` with each row as `(x_min, y_min, x_max, y_max)` coordinates for a slice. """ slice_width, slice_height = slice_wh @@ -289,10 +290,10 @@ class InferenceSlicer: return [0] if stride == slice_size: - return np.arange(0, image_size, stride).tolist() + return list(np.arange(0, image_size, stride).tolist()) last_start = image_size - slice_size - starts = np.arange(0, last_start, stride).tolist() + starts: list[int] = list(np.arange(0, last_start, stride).tolist()) if not starts or starts[-1] != last_start: starts.append(last_start) return starts @@ -312,7 +313,7 @@ class InferenceSlicer: x_max = np.clip(x_min + slice_width, 0, image_width) y_max = np.clip(y_min + slice_height, 0, image_height) - offsets = np.stack( + offsets: npt.NDArray[Any] = np.stack( [x_min, y_min, x_max, y_max], axis=-1, ).reshape(-1, 4) diff --git a/src/supervision/detection/tools/json_sink.py b/src/supervision/detection/tools/json_sink.py index a19eb33e..5a34e16e 100644 --- a/src/supervision/detection/tools/json_sink.py +++ b/src/supervision/detection/tools/json_sink.py @@ -1,9 +1,12 @@ from __future__ import annotations +import io import json import os from typing import Any +import numpy as np + from supervision.detection.core import Detections @@ -16,11 +19,11 @@ class JSONSink: !!! tip - JSONsink allow to pass custom data alongside the detection fields, providing + JSONSink allows passing custom data alongside detection fields, providing flexibility for logging various types of information. Args: - file_name (str): The name of the JSON file where the detections will be stored. + file_name: The name of the JSON file where the detections will be stored. Defaults to 'output.json'. Example: @@ -45,13 +48,10 @@ class JSONSink: Initialize the JSONSink instance. Args: - file_name (str): The name of the JSON file. - - Returns: - None + file_name: The name of the JSON file. """ self.file_name = file_name - self.file: open | None = None + self.file: io.TextIOWrapper | None = None self.data: list[dict[str, Any]] = [] def __enter__(self) -> JSONSink: @@ -69,9 +69,6 @@ class JSONSink: def open(self) -> None: """ Open the JSON file for writing. - - Returns: - None """ parent_directory = os.path.dirname(self.file_name) if parent_directory and not os.path.exists(parent_directory): @@ -82,9 +79,6 @@ class JSONSink: def write_and_close(self) -> None: """ Write and close the JSON file. - - Returns: - None """ if self.file: json.dump(self.data, self.file, indent=4) @@ -114,11 +108,14 @@ class JSONSink: if hasattr(detections, "data"): for key, value in detections.data.items(): - row[key] = ( - str(value[i]) - if hasattr(value, "__getitem__") and value.ndim != 0 - else str(value) - ) + if isinstance(value, np.ndarray): + row[key] = str(value[i]) if value.ndim != 0 else str(value) + else: + row[key] = ( + str(value[i]) + if hasattr(value, "__getitem__") + else str(value) + ) if custom_data: row.update(custom_data) @@ -132,11 +129,8 @@ class JSONSink: Append detection data to the JSON file. Args: - detections (Detections): The detection data. - custom_data (Dict[str, Any]): Custom data to include. - - Returns: - None + detections: The detection data. + custom_data: Custom data to include. """ parsed_rows = JSONSink.parse_detection_data(detections, custom_data) self.data.extend(parsed_rows) diff --git a/src/supervision/detection/tools/polygon_zone.py b/src/supervision/detection/tools/polygon_zone.py index 3319c178..336169dd 100644 --- a/src/supervision/detection/tools/polygon_zone.py +++ b/src/supervision/detection/tools/polygon_zone.py @@ -2,6 +2,7 @@ from __future__ import annotations from collections.abc import Iterable from dataclasses import replace +from typing import Any, cast import cv2 import numpy as np @@ -27,14 +28,14 @@ class PolygonZone: tracking into your inference pipeline. Attributes: - polygon (np.ndarray): A polygon represented by a numpy array of shape + polygon: A polygon represented by a numpy array of shape `(N, 2)`, containing the `x`, `y` coordinates of the points. - triggering_anchors (Iterable[sv.Position]): A list of positions specifying + triggering_anchors: A list of positions specifying which anchors of the detections bounding box to consider when deciding on whether the detection fits within the PolygonZone (default: (sv.Position.BOTTOM_CENTER,)). - current_count (int): The current count of detected objects within the zone - mask (np.ndarray): The 2D bool mask for the polygon zone + current_count: The current count of detected objects within the zone + mask: The 2D bool mask for the polygon zone Example: ```python @@ -81,12 +82,12 @@ class PolygonZone: """ Determines if the detections are within the polygon zone. - Parameters: - detections (Detections): The detections + Args: + detections: The detections to be checked against the polygon zone Returns: - np.ndarray: A boolean numpy array indicating + A boolean numpy array indicating if each detection is within the polygon zone """ @@ -107,28 +108,28 @@ class PolygonZone: .astype(bool) ) - is_in_zone: npt.NDArray[np.bool_] = np.all(is_in_zone, axis=1) + is_in_zone = np.all(is_in_zone, axis=1) self.current_count = int(np.sum(is_in_zone)) return is_in_zone.astype(bool) class PolygonZoneAnnotator: """ - A class for annotating a polygon-shaped zone within a - frame with a count of detected objects. + A class for annotating a polygon-shaped zone within a frame with a count of + detected objects. Attributes: - zone (PolygonZone): The polygon zone to be annotated - color (Color): The color to draw the polygon lines, default is white - thickness (int): The thickness of the polygon lines, default is 2 - text_color (Color): The color of the text on the polygon, default is black - text_scale (float): The scale of the text on the polygon, default is 0.5 - text_thickness (int): The thickness of the text on the polygon, default is 1 - text_padding (int): The padding around the text on the polygon, default is 10 - font (int): The font type for the text on the polygon, + zone: The polygon zone to be annotated + color: The color to draw the polygon lines, default is white + thickness: The thickness of the polygon lines, default is 2 + text_color: The color of the text on the polygon, default is black + text_scale: The scale of the text on the polygon, default is 0.5 + text_thickness: The thickness of the text on the polygon, default is 1 + text_padding: The padding around the text on the polygon, default is 10 + font: The font type for the text on the polygon, default is cv2.FONT_HERSHEY_SIMPLEX - center (Tuple[int, int]): The center of the polygon for text placement - display_in_zone_count (bool): Show the label of the zone or not. Default is True + center: The center of the polygon for text placement + display_in_zone_count: Show the label of the zone or not. Default is True opacity: The opacity of zone filling when drawn on the scene. Default is 0 """ @@ -156,17 +157,19 @@ class PolygonZoneAnnotator: self.display_in_zone_count = display_in_zone_count self.opacity = opacity - def annotate(self, scene: np.ndarray, label: str | None = None) -> np.ndarray: + def annotate( + self, scene: npt.NDArray[Any], label: str | None = None + ) -> npt.NDArray[Any]: """ Annotates the polygon zone within a frame with a count of detected objects. - Parameters: - scene (np.ndarray): The image on which the polygon zone will be annotated - label (Optional[str]): A label for the count of detected objects + Args: + scene: The image on which the polygon zone will be annotated + label: A label for the count of detected objects within the polygon zone (default: None) Returns: - np.ndarray: The image with the polygon zone and count of detected objects + The image with the polygon zone and count of detected objects """ if self.opacity == 0: annotated_frame = draw_polygon( @@ -202,4 +205,4 @@ class PolygonZoneAnnotator: text_font=self.font, ) - return annotated_frame + return cast(npt.NDArray[Any], annotated_frame) diff --git a/src/supervision/detection/tools/transformers.py b/src/supervision/detection/tools/transformers.py index 0358464e..f535217a 100644 --- a/src/supervision/detection/tools/transformers.py +++ b/src/supervision/detection/tools/transformers.py @@ -1,9 +1,10 @@ from __future__ import annotations import io -from typing import Any +from typing import Any, cast import numpy as np +import numpy.typing as npt from PIL import Image from supervision.config import CLASS_NAME_DATA_FIELD @@ -11,22 +12,22 @@ from supervision.detection.utils.converters import mask_to_xyxy def process_transformers_detection_result( - detection_result: dict, id2label: dict[int, str] | None -) -> dict: + detection_result: dict[str, Any], id2label: dict[int, str] | None +) -> dict[str, Any]: """ Process the result of Transformers object detection functions such as `post_process` (v4) and `post_process_detection` (v5). Args: - detection_result (dict): Dictionary containing detection results with keys + detection_result: Dictionary containing detection results with keys 'boxes', 'labels', and 'scores'. - id2label (Optional[Dict[int, str]]): A dictionary mapping class IDs to labels, + 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: - dict: Processed detection result including bounding boxes, confidence scores, - class IDs, and data. + Processed detection result including bounding boxes, confidence scores, + class IDs, and data. """ class_ids = detection_result["labels"].cpu().detach().numpy().astype(int) data = append_class_names_to_data(class_ids, id2label, {}) @@ -40,23 +41,23 @@ def process_transformers_detection_result( def process_transformers_v4_segmentation_result( - segmentation_result: dict, id2label: dict[int, str] | None -) -> dict: + segmentation_result: dict[str, Any], id2label: dict[int, str] | None +) -> dict[str, Any]: """ Process the result of Transformers segmentation functions such as `post_process_panoptic`, `post_process_segmentation`, and `post_process_instance` (v4). Args: - segmentation_result (dict): Dictionary containing segmentation results with keys + segmentation_result: Dictionary containing segmentation results with keys 'masks', 'labels', and 'scores'. - id2label (Optional[Dict[int, str]]): A dictionary mapping class IDs to labels, + 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: - dict: Processed segmentation result including bounding boxes, masks, confidence - scores, class IDs, and data. + Processed segmentation result including bounding boxes, masks, confidence + scores, class IDs, and data. """ if "png_string" in segmentation_result: return process_transformers_v4_panoptic_segmentation_result( @@ -79,52 +80,52 @@ def process_transformers_v4_segmentation_result( def process_transformers_v5_segmentation_result( - segmentation_result: dict, id2label: dict[int, str] | None -) -> dict: + segmentation_result: Any, id2label: dict[int, str] | None +) -> dict[str, Any]: """ Process the result of Transformers segmentation functions such as `post_process_semantic_segmentation`, `post_process_instance_segmentation`, and `post_process_panoptic_segmentation` (v5). Args: - segmentation_result (Union[dict, np.ndarray]): Either a dictionary containing - segmentation results or an ndarray representing a segmentation map. - id2label (Optional[Dict[int, str]]): A dictionary mapping class IDs to labels, + segmentation_result: Either a dictionary containing segmentation results + (`segments_info` and `segmentation`) or a tensor object + representing a panoptic segmentation map. + 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: - dict: Processed segmentation result including bounding boxes, masks, confidence - scores, class IDs, and data. + Processed segmentation result including bounding boxes, masks, confidence + scores, class IDs, and data. """ if segmentation_result.__class__.__name__ == "Tensor": segmentation_array = segmentation_result.cpu().detach().numpy() return process_transformers_v5_panoptic_segmentation_result( segmentation_array, id2label ) - return process_transformers_v5_semantic_or_instance_segmentation_result( - segmentation_result, id2label + cast(dict[str, Any], segmentation_result), id2label ) def process_transformers_v5_semantic_or_instance_segmentation_result( - segmentation_result: dict, id2label: dict[int, str] | None -) -> dict: + segmentation_result: dict[str, Any], id2label: dict[int, str] | None +) -> dict[str, Any]: """ Process the result of Transformers segmentation functions such as `post_process_semantic_segmentation` and `post_process_instance_segmentation` (v5). Args: - segmentation_result (dict): Dictionary containing segmentation results with keys + segmentation_result: Dictionary containing segmentation results with keys `segments_info` and `segmentation`. - id2label (Optional[Dict[int, str]]): A dictionary mapping class IDs to labels, + 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: - dict: Processed segmentation result including bounding boxes, masks, confidence - scores, class IDs, and data. + Processed segmentation result including bounding boxes, masks, confidence + scores, class IDs, and data. """ segments_info = segmentation_result["segments_info"] scores = np.array([segment["score"] for segment in segments_info]) @@ -145,21 +146,21 @@ def process_transformers_v5_semantic_or_instance_segmentation_result( def process_transformers_v4_panoptic_segmentation_result( - segmentation_result: dict, id2label: dict[int, str] | None -) -> dict: + segmentation_result: dict[str, Any], id2label: dict[int, str] | None +) -> dict[str, Any]: """ Process the result of the Transformers function `post_process_panoptic` (v4). Args: - segmentation_result (dict): Dictionary containing segmentation results with keys + segmentation_result: Dictionary containing segmentation results with keys such as 'png_string' and 'segments_info'. - id2label (Optional[Dict[int, str]]): A dictionary mapping class IDs to labels, + 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: - dict: Processed segmentation result including bounding boxes, masks, - class IDs, and data. + Processed segmentation result including bounding boxes, masks, + class IDs, and data. """ segments_info = segmentation_result["segments_info"] png_string = segmentation_result["png_string"] @@ -179,21 +180,21 @@ def process_transformers_v4_panoptic_segmentation_result( def process_transformers_v5_panoptic_segmentation_result( - segmentation_array: np.ndarray, id2label: dict[int, str] | None -) -> dict: + segmentation_array: npt.NDArray[Any], id2label: dict[int, str] | None +) -> dict[str, Any]: """ Process the result of the Transformers function `post_process_panoptic_segmentation` (v5). Args: - segmentation_array (np.ndarray): Segmentation array. - id2label (Optional[Dict[int, str]]): A dictionary mapping class IDs to labels, + segmentation_array: Segmentation array. + 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: - dict: Processed segmentation result including bounding boxes, masks, - class IDs, and data. + Processed segmentation result including bounding boxes, masks, + class IDs, and data. """ class_ids = np.unique(segmentation_array) masks = np.stack( @@ -203,25 +204,25 @@ def process_transformers_v5_panoptic_segmentation_result( return dict(xyxy=mask_to_xyxy(masks), mask=masks, class_id=class_ids, data=data) -def png_string_to_segmentation_array(png_string: bytes) -> np.ndarray: +def png_string_to_segmentation_array(png_string: bytes) -> npt.NDArray[Any]: """ Convert a PNG byte string to a label mask array. Args: - png_string (bytes): A byte string representing the PNG image. + png_string: A byte string representing the PNG image. Returns: - np.ndarray: A label mask array with shape (H, W), where H and W - are the height and width of the image. Each unique value in the array - represents a different object or category. + A label mask array with shape (H, W), where H and W + are the height and width of the image. Each unique value in the array + represents a different object or category. """ image = Image.open(io.BytesIO(png_string)) mask = np.array(image, dtype=np.uint8) - return mask[:, :, 0] + return cast(npt.NDArray[Any], mask[:, :, 0]) def append_class_names_to_data( - class_ids: np.ndarray, + class_ids: npt.NDArray[Any], id2label: dict[int, str] | None, data: dict[str, Any] | None = None, ) -> dict[str, Any]: @@ -230,14 +231,14 @@ def append_class_names_to_data( available. Args: - class_ids (np.ndarray): Array of class IDs. - id2label (Optional[Dict[int, str]]): A dictionary mapping class IDs to labels, + class_ids: Array of 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. - data (Optional[Dict[str, Any]]): An existing data dictionary to append to. + data: An existing data dictionary to append to. Returns: - Dict[str, Any]: Dictionary containing class names if id2label is provided. + Dictionary containing class names if id2label is provided. """ if data is None: data = {}