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

typing(dataset): resolve type annotations and sanitize docstrings (#2146) 

* resolve type annotations and sanitize docstrings
* improve error messages in `_extract_class_names` for better debugging clarity
This commit is contained in:
Jirka Borovec 2026-02-10 22:23:32 +09:00 committed by GitHub
parent 83dfc87ef5
commit 2dd1bc5df6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 210 additions and 171 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
pyproject.toml
View File

@ -205,11 +205,6 @@ module = [
"examples.*",
# TODO: fix type errors in the following modules
"supervision.classification.core",
"supervision.dataset.core",
"supervision.dataset.formats.coco",
"supervision.dataset.formats.pascal_voc",
"supervision.dataset.formats.yolo",
"supervision.dataset.utils",
"supervision.detection.core",
"supervision.detection.line_zone",
"supervision.detection.tools.csv_sink",

161
src/supervision/dataset/core.py
View File

@ -9,6 +9,7 @@ from pathlib import Path
import cv2
import numpy as np
import numpy.typing as npt
from supervision.classification.core import Classifications
from supervision.dataset.formats.coco import (
@ -58,12 +59,12 @@ class DetectionDataset(BaseDataset):
formats.
Attributes:
classes (List[str]): List containing dataset class names.
images (Union[List[str], Dict[str, np.ndarray]]):
classes: List containing dataset class names.
images:
Accepts a list of image paths, or dictionaries of loaded cv2 images
with paths as keys. If you pass a list of paths, the dataset will
lazily load images on demand, which is much more memory-efficient.
annotations (Dict[str, Detections]): Dictionary mapping
annotations: Dictionary mapping
image path to annotations. The dictionary keys match
match the keys in `images` or entries in the list of
image paths.
@ -72,7 +73,7 @@ class DetectionDataset(BaseDataset):
def __init__(
self,
classes: list[str],
images: list[str] | dict[str, np.ndarray],
images: list[str] | dict[str, npt.NDArray[np.uint8]],
annotations: dict[str, Detections],
) -> None:
self.classes = classes
@ -86,21 +87,24 @@ class DetectionDataset(BaseDataset):
# Eliminate duplicates while preserving order
self.image_paths = list(dict.fromkeys(images))
self._images_in_memory: dict[str, np.ndarray] = {}
self._images_in_memory: dict[str, npt.NDArray[np.uint8]] = {}
def _get_image(self, image_path: str) -> np.ndarray:
"""Assumes that image is in dataset"""
def _get_image(self, image_path: str) -> npt.NDArray[np.uint8]:
"""Assumes that image is in dataset."""
if self._images_in_memory:
return self._images_in_memory[image_path]
return cv2.imread(image_path)
image = cv2.imread(image_path)
if image is None:
raise ValueError(f"Could not read image from path: {image_path}")
return image
def __len__(self) -> int:
return len(self._images_in_memory) or len(self.image_paths)
def __getitem__(self, i: int) -> tuple[str, np.ndarray, Detections]:
def __getitem__(self, i: int) -> tuple[str, npt.NDArray[np.uint8], Detections]:
"""
Returns:
Tuple[str, np.ndarray, Detections]: The image path, image data,
The image path, image data,
and its corresponding annotation at index i.
"""
image_path = self.image_paths[i]
@ -108,20 +112,18 @@ class DetectionDataset(BaseDataset):
annotation = self.annotations[image_path]
return image_path, image, annotation
def __iter__(self) -> Iterator[tuple[str, np.ndarray, Detections]]:
def __iter__(self) -> Iterator[tuple[str, npt.NDArray[np.uint8], Detections]]:
"""
Iterate over the images and annotations in the dataset.
Yields:
Iterator[Tuple[str, np.ndarray, Detections]]:
An iterator that yields tuples containing the image path,
the image data, and its corresponding annotation.
Tuples containing the image path, image data, and its annotation.
"""
for i in range(len(self)):
image_path, image, annotation = self[i]
yield image_path, image, annotation
def __eq__(self, other) -> bool:
def __eq__(self, other: object) -> bool:
if not isinstance(other, DetectionDataset):
return False
@ -154,14 +156,14 @@ class DetectionDataset(BaseDataset):
using the provided split_ratio.
Args:
split_ratio (float): The ratio of the training
split_ratio: The ratio of the training
set to the entire dataset.
random_state (Optional[int]): The seed for the random number generator.
random_state: The seed for the random number generator.
This is used for reproducibility.
shuffle (bool): Whether to shuffle the data before splitting.
shuffle: Whether to shuffle the data before splitting.
Returns:
Tuple[DetectionDataset, DetectionDataset]: A tuple containing
A tuple containing
the training and testing datasets.
Examples:
@ -193,8 +195,8 @@ class DetectionDataset(BaseDataset):
shuffle=shuffle,
)
train_input: list[str] | dict[str, np.ndarray]
test_input: list[str] | dict[str, np.ndarray]
train_input: list[str] | dict[str, npt.NDArray[np.uint8]]
test_input: list[str] | dict[str, npt.NDArray[np.uint8]]
if self._images_in_memory:
train_input = {path: self._images_in_memory[path] for path in train_paths}
test_input = {path: self._images_in_memory[path] for path in test_paths}
@ -227,11 +229,11 @@ class DetectionDataset(BaseDataset):
`annotations`) into a single `DetectionDataset` object.
Args:
dataset_list (List[DetectionDataset]): A list of `DetectionDataset`
dataset_list: A list of `DetectionDataset`
objects to merge.
Returns:
(DetectionDataset): A single `DetectionDataset` object containing
A single `DetectionDataset` object containing
the merged data from the input list.
Examples:
@ -329,21 +331,21 @@ class DetectionDataset(BaseDataset):
and their corresponding annotations in PASCAL VOC format.
Args:
images_directory_path (Optional[str]): The path to the directory
images_directory_path: The path to the directory
where the images should be saved.
If not provided, images will not be saved.
annotations_directory_path (Optional[str]): The path to
annotations_directory_path: The path to
the directory where the annotations in PASCAL VOC format should be
saved. If not provided, annotations will not be saved.
min_image_area_percentage (float): The minimum percentage of
min_image_area_percentage: The minimum percentage of
detection area relative to
the image area for a detection to be included.
Argument is used only for segmentation datasets.
max_image_area_percentage (float): The maximum percentage
max_image_area_percentage: The maximum percentage
of detection area relative to
the image area for a detection to be included.
Argument is used only for segmentation datasets.
approximation_percentage (float): The percentage of
approximation_percentage: The percentage of
polygon points to be removed from the input polygon,
in the range [0, 1). Argument is used only for segmentation datasets.
"""
@ -364,7 +366,7 @@ class DetectionDataset(BaseDataset):
detections=annotations,
classes=self.classes,
filename=image_name,
image_shape=image.shape, # type: ignore
image_shape=image.shape,
min_image_area_percentage=min_image_area_percentage,
max_image_area_percentage=max_image_area_percentage,
approximation_percentage=approximation_percentage,
@ -384,14 +386,14 @@ class DetectionDataset(BaseDataset):
Creates a Dataset instance from PASCAL VOC formatted data.
Args:
images_directory_path (str): Path to the directory containing the images.
annotations_directory_path (str): Path to the directory
images_directory_path: Path to the directory containing the images.
annotations_directory_path: Path to the directory
containing the PASCAL VOC XML annotations.
force_masks (bool): If True, forces masks to
force_masks: If True, forces masks to
be loaded for all annotations, regardless of whether they are present.
Returns:
DetectionDataset: A DetectionDataset instance containing
A DetectionDataset instance containing
the loaded images and annotations.
Examples:
@ -440,21 +442,21 @@ class DetectionDataset(BaseDataset):
Creates a Dataset instance from YOLO formatted data.
Args:
images_directory_path (str): The path to the
images_directory_path: The path to the
directory containing the images.
annotations_directory_path (str): The path to the directory
annotations_directory_path: The path to the directory
containing the YOLO annotation files.
data_yaml_path (str): The path to the data
data_yaml_path: The path to the data
YAML file containing class information.
force_masks (bool): If True, forces
force_masks: If True, forces
masks to be loaded for all annotations,
regardless of whether they are present.
is_obb (bool): If True, loads the annotations in OBB format.
is_obb: If True, loads the annotations in OBB format.
OBB annotations are defined as `[class_id, x, y, x, y, x, y, x, y]`,
where pairs of [x, y] are box corners.
Returns:
DetectionDataset: A DetectionDataset instance
A DetectionDataset instance
containing the loaded images and annotations.
Examples:
@ -504,25 +506,25 @@ class DetectionDataset(BaseDataset):
images and their corresponding annotations in YOLO format.
Args:
images_directory_path (Optional[str]): The path to the
images_directory_path: The path to the
directory where the images should be saved.
If not provided, images will not be saved.
annotations_directory_path (Optional[str]): The path to the
annotations_directory_path: The path to the
directory where the annotations in
YOLO format should be saved. If not provided,
annotations will not be saved.
data_yaml_path (Optional[str]): The path where the data.yaml
data_yaml_path: The path where the data.yaml
file should be saved.
If not provided, the file will not be saved.
min_image_area_percentage (float): The minimum percentage of
min_image_area_percentage: The minimum percentage of
detection area relative to
the image area for a detection to be included.
Argument is used only for segmentation datasets.
max_image_area_percentage (float): The maximum percentage
max_image_area_percentage: The maximum percentage
of detection area relative to
the image area for a detection to be included.
Argument is used only for segmentation datasets.
approximation_percentage (float): The percentage of polygon points to
approximation_percentage: The percentage of polygon points to
be removed from the input polygon, in the range [0, 1).
This is useful for simplifying the annotations.
Argument is used only for segmentation datasets.
@ -553,14 +555,14 @@ class DetectionDataset(BaseDataset):
Creates a Dataset instance from COCO formatted data.
Args:
images_directory_path (str): The path to the
images_directory_path: The path to the
directory containing the images.
annotations_path (str): The path to the json annotation files.
force_masks (bool): If True,
annotations_path: The path to the json annotation files.
force_masks: If True,
forces masks to be loaded for all annotations,
regardless of whether they are present.
Returns:
DetectionDataset: A DetectionDataset instance containing
A DetectionDataset instance containing
the loaded images and annotations.
Examples:
@ -618,19 +620,19 @@ class DetectionDataset(BaseDataset):
standards.
Args:
images_directory_path (Optional[str]): The path to the directory
images_directory_path: The path to the directory
where the images should be saved.
If not provided, images will not be saved.
annotations_path (Optional[str]): The path to COCO annotation file.
min_image_area_percentage (float): The minimum percentage of
annotations_path: The path to COCO annotation file.
min_image_area_percentage: The minimum percentage of
detection area relative to
the image area for a detection to be included.
Argument is used only for segmentation datasets.
max_image_area_percentage (float): The maximum percentage of
max_image_area_percentage: The maximum percentage of
detection area relative to
the image area for a detection to be included.
Argument is used only for segmentation datasets.
approximation_percentage (float): The percentage of polygon points
approximation_percentage: The percentage of polygon points
to be removed from the input polygon,
in the range [0, 1). This is useful for simplifying the annotations.
Argument is used only for segmentation datasets.
@ -656,17 +658,17 @@ class ClassificationDataset(BaseDataset):
loading, dataset splitting.
Attributes:
classes (List[str]): List containing dataset class names.
images (Union[List[str], Dict[str, np.ndarray]]):
classes: List containing dataset class names.
images:
List of image paths or dictionary mapping image name to image data.
annotations (Dict[str, Classifications]): Dictionary mapping
annotations: Dictionary mapping
image name to annotations.
"""
def __init__(
self,
classes: list[str],
images: list[str] | dict[str, np.ndarray],
images: list[str] | dict[str, npt.NDArray[np.uint8]],
annotations: dict[str, Classifications],
) -> None:
self.classes = classes
@ -680,7 +682,7 @@ class ClassificationDataset(BaseDataset):
# Eliminate duplicates while preserving order
self.image_paths = list(dict.fromkeys(images))
self._images_in_memory: dict[str, np.ndarray] = {}
self._images_in_memory: dict[str, npt.NDArray[np.uint8]] = {}
if isinstance(images, dict):
self._images_in_memory = images
warn_deprecated(
@ -689,19 +691,22 @@ class ClassificationDataset(BaseDataset):
"a list of paths `List[str]` instead."
)
def _get_image(self, image_path: str) -> np.ndarray:
"""Assumes that image is in dataset"""
def _get_image(self, image_path: str) -> npt.NDArray[np.uint8]:
"""Assumes that image is in dataset."""
if self._images_in_memory:
return self._images_in_memory[image_path]
return cv2.imread(image_path)
image = cv2.imread(image_path)
if image is None:
raise ValueError(f"Could not read image from path: {image_path}")
return image
def __len__(self) -> int:
return len(self._images_in_memory) or len(self.image_paths)
def __getitem__(self, i: int) -> tuple[str, np.ndarray, Classifications]:
def __getitem__(self, i: int) -> tuple[str, npt.NDArray[np.uint8], Classifications]:
"""
Returns:
Tuple[str, np.ndarray, Classifications]: The image path, image data,
The image path, image data,
and its corresponding annotation at index i.
"""
image_path = self.image_paths[i]
@ -709,20 +714,20 @@ class ClassificationDataset(BaseDataset):
annotation = self.annotations[image_path]
return image_path, image, annotation
def __iter__(self) -> Iterator[tuple[str, np.ndarray, Classifications]]:
def __iter__(
self,
) -> Iterator[tuple[str, npt.NDArray[np.uint8], Classifications]]:
"""
Iterate over the images and annotations in the dataset.
Yields:
Iterator[Tuple[str, np.ndarray, Detections]]:
An iterator that yields tuples containing the image path,
the image data, and its corresponding annotation.
Tuples containing the image path, image data, and its annotation.
"""
for i in range(len(self)):
image_path, image, annotation = self[i]
yield image_path, image, annotation
def __eq__(self, other) -> bool:
def __eq__(self, other: object) -> bool:
if not isinstance(other, ClassificationDataset):
return False
@ -755,14 +760,14 @@ class ClassificationDataset(BaseDataset):
using the provided split_ratio.
Args:
split_ratio (float): The ratio of the training
split_ratio: The ratio of the training
set to the entire dataset.
random_state (Optional[int]): The seed for the
random_state: The seed for the
random number generator. This is used for reproducibility.
shuffle (bool): Whether to shuffle the data before splitting.
shuffle: Whether to shuffle the data before splitting.
Returns:
Tuple[ClassificationDataset, ClassificationDataset]: A tuple containing
A tuple containing
the training and testing datasets.
Examples:
@ -793,8 +798,8 @@ class ClassificationDataset(BaseDataset):
shuffle=shuffle,
)
train_input: list[str] | dict[str, np.ndarray]
test_input: list[str] | dict[str, np.ndarray]
train_input: list[str] | dict[str, npt.NDArray[np.uint8]]
test_input: list[str] | dict[str, npt.NDArray[np.uint8]]
if self._images_in_memory:
train_input = {path: self._images_in_memory[path] for path in train_paths}
test_input = {path: self._images_in_memory[path] for path in test_paths}
@ -822,7 +827,7 @@ class ClassificationDataset(BaseDataset):
Saves the dataset as a multi-class folder structure.
Args:
root_directory_path (str): The path to the directory
root_directory_path: The path to the directory
where the dataset will be saved.
"""
os.makedirs(root_directory_path, exist_ok=True)
@ -847,10 +852,10 @@ class ClassificationDataset(BaseDataset):
Load data from a multiclass folder structure into a ClassificationDataset.
Args:
root_directory_path (str): The path to the dataset directory.
root_directory_path: The path to the dataset directory.
Returns:
ClassificationDataset: The dataset.
The dataset.
Examples:
```python

53
src/supervision/dataset/formats/coco.py
View File

@ -1,7 +1,7 @@
import os
from datetime import datetime
from pathlib import Path
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Any, Union
import numpy as np
import numpy.typing as npt
@ -20,8 +20,10 @@ from supervision.utils.file import read_json_file, save_json_file
if TYPE_CHECKING:
from supervision.dataset.core import DetectionDataset
CocoDict = dict[str, Any]
def coco_categories_to_classes(coco_categories: list[dict]) -> list[str]:
def coco_categories_to_classes(coco_categories: list[CocoDict]) -> list[str]:
return [
category["name"]
for category in sorted(coco_categories, key=lambda category: category["id"])
@ -29,7 +31,7 @@ def coco_categories_to_classes(coco_categories: list[dict]) -> list[str]:
def build_coco_class_index_mapping(
coco_categories: list[dict], target_classes: list[str]
coco_categories: list[CocoDict], target_classes: list[str]
) -> dict[int, int]:
source_class_to_index = {
category["name"]: category["id"] for category in coco_categories
@ -40,7 +42,7 @@ def build_coco_class_index_mapping(
}
def classes_to_coco_categories(classes: list[str]) -> list[dict]:
def classes_to_coco_categories(classes: list[str]) -> list[CocoDict]:
return [
{
"id": class_id,
@ -52,9 +54,9 @@ def classes_to_coco_categories(classes: list[str]) -> list[dict]:
def group_coco_annotations_by_image_id(
coco_annotations: list[dict],
) -> dict[int, list[dict]]:
annotations = {}
coco_annotations: list[CocoDict],
) -> dict[int, list[CocoDict]]:
annotations: dict[int, list[CocoDict]] = {}
for annotation in coco_annotations:
image_id = annotation["image_id"]
if image_id not in annotations:
@ -64,7 +66,7 @@ def group_coco_annotations_by_image_id(
def coco_annotations_to_masks(
image_annotations: list[dict], resolution_wh: tuple[int, int]
image_annotations: list[CocoDict], resolution_wh: tuple[int, int]
) -> npt.NDArray[np.bool_]:
return np.array(
[
@ -87,7 +89,7 @@ def coco_annotations_to_masks(
def coco_annotations_to_detections(
image_annotations: list[dict],
image_annotations: list[CocoDict],
resolution_wh: tuple[int, int],
with_masks: bool,
use_iscrowd: bool = True,
@ -99,10 +101,10 @@ def coco_annotations_to_detections(
image_annotation["category_id"] for image_annotation in image_annotations
]
xyxy = [image_annotation["bbox"] for image_annotation in image_annotations]
xyxy = np.asarray(xyxy)
xyxy = np.asarray(xyxy, dtype=np.float32)
xyxy[:, 2:4] += xyxy[:, 0:2]
data = dict()
data: dict[str, npt.NDArray[np.generic]] = {}
if use_iscrowd:
iscrowd = [
image_annotation["iscrowd"] for image_annotation in image_annotations
@ -131,11 +133,13 @@ def detections_to_coco_annotations(
min_image_area_percentage: float = 0.0,
max_image_area_percentage: float = 1.0,
approximation_percentage: float = 0.75,
) -> tuple[list[dict], int]:
coco_annotations = []
) -> tuple[list[CocoDict], int]:
coco_annotations: list[CocoDict] = []
for xyxy, mask, _, class_id, _, _ in detections:
if class_id is None:
raise ValueError("Detections must include class_id for COCO export.")
box_width, box_height = xyxy[2] - xyxy[0], xyxy[3] - xyxy[1]
segmentation = []
segmentation: Union[list[list[float]], dict[str, list[int]]] = []
iscrowd = 0
if mask is not None:
iscrowd = contains_holes(mask=mask) or contains_multiple_segments(mask=mask)
@ -146,16 +150,13 @@ def detections_to_coco_annotations(
"size": list(mask.shape[:2]),
}
else:
segmentation = [
list(
approximate_mask_with_polygons(
mask=mask,
min_image_area_percentage=min_image_area_percentage,
max_image_area_percentage=max_image_area_percentage,
approximation_percentage=approximation_percentage,
)[0].flatten()
)
]
polygons = approximate_mask_with_polygons(
mask=mask,
min_image_area_percentage=min_image_area_percentage,
max_image_area_percentage=max_image_area_percentage,
approximation_percentage=approximation_percentage,
)
segmentation = [list(polygons[0].flatten())]
coco_annotation = {
"id": annotation_id,
"image_id": image_id,
@ -194,11 +195,11 @@ def get_coco_class_index_mapping(annotations_path: str) -> dict[int, int]:
- Returns a dictionary mapping: `{new_class_id: original_COCO_class_id}`.
Args:
annotations_path (str): Path to COCO JSON annotations file
annotations_path: Path to COCO JSON annotations file
(e.g., `instances_val2017.json`).
Returns:
Dict[int, int]: A mapping from new class id (sequential ranging from 0 to 79)
A mapping from new class id (sequential ranging from 0 to 79)
to original COCO class id (1 to 90 with skipped ids).
"""
coco_data = read_json_file(annotations_path)

