From 63fc377ed2619e27aa6946b9026f99996283d444 Mon Sep 17 00:00:00 2001 From: murillo-ro-silva Date: Mon, 25 May 2026 17:50:57 -0300 Subject: [PATCH 1/2] feat: add progress bars for dataset load/save operations (#183) Add `show_progress` parameter to all dataset loading and saving methods. When enabled, displays a tqdm progress bar during time-consuming operations like loading/saving COCO, YOLO, and Pascal VOC datasets. - Defaults to `False` for full backward compatibility - Uses `tqdm.auto` for terminal and Jupyter notebook support - Includes 13 new tests covering all formats and backward compatibility --- src/supervision/dataset/core.py | 42 ++- src/supervision/dataset/formats/coco.py | 18 +- src/supervision/dataset/formats/pascal_voc.py | 9 +- src/supervision/dataset/formats/yolo.py | 17 +- src/supervision/dataset/utils.py | 13 +- tests/dataset/test_progress.py | 313 ++++++++++++++++++ 6 files changed, 400 insertions(+), 12 deletions(-) create mode 100644 tests/dataset/test_progress.py diff --git a/src/supervision/dataset/core.py b/src/supervision/dataset/core.py index fdadce68c6..9bee97711c 100644 --- a/src/supervision/dataset/core.py +++ b/src/supervision/dataset/core.py @@ -37,6 +37,8 @@ from supervision.utils.internal import warn_deprecated from supervision.utils.iterables import find_duplicates +from tqdm.auto import tqdm + class BaseDataset(ABC): @abstractmethod @@ -334,6 +336,7 @@ def as_pascal_voc( min_image_area_percentage: float = 0.0, max_image_area_percentage: float = 1.0, approximation_percentage: float = 0.0, + show_progress: bool = False, ) -> None: """ Exports the dataset to PASCAL VOC format. This method saves the images @@ -357,15 +360,22 @@ def as_pascal_voc( 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. + show_progress: If True, display a progress bar during saving. """ if images_directory_path: save_dataset_images( dataset=self, images_directory_path=images_directory_path, + show_progress=show_progress, ) if annotations_directory_path: Path(annotations_directory_path).mkdir(parents=True, exist_ok=True) - for image_path, image, annotations in self: + for image_path, image, annotations in tqdm( + self, + total=len(self), + desc="Saving Pascal VOC annotations", + disable=not show_progress, + ): annotation_name = Path(image_path).stem annotations_path = os.path.join( annotations_directory_path, f"{annotation_name}.xml" @@ -390,6 +400,7 @@ def from_pascal_voc( images_directory_path: str, annotations_directory_path: str, force_masks: bool = False, + show_progress: bool = False, ) -> DetectionDataset: """ Creates a Dataset instance from PASCAL VOC formatted data. @@ -400,6 +411,7 @@ def from_pascal_voc( containing the PASCAL VOC XML annotations. force_masks: If True, forces masks to be loaded for all annotations, regardless of whether they are present. + show_progress: If True, display a progress bar during loading. Returns: A DetectionDataset instance containing @@ -420,7 +432,8 @@ def from_pascal_voc( ds = sv.DetectionDataset.from_pascal_voc( images_directory_path=f"{dataset.location}/train/images", - annotations_directory_path=f"{dataset.location}/train/labels" + annotations_directory_path=f"{dataset.location}/train/labels", + show_progress=True ) ds.classes @@ -432,6 +445,7 @@ def from_pascal_voc( images_directory_path=images_directory_path, annotations_directory_path=annotations_directory_path, force_masks=force_masks, + show_progress=show_progress, ) return DetectionDataset( @@ -446,6 +460,7 @@ def from_yolo( data_yaml_path: str, force_masks: bool = False, is_obb: bool = False, + show_progress: bool = False, ) -> DetectionDataset: """ Creates a Dataset instance from YOLO formatted data. @@ -463,6 +478,7 @@ def from_yolo( 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. + show_progress: If True, display a progress bar during loading. Returns: A DetectionDataset instance @@ -483,7 +499,8 @@ def from_yolo( ds = sv.DetectionDataset.from_yolo( images_directory_path=f"{dataset.location}/train/images", annotations_directory_path=f"{dataset.location}/train/labels", - data_yaml_path=f"{dataset.location}/data.yaml" + data_yaml_path=f"{dataset.location}/data.yaml", + show_progress=True ) ds.classes @@ -496,6 +513,7 @@ def from_yolo( data_yaml_path=data_yaml_path, force_masks=force_masks, is_obb=is_obb, + show_progress=show_progress, ) return DetectionDataset( classes=classes, images=image_paths, annotations=annotations @@ -509,6 +527,7 @@ def as_yolo( min_image_area_percentage: float = 0.0, max_image_area_percentage: float = 1.0, approximation_percentage: float = 0.0, + show_progress: bool = False, ) -> None: """ Exports the dataset to YOLO format. This method saves the @@ -537,10 +556,13 @@ def as_yolo( 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. + show_progress: If True, display a progress bar during saving. """ if images_directory_path is not None: save_dataset_images( - dataset=self, images_directory_path=images_directory_path + dataset=self, + images_directory_path=images_directory_path, + show_progress=show_progress, ) if annotations_directory_path is not None: save_yolo_annotations( @@ -549,6 +571,7 @@ def as_yolo( min_image_area_percentage=min_image_area_percentage, max_image_area_percentage=max_image_area_percentage, approximation_percentage=approximation_percentage, + show_progress=show_progress, ) if data_yaml_path is not None: save_data_yaml(data_yaml_path=data_yaml_path, classes=self.classes) @@ -559,6 +582,7 @@ def from_coco( images_directory_path: str, annotations_path: str, force_masks: bool = False, + show_progress: bool = False, ) -> DetectionDataset: """ Creates a Dataset instance from COCO formatted data. @@ -570,6 +594,7 @@ def from_coco( force_masks: If True, forces masks to be loaded for all annotations, regardless of whether they are present. + show_progress: If True, display a progress bar during loading. Returns: A DetectionDataset instance containing the loaded images and annotations. @@ -589,6 +614,7 @@ def from_coco( ds = sv.DetectionDataset.from_coco( images_directory_path=f"{dataset.location}/train", annotations_path=f"{dataset.location}/train/_annotations.coco.json", + show_progress=True ) ds.classes @@ -599,6 +625,7 @@ def from_coco( images_directory_path=images_directory_path, annotations_path=annotations_path, force_masks=force_masks, + show_progress=show_progress, ) return DetectionDataset(classes=classes, images=images, annotations=annotations) @@ -611,6 +638,7 @@ def as_coco( approximation_percentage: float = 0.0, starting_image_id: int = 1, starting_annotation_id: int = 1, + show_progress: bool = False, ) -> tuple[int, int]: """ Exports the dataset to COCO format. This method saves the @@ -654,6 +682,7 @@ def as_coco( starting_annotation_id: First annotation id to assign in the exported file. Defaults to ``1``. Override for the same multi-split reason as ``starting_image_id``. + show_progress: If True, display a progress bar during saving. Returns: A ``(next_image_id, next_annotation_id)`` tuple containing the @@ -687,7 +716,9 @@ def as_coco( """ if images_directory_path is not None: save_dataset_images( - dataset=self, images_directory_path=images_directory_path + dataset=self, + images_directory_path=images_directory_path, + show_progress=show_progress, ) if annotations_path is not None: return save_coco_annotations( @@ -698,6 +729,7 @@ def as_coco( approximation_percentage=approximation_percentage, starting_image_id=starting_image_id, starting_annotation_id=starting_annotation_id, + show_progress=show_progress, ) return starting_image_id, starting_annotation_id diff --git a/src/supervision/dataset/formats/coco.py b/src/supervision/dataset/formats/coco.py index d13440bd22..5c00e48076 100644 --- a/src/supervision/dataset/formats/coco.py +++ b/src/supervision/dataset/formats/coco.py @@ -5,6 +5,7 @@ import numpy as np import numpy.typing as npt +from tqdm.auto import tqdm from supervision.dataset.utils import ( approximate_mask_with_polygons, @@ -261,6 +262,7 @@ def load_coco_annotations( annotations_path: str, force_masks: bool = False, use_iscrowd: bool = True, + show_progress: bool = False, ) -> tuple[list[str], list[str], dict[str, Detections]]: """ Load COCO annotations and convert them to `Detections`. @@ -274,6 +276,7 @@ def load_coco_annotations( annotations_path: Path to COCO JSON annotations. force_masks: If `True`, always attempt to load masks. use_iscrowd: If `True`, include `iscrowd` and `area` in detection data. + show_progress: If `True`, display a progress bar during loading. Returns: A tuple of `(classes, image_paths, annotations)`. @@ -307,7 +310,11 @@ def load_coco_annotations( annotations = {} images_directory_resolved = Path(images_directory_path).resolve() - for coco_image in coco_images: + for coco_image in tqdm( + coco_images, + desc="Loading COCO annotations", + disable=not show_progress, + ): image_name, image_width, image_height = ( coco_image["file_name"], coco_image["width"], @@ -375,6 +382,7 @@ def save_coco_annotations( approximation_percentage: float = 0.0, starting_image_id: int = 1, starting_annotation_id: int = 1, + show_progress: bool = False, ) -> tuple[int, int]: """Save a DetectionDataset to a COCO-format ``annotations.json`` file. @@ -393,6 +401,7 @@ def save_coco_annotations( starting_annotation_id: First annotation id to assign in the exported file. Defaults to ``1``. Override for the same multi-split reason as ``starting_image_id``. + show_progress: If ``True``, display a progress bar during saving. Returns: A ``(next_image_id, next_annotation_id)`` tuple. The returned values @@ -448,7 +457,12 @@ def save_coco_annotations( coco_categories = classes_to_coco_categories(classes=dataset.classes) image_id, annotation_id = starting_image_id, starting_annotation_id - for image_path, image, annotation in dataset: + for image_path, image, annotation in tqdm( + dataset, + total=len(dataset), + desc="Saving COCO annotations", + disable=not show_progress, + ): image_height, image_width, _ = image.shape image_name = f"{Path(image_path).stem}{Path(image_path).suffix}" coco_image = { diff --git a/src/supervision/dataset/formats/pascal_voc.py b/src/supervision/dataset/formats/pascal_voc.py index 91b6664cdc..fc8696254c 100644 --- a/src/supervision/dataset/formats/pascal_voc.py +++ b/src/supervision/dataset/formats/pascal_voc.py @@ -9,6 +9,7 @@ import numpy.typing as npt from defusedxml.ElementTree import parse, tostring from defusedxml.minidom import parseString +from tqdm.auto import tqdm from supervision.dataset.utils import approximate_mask_with_polygons from supervision.detection.core import Detections @@ -149,6 +150,7 @@ def load_pascal_voc_annotations( images_directory_path: str, annotations_directory_path: str, force_masks: bool = False, + show_progress: bool = False, ) -> tuple[list[str], list[str], dict[str, Detections]]: """ Loads PASCAL VOC XML annotations and returns the image name, @@ -160,6 +162,7 @@ def load_pascal_voc_annotations( PASCAL VOC annotation files. force_masks: If True, forces masks to be loaded for all annotations, regardless of whether they are present. + show_progress: If True, display a progress bar during loading. Returns: A tuple with a list @@ -177,7 +180,11 @@ def load_pascal_voc_annotations( classes: list[str] = [] annotations = {} - for image_path in image_paths: + for image_path in tqdm( + image_paths, + desc="Loading Pascal VOC annotations", + disable=not show_progress, + ): image_stem = Path(image_path).stem annotation_path = os.path.join(annotations_directory_path, f"{image_stem}.xml") if not os.path.exists(annotation_path): diff --git a/src/supervision/dataset/formats/yolo.py b/src/supervision/dataset/formats/yolo.py index 9abaff4595..0fb443edad 100644 --- a/src/supervision/dataset/formats/yolo.py +++ b/src/supervision/dataset/formats/yolo.py @@ -8,6 +8,7 @@ import numpy as np import numpy.typing as npt from PIL import Image +from tqdm.auto import tqdm from supervision.config import ORIENTED_BOX_COORDINATES from supervision.dataset.utils import approximate_mask_with_polygons @@ -145,6 +146,7 @@ def load_yolo_annotations( data_yaml_path: str, force_masks: bool = False, is_obb: bool = False, + show_progress: bool = False, ) -> tuple[list[str], list[str], dict[str, Detections]]: """ Loads YOLO annotations and returns class names, images, @@ -163,6 +165,7 @@ def load_yolo_annotations( 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. + show_progress: If True, display a progress bar during loading. Returns: A tuple containing a list of class names, a dictionary with @@ -197,7 +200,11 @@ def load_yolo_annotations( classes = _extract_class_names(file_path=data_yaml_path) annotations = {} - for image_path in image_paths: + for image_path in tqdm( + image_paths, + desc="Loading YOLO annotations", + disable=not show_progress, + ): image_stem = Path(image_path).stem annotation_path = os.path.join(annotations_directory_path, f"{image_stem}.txt") if not os.path.exists(annotation_path): @@ -296,9 +303,15 @@ def save_yolo_annotations( min_image_area_percentage: float = 0.0, max_image_area_percentage: float = 1.0, approximation_percentage: float = 0.75, + show_progress: bool = False, ) -> None: Path(annotations_directory_path).mkdir(parents=True, exist_ok=True) - for image_path, image, annotation in dataset: + for image_path, image, annotation in tqdm( + dataset, + total=len(dataset), + desc="Saving YOLO annotations", + disable=not show_progress, + ): image_name = Path(image_path).name yolo_annotations_name = _image_name_to_annotation_name(image_name=image_name) yolo_annotations_path = os.path.join( diff --git a/src/supervision/dataset/utils.py b/src/supervision/dataset/utils.py index ab8fa98b6e..e559818c05 100644 --- a/src/supervision/dataset/utils.py +++ b/src/supervision/dataset/utils.py @@ -11,6 +11,7 @@ import numpy as np import numpy.typing as npt from deprecate import deprecated, void +from tqdm.auto import tqdm from supervision.detection.core import Detections from supervision.detection.utils.converters import mask_to_polygons @@ -125,9 +126,17 @@ def map_detections_class_id( return detections_copy -def save_dataset_images(dataset: DetectionDataset, images_directory_path: str) -> None: +def save_dataset_images( + dataset: DetectionDataset, + images_directory_path: str, + show_progress: bool = False, +) -> None: Path(images_directory_path).mkdir(parents=True, exist_ok=True) - for image_path in dataset.image_paths: + for image_path in tqdm( + dataset.image_paths, + desc="Saving images", + disable=not show_progress, + ): final_path = os.path.join(images_directory_path, Path(image_path).name) if image_path in dataset._images_in_memory: image = dataset._images_in_memory[image_path] diff --git a/tests/dataset/test_progress.py b/tests/dataset/test_progress.py new file mode 100644 index 0000000000..4bb62e1690 --- /dev/null +++ b/tests/dataset/test_progress.py @@ -0,0 +1,313 @@ +"""Tests for show_progress parameter on dataset load/save operations.""" +from __future__ import annotations + +import json +import os +import tempfile +from pathlib import Path +from unittest.mock import patch + +import cv2 +import numpy as np +import pytest + +from supervision import DetectionDataset, Detections + + +def _create_dummy_yolo_dataset( + root: str, num_images: int = 3 +) -> tuple[str, str, str]: + images_dir = os.path.join(root, "images") + labels_dir = os.path.join(root, "labels") + os.makedirs(images_dir, exist_ok=True) + os.makedirs(labels_dir, exist_ok=True) + + for i in range(num_images): + img = np.zeros((100, 100, 3), dtype=np.uint8) + cv2.imwrite(os.path.join(images_dir, f"img_{i}.jpg"), img) + with open(os.path.join(labels_dir, f"img_{i}.txt"), "w") as f: + f.write(f"0 0.5 0.5 0.2 0.2\n") + + data_yaml = os.path.join(root, "data.yaml") + with open(data_yaml, "w") as f: + f.write("names:\n - class_0\nnc: 1\n") + + return images_dir, labels_dir, data_yaml + + +def _create_dummy_coco_dataset( + root: str, num_images: int = 3 +) -> tuple[str, str]: + images_dir = os.path.join(root, "images") + os.makedirs(images_dir, exist_ok=True) + + coco = { + "images": [], + "annotations": [], + "categories": [{"id": 0, "name": "class_0", "supercategory": "none"}], + } + + for i in range(num_images): + img = np.zeros((100, 100, 3), dtype=np.uint8) + fname = f"img_{i}.jpg" + cv2.imwrite(os.path.join(images_dir, fname), img) + coco["images"].append( + { + "id": i, + "file_name": fname, + "width": 100, + "height": 100, + } + ) + coco["annotations"].append( + { + "id": i, + "image_id": i, + "category_id": 0, + "bbox": [10, 10, 20, 20], + "area": 400, + "segmentation": [], + "iscrowd": 0, + } + ) + + annotations_path = os.path.join(root, "annotations.json") + with open(annotations_path, "w") as f: + json.dump(coco, f) + + return images_dir, annotations_path + + +def _create_dummy_pascal_voc_dataset( + root: str, num_images: int = 3 +) -> tuple[str, str]: + images_dir = os.path.join(root, "images") + annotations_dir = os.path.join(root, "annotations") + os.makedirs(images_dir, exist_ok=True) + os.makedirs(annotations_dir, exist_ok=True) + + for i in range(num_images): + img = np.zeros((100, 100, 3), dtype=np.uint8) + cv2.imwrite(os.path.join(images_dir, f"img_{i}.jpg"), img) + xml_content = f""" + + images + img_{i}.jpg + + 100 + 100 + 3 + + + class_0 + + 10 + 10 + 30 + 30 + + +""" + with open(os.path.join(annotations_dir, f"img_{i}.xml"), "w") as f: + f.write(xml_content) + + return images_dir, annotations_dir + + +class TestYoloProgress: + def test_from_yolo_no_progress_by_default(self): + with tempfile.TemporaryDirectory() as tmpdir: + images_dir, labels_dir, data_yaml = _create_dummy_yolo_dataset(tmpdir) + with patch("supervision.dataset.formats.yolo.tqdm", wraps=__import__("tqdm").auto.tqdm) as mock_tqdm: + ds = DetectionDataset.from_yolo( + images_directory_path=images_dir, + annotations_directory_path=labels_dir, + data_yaml_path=data_yaml, + ) + call_kwargs = mock_tqdm.call_args + assert call_kwargs[1]["disable"] is True + assert len(ds) == 3 + + def test_from_yolo_with_progress(self): + with tempfile.TemporaryDirectory() as tmpdir: + images_dir, labels_dir, data_yaml = _create_dummy_yolo_dataset(tmpdir) + with patch("supervision.dataset.formats.yolo.tqdm", wraps=__import__("tqdm").auto.tqdm) as mock_tqdm: + ds = DetectionDataset.from_yolo( + images_directory_path=images_dir, + annotations_directory_path=labels_dir, + data_yaml_path=data_yaml, + show_progress=True, + ) + call_kwargs = mock_tqdm.call_args + assert call_kwargs[1]["disable"] is False + assert len(ds) == 3 + + def test_as_yolo_with_progress(self): + with tempfile.TemporaryDirectory() as tmpdir: + images_dir, labels_dir, data_yaml = _create_dummy_yolo_dataset(tmpdir) + ds = DetectionDataset.from_yolo( + images_directory_path=images_dir, + annotations_directory_path=labels_dir, + data_yaml_path=data_yaml, + ) + + out_dir = os.path.join(tmpdir, "output") + out_images = os.path.join(out_dir, "images") + out_labels = os.path.join(out_dir, "labels") + + with patch("supervision.dataset.formats.yolo.tqdm", wraps=__import__("tqdm").auto.tqdm) as mock_tqdm: + ds.as_yolo( + images_directory_path=out_images, + annotations_directory_path=out_labels, + data_yaml_path=os.path.join(out_dir, "data.yaml"), + show_progress=True, + ) + call_kwargs = mock_tqdm.call_args + assert call_kwargs[1]["disable"] is False + + +class TestCocoProgress: + def test_from_coco_no_progress_by_default(self): + with tempfile.TemporaryDirectory() as tmpdir: + images_dir, annotations_path = _create_dummy_coco_dataset(tmpdir) + with patch("supervision.dataset.formats.coco.tqdm", wraps=__import__("tqdm").auto.tqdm) as mock_tqdm: + ds = DetectionDataset.from_coco( + images_directory_path=images_dir, + annotations_path=annotations_path, + ) + call_kwargs = mock_tqdm.call_args + assert call_kwargs[1]["disable"] is True + assert len(ds) == 3 + + def test_from_coco_with_progress(self): + with tempfile.TemporaryDirectory() as tmpdir: + images_dir, annotations_path = _create_dummy_coco_dataset(tmpdir) + with patch("supervision.dataset.formats.coco.tqdm", wraps=__import__("tqdm").auto.tqdm) as mock_tqdm: + ds = DetectionDataset.from_coco( + images_directory_path=images_dir, + annotations_path=annotations_path, + show_progress=True, + ) + call_kwargs = mock_tqdm.call_args + assert call_kwargs[1]["disable"] is False + assert len(ds) == 3 + + def test_as_coco_with_progress(self): + with tempfile.TemporaryDirectory() as tmpdir: + images_dir, annotations_path = _create_dummy_coco_dataset(tmpdir) + ds = DetectionDataset.from_coco( + images_directory_path=images_dir, + annotations_path=annotations_path, + ) + + out_dir = os.path.join(tmpdir, "output") + with patch("supervision.dataset.formats.coco.tqdm", wraps=__import__("tqdm").auto.tqdm) as mock_tqdm: + ds.as_coco( + images_directory_path=os.path.join(out_dir, "images"), + annotations_path=os.path.join(out_dir, "annotations.json"), + show_progress=True, + ) + call_kwargs = mock_tqdm.call_args + assert call_kwargs[1]["disable"] is False + + +class TestPascalVocProgress: + def test_from_pascal_voc_no_progress_by_default(self): + with tempfile.TemporaryDirectory() as tmpdir: + images_dir, annotations_dir = _create_dummy_pascal_voc_dataset(tmpdir) + with patch("supervision.dataset.formats.pascal_voc.tqdm", wraps=__import__("tqdm").auto.tqdm) as mock_tqdm: + ds = DetectionDataset.from_pascal_voc( + images_directory_path=images_dir, + annotations_directory_path=annotations_dir, + ) + call_kwargs = mock_tqdm.call_args + assert call_kwargs[1]["disable"] is True + assert len(ds) == 3 + + def test_from_pascal_voc_with_progress(self): + with tempfile.TemporaryDirectory() as tmpdir: + images_dir, annotations_dir = _create_dummy_pascal_voc_dataset(tmpdir) + with patch("supervision.dataset.formats.pascal_voc.tqdm", wraps=__import__("tqdm").auto.tqdm) as mock_tqdm: + ds = DetectionDataset.from_pascal_voc( + images_directory_path=images_dir, + annotations_directory_path=annotations_dir, + show_progress=True, + ) + call_kwargs = mock_tqdm.call_args + assert call_kwargs[1]["disable"] is False + assert len(ds) == 3 + + def test_as_pascal_voc_with_progress(self): + with tempfile.TemporaryDirectory() as tmpdir: + images_dir, annotations_dir = _create_dummy_pascal_voc_dataset(tmpdir) + ds = DetectionDataset.from_pascal_voc( + images_directory_path=images_dir, + annotations_directory_path=annotations_dir, + ) + + out_dir = os.path.join(tmpdir, "output") + with patch("supervision.dataset.core.tqdm", wraps=__import__("tqdm").auto.tqdm) as mock_tqdm: + ds.as_pascal_voc( + images_directory_path=os.path.join(out_dir, "images"), + annotations_directory_path=os.path.join(out_dir, "annotations"), + show_progress=True, + ) + call_kwargs = mock_tqdm.call_args + assert call_kwargs[1]["disable"] is False + + +class TestSaveImagesProgress: + def test_save_images_with_progress(self): + with tempfile.TemporaryDirectory() as tmpdir: + images_dir, labels_dir, data_yaml = _create_dummy_yolo_dataset(tmpdir) + ds = DetectionDataset.from_yolo( + images_directory_path=images_dir, + annotations_directory_path=labels_dir, + data_yaml_path=data_yaml, + ) + + out_images = os.path.join(tmpdir, "output_images") + with patch("supervision.dataset.utils.tqdm", wraps=__import__("tqdm").auto.tqdm) as mock_tqdm: + from supervision.dataset.utils import save_dataset_images + save_dataset_images( + dataset=ds, + images_directory_path=out_images, + show_progress=True, + ) + call_kwargs = mock_tqdm.call_args + assert call_kwargs[1]["disable"] is False + + saved_files = os.listdir(out_images) + assert len(saved_files) == 3 + + +class TestBackwardCompatibility: + """Ensure show_progress=False (default) doesn't change behavior.""" + + def test_from_yolo_default_works(self): + with tempfile.TemporaryDirectory() as tmpdir: + images_dir, labels_dir, data_yaml = _create_dummy_yolo_dataset(tmpdir) + ds = DetectionDataset.from_yolo( + images_directory_path=images_dir, + annotations_directory_path=labels_dir, + data_yaml_path=data_yaml, + ) + assert len(ds) == 3 + + def test_from_coco_default_works(self): + with tempfile.TemporaryDirectory() as tmpdir: + images_dir, annotations_path = _create_dummy_coco_dataset(tmpdir) + ds = DetectionDataset.from_coco( + images_directory_path=images_dir, + annotations_path=annotations_path, + ) + assert len(ds) == 3 + + def test_from_pascal_voc_default_works(self): + with tempfile.TemporaryDirectory() as tmpdir: + images_dir, annotations_dir = _create_dummy_pascal_voc_dataset(tmpdir) + ds = DetectionDataset.from_pascal_voc( + images_directory_path=images_dir, + annotations_directory_path=annotations_dir, + ) + assert len(ds) == 3 From 0f080fc1e6900ea2f467066487156da15804620a Mon Sep 17 00:00:00 2001 From: murillo-ro-silva Date: Mon, 25 May 2026 17:53:44 -0300 Subject: [PATCH 2/2] style: fix ruff formatting (line length, import order) --- src/supervision/dataset/core.py | 3 +- tests/dataset/test_progress.py | 68 ++++++++++++++++++++++----------- 2 files changed, 46 insertions(+), 25 deletions(-) diff --git a/src/supervision/dataset/core.py b/src/supervision/dataset/core.py index 9bee97711c..746ff1a2f2 100644 --- a/src/supervision/dataset/core.py +++ b/src/supervision/dataset/core.py @@ -10,6 +10,7 @@ import cv2 import numpy as np import numpy.typing as npt +from tqdm.auto import tqdm from supervision.classification.core import Classifications from supervision.config import CLASS_NAME_DATA_FIELD @@ -37,8 +38,6 @@ from supervision.utils.internal import warn_deprecated from supervision.utils.iterables import find_duplicates -from tqdm.auto import tqdm - class BaseDataset(ABC): @abstractmethod diff --git a/tests/dataset/test_progress.py b/tests/dataset/test_progress.py index 4bb62e1690..572d3a71f6 100644 --- a/tests/dataset/test_progress.py +++ b/tests/dataset/test_progress.py @@ -1,22 +1,19 @@ """Tests for show_progress parameter on dataset load/save operations.""" + from __future__ import annotations import json import os import tempfile -from pathlib import Path from unittest.mock import patch import cv2 import numpy as np -import pytest -from supervision import DetectionDataset, Detections +from supervision import DetectionDataset -def _create_dummy_yolo_dataset( - root: str, num_images: int = 3 -) -> tuple[str, str, str]: +def _create_dummy_yolo_dataset(root: str, num_images: int = 3) -> tuple[str, str, str]: images_dir = os.path.join(root, "images") labels_dir = os.path.join(root, "labels") os.makedirs(images_dir, exist_ok=True) @@ -26,7 +23,7 @@ def _create_dummy_yolo_dataset( img = np.zeros((100, 100, 3), dtype=np.uint8) cv2.imwrite(os.path.join(images_dir, f"img_{i}.jpg"), img) with open(os.path.join(labels_dir, f"img_{i}.txt"), "w") as f: - f.write(f"0 0.5 0.5 0.2 0.2\n") + f.write("0 0.5 0.5 0.2 0.2\n") data_yaml = os.path.join(root, "data.yaml") with open(data_yaml, "w") as f: @@ -35,9 +32,7 @@ def _create_dummy_yolo_dataset( return images_dir, labels_dir, data_yaml -def _create_dummy_coco_dataset( - root: str, num_images: int = 3 -) -> tuple[str, str]: +def _create_dummy_coco_dataset(root: str, num_images: int = 3) -> tuple[str, str]: images_dir = os.path.join(root, "images") os.makedirs(images_dir, exist_ok=True) @@ -78,9 +73,7 @@ def _create_dummy_coco_dataset( return images_dir, annotations_path -def _create_dummy_pascal_voc_dataset( - root: str, num_images: int = 3 -) -> tuple[str, str]: +def _create_dummy_pascal_voc_dataset(root: str, num_images: int = 3) -> tuple[str, str]: images_dir = os.path.join(root, "images") annotations_dir = os.path.join(root, "annotations") os.makedirs(images_dir, exist_ok=True) @@ -118,7 +111,10 @@ class TestYoloProgress: def test_from_yolo_no_progress_by_default(self): with tempfile.TemporaryDirectory() as tmpdir: images_dir, labels_dir, data_yaml = _create_dummy_yolo_dataset(tmpdir) - with patch("supervision.dataset.formats.yolo.tqdm", wraps=__import__("tqdm").auto.tqdm) as mock_tqdm: + with patch( + "supervision.dataset.formats.yolo.tqdm", + wraps=__import__("tqdm").auto.tqdm, + ) as mock_tqdm: ds = DetectionDataset.from_yolo( images_directory_path=images_dir, annotations_directory_path=labels_dir, @@ -131,7 +127,10 @@ def test_from_yolo_no_progress_by_default(self): def test_from_yolo_with_progress(self): with tempfile.TemporaryDirectory() as tmpdir: images_dir, labels_dir, data_yaml = _create_dummy_yolo_dataset(tmpdir) - with patch("supervision.dataset.formats.yolo.tqdm", wraps=__import__("tqdm").auto.tqdm) as mock_tqdm: + with patch( + "supervision.dataset.formats.yolo.tqdm", + wraps=__import__("tqdm").auto.tqdm, + ) as mock_tqdm: ds = DetectionDataset.from_yolo( images_directory_path=images_dir, annotations_directory_path=labels_dir, @@ -155,7 +154,10 @@ def test_as_yolo_with_progress(self): out_images = os.path.join(out_dir, "images") out_labels = os.path.join(out_dir, "labels") - with patch("supervision.dataset.formats.yolo.tqdm", wraps=__import__("tqdm").auto.tqdm) as mock_tqdm: + with patch( + "supervision.dataset.formats.yolo.tqdm", + wraps=__import__("tqdm").auto.tqdm, + ) as mock_tqdm: ds.as_yolo( images_directory_path=out_images, annotations_directory_path=out_labels, @@ -170,7 +172,10 @@ class TestCocoProgress: def test_from_coco_no_progress_by_default(self): with tempfile.TemporaryDirectory() as tmpdir: images_dir, annotations_path = _create_dummy_coco_dataset(tmpdir) - with patch("supervision.dataset.formats.coco.tqdm", wraps=__import__("tqdm").auto.tqdm) as mock_tqdm: + with patch( + "supervision.dataset.formats.coco.tqdm", + wraps=__import__("tqdm").auto.tqdm, + ) as mock_tqdm: ds = DetectionDataset.from_coco( images_directory_path=images_dir, annotations_path=annotations_path, @@ -182,7 +187,10 @@ def test_from_coco_no_progress_by_default(self): def test_from_coco_with_progress(self): with tempfile.TemporaryDirectory() as tmpdir: images_dir, annotations_path = _create_dummy_coco_dataset(tmpdir) - with patch("supervision.dataset.formats.coco.tqdm", wraps=__import__("tqdm").auto.tqdm) as mock_tqdm: + with patch( + "supervision.dataset.formats.coco.tqdm", + wraps=__import__("tqdm").auto.tqdm, + ) as mock_tqdm: ds = DetectionDataset.from_coco( images_directory_path=images_dir, annotations_path=annotations_path, @@ -201,7 +209,10 @@ def test_as_coco_with_progress(self): ) out_dir = os.path.join(tmpdir, "output") - with patch("supervision.dataset.formats.coco.tqdm", wraps=__import__("tqdm").auto.tqdm) as mock_tqdm: + with patch( + "supervision.dataset.formats.coco.tqdm", + wraps=__import__("tqdm").auto.tqdm, + ) as mock_tqdm: ds.as_coco( images_directory_path=os.path.join(out_dir, "images"), annotations_path=os.path.join(out_dir, "annotations.json"), @@ -215,7 +226,10 @@ class TestPascalVocProgress: def test_from_pascal_voc_no_progress_by_default(self): with tempfile.TemporaryDirectory() as tmpdir: images_dir, annotations_dir = _create_dummy_pascal_voc_dataset(tmpdir) - with patch("supervision.dataset.formats.pascal_voc.tqdm", wraps=__import__("tqdm").auto.tqdm) as mock_tqdm: + with patch( + "supervision.dataset.formats.pascal_voc.tqdm", + wraps=__import__("tqdm").auto.tqdm, + ) as mock_tqdm: ds = DetectionDataset.from_pascal_voc( images_directory_path=images_dir, annotations_directory_path=annotations_dir, @@ -227,7 +241,10 @@ def test_from_pascal_voc_no_progress_by_default(self): def test_from_pascal_voc_with_progress(self): with tempfile.TemporaryDirectory() as tmpdir: images_dir, annotations_dir = _create_dummy_pascal_voc_dataset(tmpdir) - with patch("supervision.dataset.formats.pascal_voc.tqdm", wraps=__import__("tqdm").auto.tqdm) as mock_tqdm: + with patch( + "supervision.dataset.formats.pascal_voc.tqdm", + wraps=__import__("tqdm").auto.tqdm, + ) as mock_tqdm: ds = DetectionDataset.from_pascal_voc( images_directory_path=images_dir, annotations_directory_path=annotations_dir, @@ -246,7 +263,9 @@ def test_as_pascal_voc_with_progress(self): ) out_dir = os.path.join(tmpdir, "output") - with patch("supervision.dataset.core.tqdm", wraps=__import__("tqdm").auto.tqdm) as mock_tqdm: + with patch( + "supervision.dataset.core.tqdm", wraps=__import__("tqdm").auto.tqdm + ) as mock_tqdm: ds.as_pascal_voc( images_directory_path=os.path.join(out_dir, "images"), annotations_directory_path=os.path.join(out_dir, "annotations"), @@ -267,8 +286,11 @@ def test_save_images_with_progress(self): ) out_images = os.path.join(tmpdir, "output_images") - with patch("supervision.dataset.utils.tqdm", wraps=__import__("tqdm").auto.tqdm) as mock_tqdm: + with patch( + "supervision.dataset.utils.tqdm", wraps=__import__("tqdm").auto.tqdm + ) as mock_tqdm: from supervision.dataset.utils import save_dataset_images + save_dataset_images( dataset=ds, images_directory_path=out_images,