Skip to content

feat: add show_progress option to dataset loaders and savers #2181

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
satishkc7 wants to merge 4 commits into
base: develop
Choose a base branch
from
Open

feat: add show_progress option to dataset loaders and savers #2181

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
c9201d0
feat: add show_progress option to dataset loaders and savers
satishkc7 Mar 24, 2026
069659c
feat: propagate show_progress to annotation savers and add tests
satishkc7 Mar 26, 2026
e1b3220
fix(pre_commit): 🎨 auto format pre-commit hooks
pre-commit-ci[bot] Mar 26, 2026
a9c024b
fix: close Image.open with context manager and clean up ValueError me…
satishkc7 Mar 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
69 changes: 49 additions & 20 deletions src/supervision/dataset/core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -357,39 +359,48 @@ 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 while saving images.
"""
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:
annotation_name = Path(image_path).stem
annotations_path = os.path.join(
annotations_directory_path, f"{annotation_name}.xml"
)
image_name = Path(image_path).name
pascal_voc_xml = detections_to_pascal_voc(
detections=annotations,
classes=self.classes,
filename=image_name,
image_shape=image.shape,
min_image_area_percentage=min_image_area_percentage,
max_image_area_percentage=max_image_area_percentage,
approximation_percentage=approximation_percentage,
)

with open(annotations_path, "w") as f:
f.write(pascal_voc_xml)
with tqdm(
total=len(self),
desc="Saving Pascal VOC annotations",
disable=not show_progress,
) as progress_bar:
for image_path, image, annotations in self:
annotation_name = Path(image_path).stem
annotations_path = os.path.join(
annotations_directory_path, f"{annotation_name}.xml"
)
image_name = Path(image_path).name
pascal_voc_xml = detections_to_pascal_voc(
detections=annotations,
classes=self.classes,
filename=image_name,
image_shape=image.shape,
min_image_area_percentage=min_image_area_percentage,
max_image_area_percentage=max_image_area_percentage,
approximation_percentage=approximation_percentage,
)

with open(annotations_path, "w") as f:
f.write(pascal_voc_xml)
progress_bar.update(1)

@classmethod
def from_pascal_voc(
cls,
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.
Expand All @@ -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 while loading images.

Returns:
A DetectionDataset instance containing
Expand Down Expand Up @@ -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(
Expand All @@ -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.
Expand All @@ -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 while loading images.

Returns:
A DetectionDataset instance
Expand Down Expand Up @@ -496,6 +511,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
Expand All @@ -509,6 +525,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
Expand Down Expand Up @@ -537,10 +554,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 while saving images.
"""
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(
Expand All @@ -549,6 +569,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)
Expand All @@ -559,6 +580,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.
Expand All @@ -570,6 +592,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 while loading images.
Returns:
A DetectionDataset instance containing
the loaded images and annotations.
Expand Down Expand Up @@ -599,6 +622,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)

