feat: add show_progress to dataset load/save operations

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
@@ -334,6 +335,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 +359,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 +399,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 +410,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 +431,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 +444,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 +459,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 +477,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 +498,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 +512,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 +526,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 +555,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 +570,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 +581,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 +593,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 +613,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 +624,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 +637,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 +681,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 +715,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 +728,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
 

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,
@@ -352,6 +353,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`.
@@ -365,6 +367,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)`.
@@ -398,7 +401,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"],
@@ -466,6 +473,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.
 
@@ -484,6 +492,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
@@ -539,7 +548,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 = {

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

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(

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]