92
src/supervision/dataset/formats/pascal_voc.py
View File

@ -6,6 +6,7 @@ from xml.etree.ElementTree import Element, SubElement
import cv2
import numpy as np
import numpy.typing as npt
from defusedxml.ElementTree import parse, tostring
from defusedxml.minidom import parseString
@ -16,7 +17,9 @@ from supervision.utils.file import list_files_with_extensions
def object_to_pascal_voc(
xyxy: np.ndarray, name: str, polygon: np.ndarray | None = None
xyxy: npt.NDArray[np.number],
name: str,
polygon: npt.NDArray[np.number] | None = None,
) -> Element:
root = Element("object")
@ -63,21 +66,21 @@ def detections_to_pascal_voc(
Converts Detections object to Pascal VOC XML format.
Args:
detections (Detections): A Detections object containing bounding boxes,
detections: A Detections object containing bounding boxes,
class ids, and other relevant information.
classes (List[str]): A list of class names corresponding to the
classes: A list of class names corresponding to the
class ids in the Detections object.
filename (str): The name of the image file associated with the detections.
image_shape (Tuple[int, int, int]): The shape of the image
filename: The name of the image file associated with the detections.
image_shape: The shape of the image
file associated with the detections.
min_image_area_percentage (float): Minimum detection area
min_image_area_percentage: Minimum detection area
relative to area of image associated with it.
max_image_area_percentage (float): Maximum detection area
max_image_area_percentage: Maximum detection area
relative to area of image associated with it.
approximation_percentage (float): The percentage of
approximation_percentage: The percentage of
polygon points to be removed from the input polygon, in the range [0, 1).
Returns:
str: An XML string in Pascal VOC format representing the detections.
An XML string in Pascal VOC format representing the detections.
"""
height, width, depth = image_shape
@ -112,6 +115,8 @@ def detections_to_pascal_voc(
# Add object elements
for xyxy, mask, _, class_id, _, _ in detections:
if class_id is None:
raise ValueError("Detections must include class_id for Pascal VOC export.")
name = classes[class_id]
if mask is not None:
polygons = approximate_mask_with_polygons(
@ -131,7 +136,7 @@ def detections_to_pascal_voc(
annotation.append(next_object)
# Generate XML string
xml_string = parseString(tostring(annotation)).toprettyxml(indent=" ")
xml_string = str(parseString(tostring(annotation)).toprettyxml(indent=" "))
return xml_string
@ -145,14 +150,14 @@ def load_pascal_voc_annotations(
a Detections instance, and a list of class names.
Args:
images_directory_path (str): The path to the directory containing the images.
annotations_directory_path (str): The path to the directory containing the
images_directory_path: The path to the directory containing the images.
annotations_directory_path: The path to the directory containing the
PASCAL VOC annotation files.
force_masks (bool): If True, forces masks to be loaded for all
force_masks: If True, forces masks to be loaded for all
annotations, regardless of whether they are present.
Returns:
Tuple[List[str], List[str], Dict[str, Detections]]: A tuple with a list
A tuple with a list
of class names, a list of paths to images, and a dictionary with image
paths as keys and corresponding Detections instances as values.
"""
@ -178,6 +183,8 @@ def load_pascal_voc_annotations(
root = tree.getroot()
image = cv2.imread(image_path)
if image is None:
raise ValueError(f"Could not read image from path: {image_path}")
resolution_wh = (image.shape[1], image.shape[0])
annotation, classes = detections_from_xml_obj(
root, classes, resolution_wh, force_masks
@ -188,7 +195,10 @@ def load_pascal_voc_annotations(
def detections_from_xml_obj(
root: Element, classes: list[str], resolution_wh, force_masks: bool = False
root: Element,
classes: list[str],
resolution_wh: tuple[int, int],
force_masks: bool = False,
) -> tuple[Detections, list[str]]:
"""
Converts an XML object in Pascal VOC format to a Detections object.
@ -217,32 +227,34 @@ def detections_from_xml_obj(
</annotation>
Returns:
Tuple[Detections, List[str]]: A tuple containing a Detections object and an
A tuple containing a Detections object and an
updated list of class names, extended with the class names
from the XML object.
"""
xyxy = []
class_names = []
masks = []
xyxy: list[list[int]] = []
class_names: list[str] = []
masks: list[npt.NDArray[np.bool_]] = []
with_masks = False
extended_classes = classes[:]
for obj in root.findall("object"):
class_name = obj.find("name").text
class_name = _get_required_text(obj, "name")
class_names.append(class_name)
bbox = obj.find("bndbox")
x1 = int(bbox.find("xmin").text)
y1 = int(bbox.find("ymin").text)
x2 = int(bbox.find("xmax").text)
y2 = int(bbox.find("ymax").text)
if bbox is None:
raise ValueError("Missing bndbox in Pascal VOC annotation.")
x1 = int(_get_required_text(bbox, "xmin"))
y1 = int(_get_required_text(bbox, "ymin"))
x2 = int(_get_required_text(bbox, "xmax"))
y2 = int(_get_required_text(bbox, "ymax"))
xyxy.append([x1, y1, x2, y2])
with_masks = obj.find("polygon") is not None
with_masks = force_masks if force_masks else with_masks
for polygon in obj.findall("polygon"):
polygon = parse_polygon_points(polygon)
for polygon_element in obj.findall("polygon"):
polygon = parse_polygon_points(polygon_element)
# https://github.com/roboflow/supervision/issues/144
polygon -= 1
@ -252,10 +264,14 @@ def detections_from_xml_obj(
)
masks.append(mask_from_polygon)
xyxy = np.array(xyxy) if len(xyxy) > 0 else np.empty((0, 4))
xyxy_arr: npt.NDArray[np.float32]
if xyxy:
xyxy_arr = np.array(xyxy, dtype=np.float32)
else:
xyxy_arr = np.empty((0, 4), dtype=np.float32)
# https://github.com/roboflow/supervision/issues/144
xyxy -= 1
xyxy_arr -= 1
for k in set(class_names):
if k not in extended_classes:
@ -265,7 +281,7 @@ def detections_from_xml_obj(
)
annotation = Detections(
xyxy=xyxy.astype(np.float32),
xyxy=xyxy_arr,
mask=np.array(masks).astype(bool) if with_masks else None,
class_id=class_id,
)
@ -273,8 +289,20 @@ def detections_from_xml_obj(
return annotation, extended_classes
def parse_polygon_points(polygon: Element) -> np.ndarray:
coordinates = [int(coord.text) for coord in polygon.findall(".//*")]
def parse_polygon_points(polygon: Element) -> npt.NDArray[np.int_]:
coordinates: list[int] = []
for coord in polygon.findall(".//*"):
if coord.text is None:
raise ValueError("Missing polygon coordinate value in Pascal VOC.")
coordinates.append(int(coord.text))
return np.array(
[(coordinates[i], coordinates[i + 1]) for i in range(0, len(coordinates), 2)]
[(coordinates[i], coordinates[i + 1]) for i in range(0, len(coordinates), 2)],
dtype=int,
)
def _get_required_text(element: Element, tag: str) -> str:
child = element.find(tag)
if child is None or child.text is None:
raise ValueError(f"Missing '{tag}' in Pascal VOC annotation.")
return child.text

50
src/supervision/dataset/formats/yolo.py
View File

@ -2,9 +2,10 @@ from __future__ import annotations
import os
from pathlib import Path
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Any
import numpy as np
import numpy.typing as npt
from PIL import Image
from supervision.config import ORIENTED_BOX_COORDINATES
@ -23,7 +24,7 @@ if TYPE_CHECKING:
from supervision.dataset.core import DetectionDataset
def _parse_box(values: list[str]) -> np.ndarray:
def _parse_box(values: list[str]) -> npt.NDArray[np.float32]:
x_center, y_center, width, height = values
return np.array(
[
@ -36,19 +37,19 @@ def _parse_box(values: list[str]) -> np.ndarray:
)
def _box_to_polygon(box: np.ndarray) -> np.ndarray:
def _box_to_polygon(box: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
return np.array(
[[box[0], box[1]], [box[2], box[1]], [box[2], box[3]], [box[0], box[3]]]
)
def _parse_polygon(values: list[str]) -> np.ndarray:
def _parse_polygon(values: list[str]) -> npt.NDArray[np.float32]:
return np.array(values, dtype=np.float32).reshape(-1, 2)
def _polygons_to_masks(
polygons: list[np.ndarray], resolution_wh: tuple[int, int]
) -> np.ndarray:
polygons: list[npt.NDArray[np.number]], resolution_wh: tuple[int, int]
) -> npt.NDArray[np.bool_]:
return np.array(
[
polygon_to_mask(polygon=polygon, resolution_wh=resolution_wh)
@ -63,11 +64,21 @@ def _with_mask(lines: list[str]) -> bool:
def _extract_class_names(file_path: str) -> list[str]:
data = read_yaml_file(file_path=file_path)
names = data["names"]
data: dict[str, Any] = read_yaml_file(file_path=file_path)
if not isinstance(data, dict):
raise ValueError(
f"Expected mapping in data.yaml at '{file_path}',"
f" got {type(data).__name__}."
)
names = data.get("names")
if isinstance(names, dict):
names = [names[key] for key in sorted(names.keys())]
return names
return [str(names[key]) for key in sorted(names.keys())]
if isinstance(names, list):
return [str(name) for name in names]
raise ValueError(
"Expected 'names' to be a list or dict in data.yaml at "
f"'{file_path}', got {type(names).__name__}."
)
def _image_name_to_annotation_name(image_name: str) -> str:
@ -136,20 +147,19 @@ def load_yolo_annotations(
and their corresponding detections.
Args:
images_directory_path (str): The path to the directory containing the images.
annotations_directory_path (str): The path to the directory
images_directory_path: The path to the directory containing the images.
annotations_directory_path: The path to the directory
containing the YOLO annotation files.
data_yaml_path (str): The path to the data
data_yaml_path: The path to the data
YAML file containing class information.
force_masks (bool): If True, forces masks to be loaded
force_masks: If True, forces masks to be loaded
for all annotations, regardless of whether they are present.
is_obb (bool): If True, loads the annotations in OBB format.
is_obb: If True, loads the annotations in OBB format.
OBB annotations are defined as `[class_id, x, y, x, y, x, y, x, y]`,
where pairs of [x, y] are box corners.
Returns:
Tuple[List[str], List[str], Dict[str, Detections]]:
A tuple containing a list of class names, a dictionary with
A tuple containing a list of class names, a dictionary with
image names as keys and images as values, and a dictionary
with image names as keys and corresponding Detections instances as values.
"""
@ -205,10 +215,10 @@ def load_yolo_annotations(
def object_to_yolo(
xyxy: np.ndarray,
xyxy: npt.NDArray[np.number],
class_id: int,
image_shape: tuple[int, int, int],
polygon: np.ndarray | None = None,
polygon: npt.NDArray[np.number] | None = None,
) -> str:
h, w, _ = image_shape
if polygon is None:
@ -278,7 +288,7 @@ def save_yolo_annotations(
)
lines = detections_to_yolo_annotations(
detections=annotation,
image_shape=image.shape, # type: ignore
image_shape=image.shape,
min_image_area_percentage=min_image_area_percentage,
max_image_area_percentage=max_image_area_percentage,
approximation_percentage=approximation_percentage,

20
src/supervision/dataset/utils.py
View File

@ -25,11 +25,11 @@ T = TypeVar("T")
def approximate_mask_with_polygons(
mask: np.ndarray,
mask: npt.NDArray[np.bool_],
min_image_area_percentage: float = 0.0,
max_image_area_percentage: float = 1.0,
approximation_percentage: float = 0.75,
) -> list[np.ndarray]:
) -> list[npt.NDArray[np.number]]:
height, width = mask.shape
image_area = height * width
minimum_detection_area = min_image_area_percentage * image_area
@ -121,13 +121,13 @@ def train_test_split(
Splits the data into two parts using the provided train_ratio.
Args:
data (List[T]): The data to split.
train_ratio (float): The ratio of the training set to the entire dataset.
random_state (Optional[int]): The seed for the random number generator.
shuffle (bool): Whether to shuffle the data before splitting.
data: The data to split.
train_ratio: The ratio of the training set to the entire dataset.
random_state: The seed for the random number generator.
shuffle: Whether to shuffle the data before splitting.
Returns:
Tuple[List[T], List[T]]: The split data.
The split data.
"""
if random_state is not None:
random.seed(random_state)
@ -146,12 +146,12 @@ def rle_to_mask(
Converts run-length encoding (RLE) to a binary mask.
Args:
rle (Union[npt.NDArray[np.int_], List[int]]): The 1D RLE array, the format
rle: The 1D RLE array, the format
used in the COCO dataset (column-wise encoding, values of an array with
even indices represent the number of pixels assigned as background,
values of an array with odd indices represent the number of pixels
assigned as foreground object).
resolution_wh (Tuple[int, int]): The width (w) and height (h)
resolution_wh: The width (w) and height (h)
of the desired binary mask.
Returns:
@ -200,7 +200,7 @@ def mask_to_rle(mask: npt.NDArray[np.bool_]) -> list[int]:
Converts a binary mask into a run-length encoding (RLE).
Args:
mask (npt.NDArray[np.bool_]): 2D binary mask where `True` indicates foreground
mask: 2D binary mask where `True` indicates foreground
object and `False` indicates background.
Returns: