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

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>
This commit is contained in:
Omkar Kabde 2026-03-10 23:21:27 +05:30 committed by GitHub
parent 664bf3cf79
commit e4cedb6861
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 157 additions and 164 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
pyproject.toml
View File

@ -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",
]

49
src/supervision/detection/tools/csv_sink.py
View File

@ -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

69
src/supervision/detection/tools/inference_slicer.py
View File

@ -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)

40
src/supervision/detection/tools/json_sink.py
View File

@ -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)

55
src/supervision/detection/tools/polygon_zone.py
View File

@ -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)

103
src/supervision/detection/tools/transformers.py
View File

@ -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 = {}