Expand All @@ -609,6 +633,7 @@ def as_coco(
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 COCO format. This method saves the
Expand Down Expand Up @@ -645,10 +670,13 @@ def as_coco(
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.
show_progress: If `True`, display a progress bar while saving images.
"""
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:
save_coco_annotations(
Expand All @@ -657,6 +685,7 @@ def as_coco(
min_image_area_percentage=min_image_area_percentage,
max_image_area_percentage=max_image_area_percentage,
approximation_percentage=approximation_percentage,
show_progress=show_progress,
)


Expand Down
119 changes: 73 additions & 46 deletions src/supervision/dataset/formats/coco.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

import numpy as np
import numpy.typing as npt
from tqdm.auto import tqdm

from supervision.dataset.utils import (
approximate_mask_with_polygons,
Expand Down Expand Up @@ -254,6 +255,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`.
Expand All @@ -267,9 +269,21 @@ 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 while loading images.

Returns:
A tuple of `(classes, image_paths, annotations)`.

Examples:
```python
import supervision as sv

ds = sv.DetectionDataset.from_coco(
images_directory_path="images/train",
annotations_path="images/train/_annotations.coco.json",
show_progress=True,
)
```
"""
coco_data = read_json_file(file_path=annotations_path)
classes = coco_categories_to_classes(coco_categories=coco_data["categories"])
Expand All @@ -286,32 +300,38 @@ def load_coco_annotations(
images = []
annotations = {}

for coco_image in coco_images:
image_name, image_width, image_height = (
coco_image["file_name"],
coco_image["width"],
coco_image["height"],
)
image_annotations = coco_annotations_groups.get(coco_image["id"], [])
image_path = os.path.join(images_directory_path, image_name)
with tqdm(
total=len(coco_images),
desc="Loading COCO annotations",
disable=not show_progress,
) as progress_bar:
for coco_image in coco_images:
image_name, image_width, image_height = (
coco_image["file_name"],
coco_image["width"],
coco_image["height"],
)
image_annotations = coco_annotations_groups.get(coco_image["id"], [])
image_path = os.path.join(images_directory_path, image_name)

with_masks = force_masks or any(
_with_seg_mask(annotation) for annotation in image_annotations
)
annotation = coco_annotations_to_detections(
image_annotations=image_annotations,
resolution_wh=(image_width, image_height),
with_masks=with_masks,
use_iscrowd=use_iscrowd,
)
with_masks = force_masks or any(
_with_seg_mask(annotation) for annotation in image_annotations
)
annotation = coco_annotations_to_detections(
image_annotations=image_annotations,
resolution_wh=(image_width, image_height),
with_masks=with_masks,
use_iscrowd=use_iscrowd,
)

annotation = map_detections_class_id(
source_to_target_mapping=class_index_mapping,
detections=annotation,
)
annotation = map_detections_class_id(
source_to_target_mapping=class_index_mapping,
detections=annotation,
)

images.append(image_path)
annotations[image_path] = annotation
images.append(image_path)
annotations[image_path] = annotation
progress_bar.update(1)

return classes, images, annotations

Expand All @@ -326,6 +346,7 @@ def save_coco_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(annotation_path).parent.mkdir(parents=True, exist_ok=True)
licenses = [
Expand All @@ -341,30 +362,36 @@ def save_coco_annotations(
coco_categories = classes_to_coco_categories(classes=dataset.classes)

image_id, annotation_id = 1, 1
for image_path, image, annotation in dataset:
image_height, image_width, _ = image.shape
image_name = f"{Path(image_path).stem}{Path(image_path).suffix}"
coco_image = {
"id": image_id,
"license": 1,
"file_name": image_name,
"height": image_height,
"width": image_width,
"date_captured": datetime.now().strftime("%m/%d/%Y,%H:%M:%S"),
}

coco_images.append(coco_image)
coco_annotation, annotation_id = detections_to_coco_annotations(
detections=annotation,
image_id=image_id,
annotation_id=annotation_id,
min_image_area_percentage=min_image_area_percentage,
max_image_area_percentage=max_image_area_percentage,
approximation_percentage=approximation_percentage,
)
with tqdm(
total=len(dataset),
desc="Saving COCO annotations",
disable=not show_progress,
) as progress_bar:
for image_path, image, annotation in dataset:
image_height, image_width, _ = image.shape
image_name = f"{Path(image_path).stem}{Path(image_path).suffix}"
coco_image = {
"id": image_id,
"license": 1,
"file_name": image_name,
"height": image_height,
"width": image_width,
"date_captured": datetime.now().strftime("%m/%d/%Y,%H:%M:%S"),
}

coco_images.append(coco_image)
coco_annotation, annotation_id = detections_to_coco_annotations(
detections=annotation,
image_id=image_id,
annotation_id=annotation_id,
min_image_area_percentage=min_image_area_percentage,
max_image_area_percentage=max_image_area_percentage,
approximation_percentage=approximation_percentage,
)

coco_annotations.extend(coco_annotation)
image_id += 1
coco_annotations.extend(coco_annotation)
image_id += 1
progress_bar.update(1)

annotation_dict = {
"info": {},
Expand Down
Loading
Loading