From 9ff41b970690c04b25549cf148b0ce6269ee3cff Mon Sep 17 00:00:00 2001 From: Madhav-C Date: Thu, 25 Jun 2026 08:26:34 -0500 Subject: [PATCH] feat(dataset): add CreateML format support to DetectionDataset (#2284) - Added CreateML import/export support for detection datasets, including pixel-space center/width/height box conversion, class-name inference, global class-id consistency, and image path safety validation. - Added optional progress bars for CreateML loading, exporting, and image saving. - Improved CreateML validation with clear errors for malformed JSON, missing fields, duplicate images, null annotations, and unsafe image paths. - Updated dataset documentation and references to include CreateML workflows and `from_createml` / `as_createml` usage. --------- Co-authored-by: jirka <6035284+Borda@users.noreply.github.com> Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Co-authored-by: Codex --- docs/changelog.md | 2 + docs/datasets/core.md | 2 +- docs/how_to/process_datasets.md | 97 ++++- docs/llms.txt | 4 +- src/supervision/dataset/core.py | 102 ++++++ src/supervision/dataset/formats/createml.py | 320 ++++++++++++++++ tests/dataset/formats/test_createml.py | 382 ++++++++++++++++++++ 7 files changed, 901 insertions(+), 8 deletions(-) create mode 100644 src/supervision/dataset/formats/createml.py create mode 100644 tests/dataset/formats/test_createml.py diff --git a/docs/changelog.md b/docs/changelog.md index 31bcded1..a9a909a8 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -31,6 +31,8 @@ date_modified: 2026-06-16 - Performance [#2339](https://github.com/roboflow/supervision/pull/2339): `sv.HaloAnnotator` now uses the same CompactMask painting path as `sv.MaskAnnotator` via a shared `_paint_masks_by_area` helper. On a 1080p frame with 30 CompactMask detections, `HaloAnnotator` runs approximately 4× faster; annotated output is unchanged. +- Added [#2284](https://github.com/roboflow/supervision/pull/2284): [`DetectionDataset.from_createml`](https://supervision.roboflow.com/latest/datasets/core/#supervision.dataset.core.DetectionDataset.from_createml) and [`DetectionDataset.as_createml`](https://supervision.roboflow.com/latest/datasets/core/#supervision.dataset.core.DetectionDataset.as_createml) add load and export support for the CreateML object-detection JSON format, alongside the existing COCO, YOLO, and Pascal VOC formats. + - Performance [#2330](https://github.com/roboflow/supervision/pull/2330): `sv.mask_to_xyxy` and `sv.KeyPoints.as_detections` are now vectorized. `mask_to_xyxy` uses batched occupancy-profile reductions instead of per-mask pixel scans; `KeyPoints.as_detections` computes all bounding boxes in a single batch operation. Both produce bit-identical results. - Performance [#2323](https://github.com/roboflow/supervision/pull/2323): Mask IoU computation now uses matrix multiplication on flattened masks instead of an explicit `(N, M, H, W)` intersection tensor, reducing peak memory for large mask sets. For masks larger than 4096×4096 pixels, computation automatically promotes to float64 to preserve exact pixel counts. Results are numerically identical. diff --git a/docs/datasets/core.md b/docs/datasets/core.md index f98712c9..66216c2d 100644 --- a/docs/datasets/core.md +++ b/docs/datasets/core.md @@ -1,6 +1,6 @@ --- comments: true -description: API reference for supervision's DetectionDataset and ClassificationDataset — load, merge, split, and convert datasets in YOLO, COCO, and VOC formats. +description: API reference for supervision's DetectionDataset and ClassificationDataset — load, merge, split, and convert datasets in YOLO, COCO, VOC, and CreateML formats. --- # Datasets diff --git a/docs/how_to/process_datasets.md b/docs/how_to/process_datasets.md index cae34e86..40241136 100644 --- a/docs/how_to/process_datasets.md +++ b/docs/how_to/process_datasets.md @@ -1,24 +1,24 @@ --- comments: true -description: Load, split, merge, and convert computer vision datasets between YOLO, COCO, and Pascal VOC formats using supervision's DetectionDataset. +description: Load, split, merge, and convert computer vision datasets between YOLO, COCO, Pascal VOC, and CreateML formats using supervision's DetectionDataset. authors: - name: Piotr Skalski role: Computer Vision Engineer, Roboflow github: https://github.com/SkalskiP -date_modified: 2026-04-22 +date_modified: 2026-06-25 --- With Supervision, you can load and manipulate classification, object detection, and segmentation datasets. This tutorial will walk you through how to load, split, merge, visualize, and augment datasets in Supervision. ## Download Dataset -In this tutorial, we will use a dataset from [Roboflow Universe](https://universe.roboflow.com/), a public repository of thousands of computer vision datasets. If you already have your dataset in [COCO](https://roboflow.com/formats/coco-json), [YOLO](https://roboflow.com/formats/yolov8-pytorch-txt), or [Pascal VOC](https://roboflow.com/formats/pascal-voc-xml) format, you can skip this section. +In this tutorial, we will use a dataset from [Roboflow Universe](https://universe.roboflow.com/), a public repository of thousands of computer vision datasets. If you already have your dataset in [COCO](https://roboflow.com/formats/coco-json), [YOLO](https://roboflow.com/formats/yolov8-pytorch-txt), [Pascal VOC](https://roboflow.com/formats/pascal-voc-xml), or [CreateML](https://roboflow.com/formats/createml-json) format, you can skip this section. ```bash pip install roboflow ``` -Next, log into your Roboflow account and download the dataset of your choice in the COCO, YOLO, or Pascal VOC format. You can customize the following code snippet with your workspace ID, project ID, and version number. +Next, log into your Roboflow account and download the dataset of your choice in the COCO, YOLO, Pascal VOC, or CreateML format. You can customize the following code snippet with your workspace ID, project ID, and version number. === "COCO" @@ -56,6 +56,18 @@ Next, log into your Roboflow account and download the dataset of your choice in dataset = project.version("").download("voc") ``` +=== "CreateML" + + ```python + import roboflow + + roboflow.login() + + rf = roboflow.Roboflow() + project = rf.workspace("").project("") + dataset = project.version("").download("createml") + ``` + ## Load Dataset The Supervision library provides convenient functions to load datasets in various formats. If your dataset is already split into train, test, and valid subsets, you can load each of those as separate [`sv.DetectionDataset`](https://supervision.roboflow.com/latest/datasets/core/#supervision.dataset.core.DetectionDataset) instances. @@ -144,6 +156,33 @@ The Supervision library provides convenient functions to load datasets in variou # 800, 100, 100 ``` +=== "CreateML" + + We can do so using the [`sv.DetectionDataset.from_createml`](https://supervision.roboflow.com/latest/datasets/core/#supervision.dataset.core.DetectionDataset.from_createml) to load annotations in [CreateML](https://roboflow.com/formats/createml-json) format. + + ```python + import supervision as sv + + ds_train = sv.DetectionDataset.from_createml( + images_directory_path=f"{dataset.location}/train", + annotations_path=f"{dataset.location}/train/_annotations.createml.json", + ) + ds_valid = sv.DetectionDataset.from_createml( + images_directory_path=f"{dataset.location}/valid", + annotations_path=f"{dataset.location}/valid/_annotations.createml.json", + ) + ds_test = sv.DetectionDataset.from_createml( + images_directory_path=f"{dataset.location}/test", + annotations_path=f"{dataset.location}/test/_annotations.createml.json", + ) + + ds_train.classes + # ['person', 'bicycle', 'car', ...] + + len(ds_train), len(ds_valid), len(ds_test) + # 800, 100, 100 + ``` + ## Split Dataset If your dataset is not already split into train, test, and valid subsets, you can easily do so using the [`sv.DetectionDataset.split`](https://supervision.roboflow.com/latest/datasets/core/#supervision.dataset.core.DetectionDataset.split) method. We can split it as follows, ensuring a random shuffle of the data. @@ -269,6 +308,39 @@ If you have multiple datasets that you would like to merge, you can do so using # 1000 ``` +=== "CreateML" + + ```{ .py hl_lines="22-28" } + import supervision as sv + + ds_train = sv.DetectionDataset.from_createml( + images_directory_path=f'{dataset.location}/train', + annotations_path=f'{dataset.location}/train/_annotations.createml.json', + ) + ds_valid = sv.DetectionDataset.from_createml( + images_directory_path=f'{dataset.location}/valid', + annotations_path=f'{dataset.location}/valid/_annotations.createml.json', + ) + ds_test = sv.DetectionDataset.from_createml( + images_directory_path=f'{dataset.location}/test', + annotations_path=f'{dataset.location}/test/_annotations.createml.json', + ) + + ds_train.classes + # ['person', 'bicycle', 'car', ...] + + len(ds_train), len(ds_valid), len(ds_test) + # 800, 100, 100 + + ds = sv.DetectionDataset.merge([ds_train, ds_valid, ds_test]) + + ds.classes + # ['person', 'bicycle', 'car', ...] + + len(ds) + # 1000 + ``` + ## Iterate over Dataset There are two ways to loop over a `sv.DetectionDataset`: using a direct [for loop](https://supervision.roboflow.com/latest/datasets/core/#supervision.dataset.core.DetectionDataset.__iter__) called on the `sv.DetectionDataset` instance or loading `sv.DetectionDataset` entries [by index](https://supervision.roboflow.com/latest/datasets/core/#supervision.dataset.core.DetectionDataset.__getitem__). @@ -367,6 +439,21 @@ sv.plot_images_grid( ) ``` +=== "CreateML" + + We can do so using the [`sv.DetectionDataset.as_createml`](https://supervision.roboflow.com/latest/datasets/core/#supervision.dataset.core.DetectionDataset.as_createml) method to save annotations in [CreateML](https://roboflow.com/formats/createml-json) format. + + ```python + import supervision as sv + + ds = sv.DetectionDataset(...) + + ds.as_createml( + images_directory_path="", + annotations_path="", + ) + ``` + ## Augment Dataset In this section, we'll explore using Supervision in combination with Albumentations to augment our dataset. Data augmentation is a common technique in computer vision to increase the size and diversity of training datasets, leading to improved model performance and generalization. @@ -424,7 +511,7 @@ augmented_annotations = replace( ### What dataset formats does supervision support? -For detection datasets, supervision supports YOLO, COCO JSON, and Pascal VOC. Use `DetectionDataset.from_yolo()`, `from_coco()`, or `from_pascal_voc()` to load, and `as_yolo()`, `as_coco()`, or `as_pascal_voc()` to save. Classification datasets use `ClassificationDataset.from_folder_structure()` and `as_folder_structure()`. +For detection datasets, supervision supports YOLO, COCO JSON, Pascal VOC, and CreateML. Use `DetectionDataset.from_yolo()`, `from_coco()`, `from_pascal_voc()`, or `from_createml()` to load, and `as_yolo()`, `as_coco()`, `as_pascal_voc()`, or `as_createml()` to save. Classification datasets use `ClassificationDataset.from_folder_structure()` and `as_folder_structure()`. ### Can I split a dataset into train/val/test sets? diff --git a/docs/llms.txt b/docs/llms.txt index b03208f2..8eba37af 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -47,7 +47,7 @@ Object tracker wrapper that assigns persistent IDs across video frames. The buil Zone-based counting. `PolygonZone.trigger(detections)` returns a boolean mask for detections currently inside an arbitrary polygon. `LineZone.trigger(detections)` returns `(crossed_in, crossed_out)` arrays for line crossings and requires `detections.tracker_id` so objects can be matched across frames. Both are commonly paired with zone annotators for visualization. ### sv.DetectionDataset and sv.ClassificationDataset -For detection datasets, load, merge, split, and convert between YOLO, COCO JSON, and Pascal VOC formats. Classification datasets use folder-structure import and export via `ClassificationDataset.from_folder_structure()` and `as_folder_structure()`. +For detection datasets, load, merge, split, and convert between YOLO, COCO JSON, Pascal VOC, and CreateML formats. Classification datasets use folder-structure import and export via `ClassificationDataset.from_folder_structure()` and `as_folder_structure()`. ### sv.InferenceSlicer SAHI-style inference slicing: split high-resolution images into overlapping tiles, run detection on each tile, merge results with non-maximum suppression or non-maximum merge. Configure tile overlap in pixels with `overlap_wh`. @@ -144,7 +144,7 @@ Use a tracker to assign persistent IDs. The built-in `sv.ByteTrack` wrapper acce ### What dataset formats does supervision support? -For detection datasets, supervision supports YOLO, COCO JSON, and Pascal VOC. Use `DetectionDataset.from_yolo()`, `from_coco()`, or `from_pascal_voc()` to load, and `as_yolo()`, `as_coco()`, or `as_pascal_voc()` to save. For classification datasets, use `ClassificationDataset.from_folder_structure()` and `as_folder_structure()`. +For detection datasets, supervision supports YOLO, COCO JSON, Pascal VOC, and CreateML. Use `DetectionDataset.from_yolo()`, `from_coco()`, `from_pascal_voc()`, or `from_createml()` to load, and `as_yolo()`, `as_coco()`, `as_pascal_voc()`, or `as_createml()` to save. For classification datasets, use `ClassificationDataset.from_folder_structure()` and `as_folder_structure()`. ### How do I count objects in a zone? diff --git a/src/supervision/dataset/core.py b/src/supervision/dataset/core.py index e8d5a9ae..51c89aa0 100644 --- a/src/supervision/dataset/core.py +++ b/src/supervision/dataset/core.py @@ -18,6 +18,10 @@ from supervision.dataset.formats.coco import ( load_coco_annotations, save_coco_annotations, ) +from supervision.dataset.formats.createml import ( + load_createml_annotations, + save_createml_annotations, +) from supervision.dataset.formats.pascal_voc import ( detections_to_pascal_voc, load_pascal_voc_annotations, @@ -596,6 +600,104 @@ class DetectionDataset(BaseDataset): if data_yaml_path is not None: save_data_yaml(data_yaml_path=data_yaml_path, classes=self.classes) + @classmethod + def from_createml( + cls, + images_directory_path: str, + annotations_path: str, + show_progress: bool = False, + ) -> DetectionDataset: + """ + Creates a Dataset instance from CreateML formatted data. + + CreateML stores object-detection annotations in a single JSON file as a + list of per-image entries, with each box expressed as a pixel-space + centre point plus width and height. Class names are inferred from the + labels present in the file. + + Args: + images_directory_path: The path to the directory containing the + images. + annotations_path: The path to the CreateML json annotation file. + show_progress: If True, display a progress bar during loading. + + Returns: + A DetectionDataset instance containing the loaded images and + annotations. + + Examples: + ```python + import roboflow + from roboflow import Roboflow + import supervision as sv + + roboflow.login() + rf = Roboflow() + + project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID) + dataset = project.version(PROJECT_VERSION).download("createml") + + ds = sv.DetectionDataset.from_createml( + images_directory_path=f"{dataset.location}/train", + annotations_path=f"{dataset.location}/train/_annotations.createml.json", + ) + + ds.classes + # ['dog', 'person'] + ``` + """ + classes, image_paths, annotations = load_createml_annotations( + images_directory_path=images_directory_path, + annotations_path=annotations_path, + show_progress=show_progress, + ) + return DetectionDataset( + classes=classes, images=image_paths, annotations=annotations + ) + + def as_createml( + self, + images_directory_path: str | None = None, + annotations_path: str | None = None, + show_progress: bool = False, + ) -> None: + """ + Exports the dataset to CreateML format. This method saves the + images and their corresponding annotations in CreateML format. + + Args: + images_directory_path: The path to the directory where the images + should be saved. If not provided, images will not be saved. + annotations_path: The path to the CreateML json annotation file. + If not provided, the annotations will not be saved. + show_progress: If True, display a progress bar while saving images. + + Returns: + None. Side-effects only: writes images and/or annotation file. + + Examples: + ```python + import supervision as sv + + ds = sv.DetectionDataset(classes=["dog"], images=[], annotations={}) + ds.as_createml( + images_directory_path="/tmp/images", + annotations_path="/tmp/annotations.json", + ) + ``` + """ + if images_directory_path is not None: + save_dataset_images( + dataset=self, + images_directory_path=images_directory_path, + show_progress=show_progress, + ) + if annotations_path is not None: + save_createml_annotations( + dataset=self, + annotations_path=annotations_path, + ) + @classmethod def from_coco( cls, diff --git a/src/supervision/dataset/formats/createml.py b/src/supervision/dataset/formats/createml.py new file mode 100644 index 00000000..a48e74a0 --- /dev/null +++ b/src/supervision/dataset/formats/createml.py @@ -0,0 +1,320 @@ +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING, Any, cast + +import numpy as np +from tqdm.auto import tqdm + +from supervision.detection.core import Detections +from supervision.utils.file import read_json_file, save_json_file + +if TYPE_CHECKING: + from supervision.dataset.core import DetectionDataset + +CreateMLDict = dict[str, Any] + + +def _resolve_image_path(images_directory_path: str, image_name: str) -> str: + """Resolve and validate an image path against the images directory. + + Rejects annotations whose ``image`` field escapes ``images_directory_path`` + (via ``..`` traversal, an absolute path, or a symlink pointing outside), + mirroring the protection used by the COCO loader. + """ + images_directory_resolved = Path(images_directory_path).resolve() + image_path = Path(images_directory_path) / Path(image_name) + try: + resolved_image_path = image_path.resolve() + except (OSError, ValueError) as exc: + raise ValueError( + f"CreateML annotation refers to image {image_name!r}, which " + f"produces an invalid path: {exc}" + ) from exc + if resolved_image_path == images_directory_resolved: + raise ValueError( + f"CreateML annotation refers to image {image_name!r}, which " + f"resolves to the images directory itself " + f"({images_directory_resolved}). Expected a path to an image file." + ) + if images_directory_resolved not in resolved_image_path.parents: + raise ValueError( + f"CreateML annotation refers to image {image_name!r}, which " + f"resolves to {resolved_image_path} — outside the images " + f"directory {images_directory_resolved}." + ) + if resolved_image_path.is_dir(): + raise ValueError( + f"CreateML annotation refers to image {image_name!r}, which " + f"resolves to directory {resolved_image_path}. Expected a path " + "to an image file." + ) + return str(image_path) + + +def createml_annotations_to_detections( + image_annotations: list[CreateMLDict], class_to_index: dict[str, int] +) -> Detections: + """Convert a single image's CreateML annotations into ``Detections``. + + CreateML stores each box as a pixel-space centre point plus width/height + (``{"x", "y", "width", "height"}``); they are converted to ``xyxy`` corners. + + Args: + image_annotations: List of annotation dicts for one image, each containing + a ``"label"`` key and a ``"coordinates"`` dict with ``"x"``, ``"y"``, + ``"width"``, and ``"height"`` keys. + class_to_index: Mapping from class name to zero-based integer id. + + Returns: + A ``Detections`` instance with ``xyxy`` boxes and ``class_id`` set. + Returns ``Detections.empty()`` when ``image_annotations`` is empty. + + Raises: + ValueError: If an annotation is missing required keys (``"coordinates"``, + ``"label"``, or any coordinate sub-key), or if a coordinate value + cannot be converted to float. + + Examples: + ```python + import supervision as sv + from supervision.dataset.formats.createml import ( + createml_annotations_to_detections, + ) + + annotations = [ + { + "label": "dog", + "coordinates": {"x": 50, "y": 50, "width": 20, "height": 20}, + } + ] + detections = createml_annotations_to_detections(annotations, {"dog": 0}) + # detections.xyxy → [[40, 40, 60, 60]] + ``` + """ + if not image_annotations: + return Detections.empty() + + xyxy = [] + class_ids = [] + for annotation in image_annotations: + try: + coordinates = annotation["coordinates"] + x_center = float(coordinates["x"]) + y_center = float(coordinates["y"]) + width = float(coordinates["width"]) + height = float(coordinates["height"]) + label = annotation["label"] + except (KeyError, TypeError) as exc: + raise ValueError( + f"Malformed CreateML annotation entry {annotation!r}: {exc}" + ) from exc + xyxy.append( + [ + x_center - width / 2, + y_center - height / 2, + x_center + width / 2, + y_center + height / 2, + ] + ) + class_ids.append(class_to_index[label]) + + return Detections( + xyxy=np.array(xyxy, dtype=np.float32), + class_id=np.array(class_ids, dtype=int), + ) + + +def load_createml_annotations( + images_directory_path: str, + annotations_path: str, + show_progress: bool = False, +) -> tuple[list[str], list[str], dict[str, Detections]]: + """Load CreateML object-detection annotations and convert them to ``Detections``. + + CreateML uses a single JSON file containing a list of per-image entries, each + holding axis-aligned bounding boxes. Class names are inferred from the labels + present in the file and assigned stable, sorted, zero-based ids. Because the + format has no explicit category list, a class with no boxes anywhere in the + file will not appear in the returned ``classes``. + + Args: + images_directory_path: Path to the directory containing the images. + annotations_path: Path to the CreateML JSON annotation file. + show_progress: If ``True``, display a tqdm progress bar while loading + annotations. + + Returns: + A tuple of three elements: + + - ``classes`` (``list[str]``): globally sorted class names inferred from + all labels present in the file. + - ``image_paths`` (``list[str]``): joined (but not fully resolved) path + for every entry in the JSON, in file order. + - ``annotations`` (``dict[str, Detections]``): mapping from joined image + path to its ``Detections``. + + Raises: + ValueError: If the JSON root is not a list. + ValueError: If an entry is missing the required ``"image"`` key. + ValueError: If an annotation is missing required coordinate or label keys. + ValueError: If the same image filename appears more than once in the file. + ValueError: If an annotation's ``image`` field resolves to the images + directory itself or to a path outside it (e.g. via ``..`` traversal + or an absolute path). + """ + createml_data = cast( + "list[CreateMLDict]", read_json_file(file_path=annotations_path) + ) + if not isinstance(createml_data, list): + raise ValueError( + f"CreateML annotation file must contain a JSON list at the root, " + f"got {type(createml_data).__name__}." + ) + + try: + classes = sorted( + { + annotation["label"] + for entry in createml_data + for annotation in (entry.get("annotations") or []) + } + ) + except (KeyError, TypeError) as exc: + raise ValueError( + f"Malformed CreateML annotation entry " + f"(missing or non-string 'label'): {exc}" + ) from exc + class_to_index = {class_name: index for index, class_name in enumerate(classes)} + + image_paths: list[str] = [] + annotations: dict[str, Detections] = {} + for entry in tqdm( + createml_data, + desc="Loading CreateML annotations", + disable=not show_progress, + ): + image_name = entry.get("image") + if image_name is None: + raise ValueError( + f"CreateML annotation entry is missing the required 'image' key: " + f"{entry!r}" + ) + image_path = _resolve_image_path( + images_directory_path=images_directory_path, image_name=image_name + ) + if image_path in annotations: + raise ValueError( + f"CreateML annotation file contains duplicate entries for image " + f"{image_name!r}. Each image must appear at most once." + ) + annotations[image_path] = createml_annotations_to_detections( + image_annotations=entry.get("annotations") or [], + class_to_index=class_to_index, + ) + image_paths.append(image_path) + + return classes, image_paths, annotations + + +def detections_to_createml_annotations( + detections: Detections, classes: list[str] +) -> list[CreateMLDict]: + """Convert ``Detections`` into a list of CreateML annotation dicts. + + Each bounding box is stored as a pixel-space centre point plus width and + height, which is the CreateML object-detection convention. + + Args: + detections: The detections to convert. ``class_id`` must not be ``None``. + classes: Ordered list of class names; ``detections.class_id`` values are + used as indices into this list. + + Returns: + A list of dicts, each with a ``"label"`` key (class name) and a + ``"coordinates"`` dict containing ``"x"``, ``"y"``, ``"width"``, and + ``"height"`` in pixel space. + + Raises: + ValueError: If ``detections.class_id`` is ``None``. + + Examples: + ```python + import numpy as np + import supervision as sv + from supervision.dataset.formats.createml import ( + detections_to_createml_annotations, + ) + + detections = sv.Detections( + xyxy=np.array([[40, 40, 60, 60]], dtype=np.float32), + class_id=np.array([0], dtype=int), + ) + detections_to_createml_annotations(detections, classes=["dog"]) + # [{"label": "dog", "coordinates": {"x": 50.0, "y": 50.0, ...}}] + ``` + """ + class_ids = detections.class_id + if class_ids is None: + raise ValueError( + "class_id is required for CreateML export, but the provided " + "Detections has class_id=None." + ) + annotations: list[CreateMLDict] = [] + for xyxy, class_id in zip(detections.xyxy, class_ids): + x_min, y_min, x_max, y_max = (float(value) for value in xyxy) + annotations.append( + { + "label": classes[int(class_id)], + "coordinates": { + "x": (x_min + x_max) / 2, + "y": (y_min + y_max) / 2, + "width": x_max - x_min, + "height": y_max - y_min, + }, + } + ) + return annotations + + +def save_createml_annotations( + dataset: DetectionDataset, + annotations_path: str, +) -> None: + """Export a ``DetectionDataset`` to a CreateML object-detection JSON file. + + Only the filename component of each image path is stored in the JSON (e.g. + ``"img.jpg"`` rather than ``"/data/train/img.jpg"``). This matches CreateML + convention and means the loader reconstructs paths relative to + ``images_directory_path``. As a consequence, two images with the same + basename from different directories will produce duplicate ``"image"`` keys + in the output and cannot be round-tripped correctly. + + Args: + dataset: The ``DetectionDataset`` to write. + annotations_path: Output path for the CreateML JSON file. Parent + directories are created if they do not already exist. + + Examples: + ```python + import supervision as sv + from supervision.dataset.formats.createml import save_createml_annotations + + dataset = sv.DetectionDataset(classes=["dog"], images=[], annotations={}) + save_createml_annotations(dataset, "/tmp/annotations.json") + ``` + """ + Path(annotations_path).parent.mkdir(parents=True, exist_ok=True) + createml_data: list[CreateMLDict] = [ + { + "image": Path(image_path).name, + "annotations": detections_to_createml_annotations( + detections=dataset.annotations[image_path], classes=dataset.classes + ), + } + for image_path in dataset.image_paths + ] + save_json_file( + data=createml_data, # type: ignore[arg-type] # save_json_file accepts list at runtime + file_path=annotations_path, + ) diff --git a/tests/dataset/formats/test_createml.py b/tests/dataset/formats/test_createml.py new file mode 100644 index 00000000..addc290b --- /dev/null +++ b/tests/dataset/formats/test_createml.py @@ -0,0 +1,382 @@ +"""Tests for CreateML object-detection annotation load/save and conversion helpers.""" + +from __future__ import annotations + +import json +from contextlib import ExitStack as DoesNotRaise +from pathlib import Path + +import numpy as np +import pytest + +from supervision.dataset.core import DetectionDataset +from supervision.dataset.formats.createml import ( + createml_annotations_to_detections, + detections_to_createml_annotations, + load_createml_annotations, + save_createml_annotations, +) +from supervision.detection.core import Detections + + +class TestCreatemlAnnotationsToDetections: + @pytest.mark.parametrize( + ("image_annotations", "class_to_index", "expected_result", "exception"), + [ + pytest.param( + [], + {}, + Detections.empty(), + DoesNotRaise(), + id="empty-annotations", + ), + pytest.param( + [ + { + "label": "dog", + "coordinates": {"x": 50, "y": 50, "width": 20, "height": 20}, + } + ], + {"dog": 0}, + Detections( + xyxy=np.array([[40, 40, 60, 60]], dtype=np.float32), + class_id=np.array([0], dtype=int), + ), + DoesNotRaise(), + id="single-centre-box-to-xyxy", + ), + pytest.param( + [ + { + "label": "cat", + "coordinates": {"x": 10, "y": 10, "width": 4, "height": 4}, + }, + { + "label": "dog", + "coordinates": {"x": 30, "y": 20, "width": 10, "height": 8}, + }, + ], + {"cat": 0, "dog": 1}, + Detections( + xyxy=np.array([[8, 8, 12, 12], [25, 16, 35, 24]], dtype=np.float32), + class_id=np.array([0, 1], dtype=int), + ), + DoesNotRaise(), + id="multi-class-distinct-ids", + ), + pytest.param( + [ + { + "label": "dog", + "coordinates": {"x": 10, "y": 10, "width": 4, "height": 4}, + }, + { + "label": "dog", + "coordinates": {"x": 30, "y": 30, "width": 4, "height": 4}, + }, + ], + {"dog": 0}, + Detections( + xyxy=np.array([[8, 8, 12, 12], [28, 28, 32, 32]], dtype=np.float32), + class_id=np.array([0, 0], dtype=int), + ), + DoesNotRaise(), + id="duplicate-labels-two-detections-same-id", + ), + ], + ) + def test_converts_annotations( + self, + image_annotations: list[dict], + class_to_index: dict[str, int], + expected_result: Detections, + exception: Exception, + ) -> None: + """Converts CreateML annotation list to Detections with correct xyxy and ids.""" + with exception: + result = createml_annotations_to_detections( + image_annotations=image_annotations, class_to_index=class_to_index + ) + np.testing.assert_array_almost_equal(result.xyxy, expected_result.xyxy) + assert (result.class_id is None) == (expected_result.class_id is None) + if expected_result.class_id is not None: + np.testing.assert_array_equal(result.class_id, expected_result.class_id) + + def test_raises_on_missing_coordinates_key(self) -> None: + """Raises ValueError when an annotation entry lacks the 'coordinates' key.""" + with pytest.raises(ValueError, match="Malformed"): + createml_annotations_to_detections( + image_annotations=[{"label": "dog"}], + class_to_index={"dog": 0}, + ) + + def test_raises_on_missing_label_key(self) -> None: + """Raises ValueError when an annotation entry lacks the 'label' key.""" + with pytest.raises(ValueError, match="Malformed"): + createml_annotations_to_detections( + image_annotations=[ + {"coordinates": {"x": 10, "y": 10, "width": 4, "height": 4}} + ], + class_to_index={"dog": 0}, + ) + + def test_raises_on_missing_coordinate_subkey(self) -> None: + """Raises ValueError when a coordinates dict is missing a required sub-key.""" + with pytest.raises(ValueError, match="Malformed"): + createml_annotations_to_detections( + image_annotations=[ + {"label": "dog", "coordinates": {"x": 10, "y": 10, "width": 4}} + ], + class_to_index={"dog": 0}, + ) + + def test_raises_when_coordinates_is_none(self) -> None: + """Raises ValueError when the coordinates value is None.""" + with pytest.raises(ValueError, match="Malformed"): + createml_annotations_to_detections( + image_annotations=[{"label": "dog", "coordinates": None}], + class_to_index={"dog": 0}, + ) + + +class TestDetectionsToCreatemlAnnotations: + def test_round_trips_coordinates(self) -> None: + """Round-trip: xyxy corners convert to CreateML centre+wh and back correctly.""" + detections = Detections( + xyxy=np.array([[40, 40, 60, 60]], dtype=np.float32), + class_id=np.array([1], dtype=int), + ) + + result = detections_to_createml_annotations( + detections=detections, classes=["cat", "dog"] + ) + + assert result == [ + { + "label": "dog", + "coordinates": {"x": 50.0, "y": 50.0, "width": 20.0, "height": 20.0}, + } + ] + + def test_raises_when_class_id_is_none(self) -> None: + """Raises ValueError when Detections.class_id is None.""" + detections = Detections(xyxy=np.array([[0, 0, 10, 10]], dtype=np.float32)) + + with pytest.raises(ValueError, match="class_id"): + detections_to_createml_annotations(detections=detections, classes=["dog"]) + + +class TestLoadCreatemlAnnotations: + def test_loads_basic_annotations(self, tmp_path: Path) -> None: + """Loads classes, image_paths, and Detections from a valid CreateML file.""" + annotations_path = tmp_path / "annotations.json" + payload = [ + { + "image": "a.jpg", + "annotations": [ + { + "label": "dog", + "coordinates": {"x": 50, "y": 50, "width": 20, "height": 20}, + } + ], + }, + {"image": "b.jpg", "annotations": []}, + ] + annotations_path.write_text(json.dumps(payload)) + + classes, image_paths, annotations = load_createml_annotations( + images_directory_path=str(tmp_path), + annotations_path=str(annotations_path), + ) + + assert classes == ["dog"] + assert image_paths == [str(tmp_path / "a.jpg"), str(tmp_path / "b.jpg")] + detections = annotations[str(tmp_path / "a.jpg")] + np.testing.assert_array_almost_equal( + detections.xyxy, np.array([[40, 40, 60, 60]], dtype=np.float32) + ) + np.testing.assert_array_equal(detections.class_id, np.array([0], dtype=int)) + assert len(annotations[str(tmp_path / "b.jpg")]) == 0 + + def test_assigns_global_sorted_class_ids(self, tmp_path: Path) -> None: + """Class ids are globally sorted regardless of per-image label order.""" + annotations_path = tmp_path / "annotations.json" + payload = [ + { + "image": "a.jpg", + "annotations": [ + { + "label": "zebra", + "coordinates": {"x": 10, "y": 10, "width": 4, "height": 4}, + } + ], + }, + { + "image": "b.jpg", + "annotations": [ + { + "label": "ant", + "coordinates": {"x": 20, "y": 20, "width": 6, "height": 6}, + } + ], + }, + ] + annotations_path.write_text(json.dumps(payload)) + + classes, image_paths, annotations = load_createml_annotations( + images_directory_path=str(tmp_path), + annotations_path=str(annotations_path), + ) + + assert classes == ["ant", "zebra"] + assert image_paths == [str(tmp_path / "a.jpg"), str(tmp_path / "b.jpg")] + np.testing.assert_array_equal( + annotations[str(tmp_path / "a.jpg")].class_id, np.array([1], dtype=int) + ) + np.testing.assert_array_equal( + annotations[str(tmp_path / "b.jpg")].class_id, np.array([0], dtype=int) + ) + + def test_raises_on_path_traversal(self, tmp_path: Path) -> None: + """Raises ValueError when 'image' field attempts directory traversal.""" + annotations_path = tmp_path / "annotations.json" + payload = [{"image": "../evil.jpg", "annotations": []}] + annotations_path.write_text(json.dumps(payload)) + + with pytest.raises(ValueError, match="outside"): + load_createml_annotations( + images_directory_path=str(tmp_path / "images"), + annotations_path=str(annotations_path), + ) + + def test_raises_on_absolute_path(self, tmp_path: Path) -> None: + """Raises ValueError when 'image' is an absolute path outside images dir.""" + annotations_path = tmp_path / "annotations.json" + outside = tmp_path.parent / "evil.jpg" + payload = [{"image": str(outside), "annotations": []}] + annotations_path.write_text(json.dumps(payload)) + + with pytest.raises(ValueError, match="outside"): + load_createml_annotations( + images_directory_path=str(tmp_path), + annotations_path=str(annotations_path), + ) + + def test_raises_when_image_is_the_directory_itself(self, tmp_path: Path) -> None: + """Raises ValueError when 'image' resolves to the images directory itself.""" + annotations_path = tmp_path / "annotations.json" + payload = [{"image": ".", "annotations": []}] + annotations_path.write_text(json.dumps(payload)) + + with pytest.raises(ValueError, match="directory"): + load_createml_annotations( + images_directory_path=str(tmp_path), + annotations_path=str(annotations_path), + ) + + def test_raises_when_json_root_is_dict(self, tmp_path: Path) -> None: + """Raises ValueError when the JSON root is a dict instead of a list.""" + annotations_path = tmp_path / "annotations.json" + annotations_path.write_text(json.dumps({"image": "a.jpg", "annotations": []})) + + with pytest.raises(ValueError, match="JSON list"): + load_createml_annotations( + images_directory_path=str(tmp_path), + annotations_path=str(annotations_path), + ) + + def test_raises_on_missing_image_key(self, tmp_path: Path) -> None: + """Raises ValueError when an entry lacks the required 'image' key.""" + annotations_path = tmp_path / "annotations.json" + annotations_path.write_text(json.dumps([{"annotations": []}])) + + with pytest.raises(ValueError, match="'image'"): + load_createml_annotations( + images_directory_path=str(tmp_path), + annotations_path=str(annotations_path), + ) + + def test_raises_on_duplicate_image_entry(self, tmp_path: Path) -> None: + """Raises ValueError when the same image filename appears more than once.""" + annotations_path = tmp_path / "annotations.json" + payload = [ + {"image": "a.jpg", "annotations": []}, + {"image": "a.jpg", "annotations": []}, + ] + annotations_path.write_text(json.dumps(payload)) + + with pytest.raises(ValueError, match="duplicate"): + load_createml_annotations( + images_directory_path=str(tmp_path), + annotations_path=str(annotations_path), + ) + + +class TestSaveCreatemlAnnotations: + def test_empty_dataset_writes_empty_list(self, tmp_path: Path) -> None: + """Empty dataset serialises to an empty JSON array.""" + annotations_path = tmp_path / "nested" / "annotations.json" + dataset = DetectionDataset(classes=[], images=[], annotations={}) + + save_createml_annotations( + dataset=dataset, annotations_path=str(annotations_path) + ) + + assert json.loads(annotations_path.read_text()) == [] + + def test_save_load_round_trip(self, tmp_path: Path) -> None: + """Save then load preserves class names, image paths, and bounding boxes.""" + images_directory_path = tmp_path / "images" + annotations_path = tmp_path / "annotations.json" + classes = ["cat", "dog"] + image_paths = [str(images_directory_path / "a.jpg")] + annotations = { + image_paths[0]: Detections( + xyxy=np.array([[8, 8, 12, 12], [25, 16, 35, 24]], dtype=np.float32), + class_id=np.array([0, 1], dtype=int), + ) + } + dataset = DetectionDataset( + classes=classes, images=image_paths, annotations=annotations + ) + + save_createml_annotations( + dataset=dataset, annotations_path=str(annotations_path) + ) + loaded_classes, _, loaded_annotations = load_createml_annotations( + images_directory_path=str(images_directory_path), + annotations_path=str(annotations_path), + ) + + assert loaded_classes == classes + loaded = loaded_annotations[str(images_directory_path / "a.jpg")] + np.testing.assert_array_almost_equal( + loaded.xyxy, annotations[image_paths[0]].xyxy + ) + np.testing.assert_array_equal( + loaded.class_id, annotations[image_paths[0]].class_id + ) + + def test_save_load_round_trip_float_coordinates(self, tmp_path: Path) -> None: + """Float32 coordinates survive a save/load cycle within float32 precision.""" + images_directory_path = tmp_path / "images" + annotations_path = tmp_path / "annotations.json" + xyxy = np.array([[10.3, 7.9, 44.1, 88.6]], dtype=np.float32) + image_paths = [str(images_directory_path / "a.jpg")] + annotations = { + image_paths[0]: Detections(xyxy=xyxy, class_id=np.array([0], dtype=int)) + } + dataset = DetectionDataset( + classes=["dog"], images=image_paths, annotations=annotations + ) + + save_createml_annotations( + dataset=dataset, annotations_path=str(annotations_path) + ) + _, _, loaded_annotations = load_createml_annotations( + images_directory_path=str(images_directory_path), + annotations_path=str(annotations_path), + ) + + loaded = loaded_annotations[str(images_directory_path / "a.jpg")] + np.testing.assert_array_almost_equal(loaded.xyxy, xyxy, decimal=4)