diff --git a/supervision/dataset/core.py b/supervision/dataset/core.py index 00248fcd..358b4e59 100644 --- a/supervision/dataset/core.py +++ b/supervision/dataset/core.py @@ -2,13 +2,13 @@ from __future__ import annotations from dataclasses import dataclass from pathlib import Path -from typing import Dict, List, Optional +from typing import Dict, List, Optional, Tuple import cv2 import numpy as np from supervision.file import list_files_with_extensions -from supervision.dataset.formats.pascal_voc import detections_to_pascal_voc +from supervision.dataset.formats.pascal_voc import detections_to_pascal_voc, load_pascal_voc_annotations from supervision.detection.core import Detections @@ -82,10 +82,55 @@ class Dataset: @classmethod def from_pascal_voc(cls, images_directory_path: str, annotations_directory_path: str) -> Dataset: + """ + Creates a Dataset instance from PASCAL VOC formatted data. + + Args: + images_directory_path (str): The path to the directory containing the images. + annotations_directory_path (str): The path to the directory containing the PASCAL VOC XML annotations. + + Returns: + Dataset: A Dataset instance containing the loaded images and annotations. + """ image_paths = list_files_with_extensions( directory=images_directory_path, extensions=['jpg', 'jpeg', 'png']) annotation_paths = list_files_with_extensions( directory=annotations_directory_path, extensions=['xml']) - return image_paths, annotation_paths + + raw_annotations: List[Tuple[str, Detections, List[str]]] = [ + load_pascal_voc_annotations(annotation_path=str(annotation_path)) + for annotation_path + in annotation_paths + ] + + classes = [] + for annotation in raw_annotations: + classes.extend(annotation[2]) + classes = list(set(classes)) + + for annotation in raw_annotations: + class_id = [ + classes.index(class_name) + for class_name + in annotation[2] + ] + annotation[1].class_id = np.array(class_id) + + images = { + image_path.name: cv2.imread(str(image_path)) + for image_path + in image_paths + } + + annotations = { + image_name: detections + for image_name, detections, _ + in raw_annotations + } + return Dataset( + classes=classes, + images=images, + annotations=annotations + ) diff --git a/supervision/dataset/formats/pascal_voc.py b/supervision/dataset/formats/pascal_voc.py index f453078e..ac3d3562 100644 --- a/supervision/dataset/formats/pascal_voc.py +++ b/supervision/dataset/formats/pascal_voc.py @@ -134,17 +134,17 @@ def detections_to_pascal_voc( return xml_string -def load_pascal_voc_annotations(xml_path: str) -> Tuple[str, Detections, List[str]]: +def load_pascal_voc_annotations(annotation_path: str) -> Tuple[str, Detections, List[str]]: """ Loads PASCAL VOC XML annotations and returns the image name, a Detections instance, and a list of class names. Args: - xml_path (str): The path to the PASCAL VOC XML annotations file. + annotation_path (str): The path to the PASCAL VOC XML annotations file. Returns: Tuple[str, Detections, List[str]]: A tuple containing the image name, a Detections instance, and a list of class names of objects in the detections. """ - tree = parse(xml_path) + tree = parse(annotation_path) root = tree.getroot() image_name = root.find("filename").text