From 2dd1bc5df6cb2ea066bd77d7ea57f384715ef48f Mon Sep 17 00:00:00 2001 From: Jirka Borovec <6035284+Borda@users.noreply.github.com> Date: Tue, 10 Feb 2026 22:23:32 +0900 Subject: [PATCH] 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 --- pyproject.toml | 5 - src/supervision/dataset/core.py | 161 +++++++++--------- src/supervision/dataset/formats/coco.py | 53 +++--- src/supervision/dataset/formats/pascal_voc.py | 92 ++++++---- src/supervision/dataset/formats/yolo.py | 50 +++--- src/supervision/dataset/utils.py | 20 +-- 6 files changed, 210 insertions(+), 171 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 55dc2681..05e33eb9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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", diff --git a/src/supervision/dataset/core.py b/src/supervision/dataset/core.py index f0b3d339..823339ab 100644 --- a/src/supervision/dataset/core.py +++ b/src/supervision/dataset/core.py @@ -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 diff --git a/src/supervision/dataset/formats/coco.py b/src/supervision/dataset/formats/coco.py index b4827f29..59603ee0 100644 --- a/src/supervision/dataset/formats/coco.py +++ b/src/supervision/dataset/formats/coco.py @@ -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) diff --git a/src/supervision/dataset/formats/pascal_voc.py b/src/supervision/dataset/formats/pascal_voc.py index c5172c33..83ec55a8 100644 --- a/src/supervision/dataset/formats/pascal_voc.py +++ b/src/supervision/dataset/formats/pascal_voc.py @@ -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( 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 diff --git a/src/supervision/dataset/formats/yolo.py b/src/supervision/dataset/formats/yolo.py index 6e8ffb04..1b8a5feb 100644 --- a/src/supervision/dataset/formats/yolo.py +++ b/src/supervision/dataset/formats/yolo.py @@ -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, diff --git a/src/supervision/dataset/utils.py b/src/supervision/dataset/utils.py index 37d9bf91..f52111a3 100644 --- a/src/supervision/dataset/utils.py +++ b/src/supervision/dataset/utils.py @@ -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: