airo-dataset-tools 2025.4.0__tar.gz

airo_dataset_tools-2025.4.0/PKG-INFO

@@ -0,0 +1,24 @@
+Metadata-Version: 2.4
+Name: airo-dataset-tools
+Version: 2025.4.0
+Summary: Scripts for loading and converting datasets for the Ghent University AI and Robotics Lab
+Author: Victor-Louis De Gusseme
+Author-email: victorlouisdg@gmail.com
+Requires-Dist: numpy<2.0
+Requires-Dist: pydantic>2.0.0
+Requires-Dist: opencv-contrib-python==4.8.1.78
+Requires-Dist: opencv-python-headless==4.8.1.78
+Requires-Dist: pycocotools
+Requires-Dist: xmltodict
+Requires-Dist: tqdm
+Requires-Dist: fiftyone
+Requires-Dist: Pillow
+Requires-Dist: types-Pillow
+Requires-Dist: albumentations
+Requires-Dist: click
+Requires-Dist: airo-typing==2025.4.0
+Requires-Dist: airo-spatial-algebra==2025.4.0
+Dynamic: author
+Dynamic: author-email
+Dynamic: requires-dist
+Dynamic: summary

airo_dataset_tools-2025.4.0/README.md

@@ -0,0 +1,20 @@
+# airo-dataset-tools
+Tools for working with datasets.
+They fall into two categories:
+
+[**COCO related tools**](airo_dataset_tools/coco_tools/README.md):
+* COCO dataset loading (and creation)
+* FiftyOne visualisation
+* Albumentation transforms
+* COCO  to YOLO conversion.
+* CVAT labeling workflow
+
+[**Data formats**](airo_dataset_tools/data_parsers/README.md):
+* 3D poses
+* Camera instrinsics
+
+
+> [Pydantic](https://docs.pydantic.dev/latest/) is used heavily throughout this package.
+It allows you to easily create Python objects that can be saved and loaded to and from JSON files.
+
+[**CVAT labeling workflow & tools**](airo_dataset_tools/cvat_labeling/readme.md)

airo_dataset_tools-2025.4.0/airo_dataset_tools/__init__.py

@@ -0,0 +1,5 @@
+# init files are still required in modern python!
+# https://peps.python.org/pep-0420/ introduced implicit namespace packages
+# but for building and many toolings, you still need to have __init__ files (at least in the root of the package).
+# e.g. if you remove this init file and try to build with pip install .
+# you won't be able to import the dummy module.

airo_dataset_tools-2025.4.0/airo_dataset_tools/cli.py

@@ -0,0 +1,136 @@
+"""CLI interface for this package"""
+
+import json
+import os
+from typing import List, Optional
+
+import click
+from airo_dataset_tools.coco_tools.change_coco_images_prefix import change_coco_json_image_prefix
+from airo_dataset_tools.coco_tools.coco_instances_to_yolo import create_yolo_dataset_from_coco_instances_dataset
+from airo_dataset_tools.coco_tools.fiftyone_viewer import view_coco_dataset
+from airo_dataset_tools.coco_tools.merge_datasets import merge_coco_datasets
+from airo_dataset_tools.coco_tools.split_dataset import split_and_save_coco_dataset
+from airo_dataset_tools.coco_tools.transform_dataset import resize_coco_dataset
+from airo_dataset_tools.cvat_labeling.convert_cvat_to_coco import cvat_image_to_coco
+
+
+@click.group()
+def cli() -> None:
+    """CLI entrypoint for airo-dataset-tools"""
+
+
+@cli.command(name="fiftyone-coco-viewer")  # no help, takes the docstring of the function.
+@click.argument("annotations-json-path", type=click.Path(exists=True))
+@click.option(
+    "--dataset-dir",
+    required=False,
+    type=click.Path(exists=True),
+    help="optional directory relative to which the image paths in the coco dataset are specified",
+)
+@click.option(
+    "--label-types",
+    "-l",
+    multiple=True,
+    type=click.Choice(["detections", "segmentations", "keypoints"]),
+    help="add an argument for each label type you want to load (default: all)",
+)
+def view_coco_dataset_cli(
+    annotations_json_path: str, dataset_dir: str, label_types: Optional[List[str]] = None
+) -> None:
+    """Explore COCO dataset with FiftyOne"""
+    view_coco_dataset(annotations_json_path, dataset_dir, label_types)
+
+
+@cli.command(name="convert-cvat-to-coco-keypoints")
+@click.argument("cvat_xml_file", type=str, required=True)
+@click.argument("coco-categories-json-file", type=str, required=True)
+@click.option("--add-bbox", is_flag=True, default=False, help="include bounding box in coco annotations")
+@click.option("--add-segmentation", is_flag=True, default=False, help="include segmentation in coco annotations")
+def convert_cvat_to_coco_cli(
+    cvat_xml_file: str, coco_categories_json_file: str, add_bbox: bool, add_segmentation: bool
+) -> None:
+    """Convert CVAT XML to COCO keypoints json according to specified coco categories"""
+    coco = cvat_image_to_coco(
+        cvat_xml_file, coco_categories_json_file, add_bbox=add_bbox, add_segmentation=add_segmentation
+    )
+    path = os.path.dirname(cvat_xml_file)
+    filename = os.path.basename(cvat_xml_file)
+    path = os.path.join(path, filename.split(".")[0] + ".json")
+    with open(path, "w") as file:
+        json.dump(coco, file)
+
+
+@cli.command(name="resize-coco-dataset")
+@click.argument("annotations-json-path", type=click.Path(exists=True))
+@click.option("--width", type=int, required=True)
+@click.option("--height", type=int, required=True)
+@click.option("--target-dataset-dir", type=str, required=False)
+def resize_coco_dataset_cli(
+    annotations_json_path: str, width: int, height: int, target_dataset_dir: Optional[str]
+) -> None:
+    """Resize a COCO dataset. Will create a new directory with the resized dataset at the specified target_dataset_dir.
+    Dataset is assumed to be
+    /dir
+        annotations.json # contains relative paths w.r.t. /dir
+        ...
+    """
+    resize_coco_dataset(annotations_json_path, width, height, target_dataset_dir=target_dataset_dir)
+
+
+@cli.command(name="coco-instances-to-yolo")
+@click.option("--coco-json", type=str)
+@click.option("--target-dir", type=str)
+@click.option("--use-segmentation", is_flag=True)
+def coco_intances_to_yolo(coco_json: str, target_dir: str, use_segmentation: bool) -> None:
+    """Create a YOLO detections/segmentations dataset from a coco instances dataset"""
+    print(f"converting coco instances dataset {coco_json} to yolo dataset {target_dir}")
+    create_yolo_dataset_from_coco_instances_dataset(coco_json, target_dir, use_segmentation=use_segmentation)
+
+
+@cli.command(name="split-coco-dataset")
+@click.argument("json-path", type=click.Path(exists=True))
+@click.option("--split-ratio", type=float, multiple=True, required=True)
+@click.option("--shuffle-before-splitting", is_flag=True, default=True)
+def split_coco_dataset_cli(json_path: str, split_ratio: List[float], shuffle_before_splitting: bool) -> None:
+    """Split a COCO dataset into subsets according to the specified relative ratios and save them to disk.
+    Images are split with their corresponding annotations. No guarantees on class balance or annotation ratios.
+
+    If two ratios are specified, the dataset will be split into two subsets. these will be called train/val by default.
+    If three ratios are specified, the dataset will be split into three subsets. these will be called train/val/test by default.
+
+    e.g. split-coco-dataset <path> --split-ratio 0.8 --split-ratio 0.2
+    """
+    split_and_save_coco_dataset(json_path, split_ratio, shuffle_before_splitting=shuffle_before_splitting)
+
+
+@cli.command(name="change-coco-images-prefix")
+@click.argument("coco-json", type=click.Path(exists=True))
+@click.option("--current-prefix", type=str, required=True)
+@click.option("--new-prefix", type=str, required=True)
+@click.option("--target-json-path", type=str)
+def change_coco_images_prefix_cli(
+    coco_json: str, current_prefix: str, new_prefix: str, target_json_path: Optional[str] = None
+) -> None:
+    """change the prefix of images in a coco dataset."""
+    return change_coco_json_image_prefix(coco_json, current_prefix, new_prefix, target_json_path)
+
+
+@cli.command(name="merge-coco-datasets")
+@click.argument("coco-json-1", type=click.Path(exists=True))
+@click.argument("coco-json-2", type=click.Path(exists=True))
+@click.option(
+    "--target-json-path",
+    type=str,
+    help="optional path to save the merged dataset to. If none is provided, a new directory will be created in the parent directory of coco-json-1 ",
+)
+def merge_coco_datasets_cli(coco_json_1: str, coco_json_2: str, target_json_path: Optional[str] = None) -> None:
+    """merge two coco datasets into a single dataset."""
+    if not target_json_path:
+        target_json_path = os.path.join(os.path.dirname(coco_json_1), "merged")
+        os.makedirs(target_json_path, exist_ok=True)
+        target_json_path = os.path.join(target_json_path, "annotations.json")
+    return merge_coco_datasets(coco_json_1, coco_json_2, target_json_path)
+
+
+if __name__ == "__main__":
+    cli()

airo_dataset_tools-2025.4.0/airo_dataset_tools/coco_tools/change_coco_images_prefix.py

@@ -0,0 +1,40 @@
+import json
+from typing import Optional
+
+from airo_dataset_tools.data_parsers.coco import CocoInstancesDataset
+
+
+def change_coco_dataset_image_prefix(
+    coco_dataset: CocoInstancesDataset, current_prefix: str, target_prefix: str
+) -> CocoInstancesDataset:
+    """Change the prefix of the image paths in a COCO dataset. Can be used to change the base directory of the images in a dataset.
+
+    Args:
+        coco_dataset: the dataset to modify
+        current_prefix: the current prefix of the image paths in the dataset
+        target_prefix: the target prefix of the image paths in the dataset
+
+    e.g.
+
+    if the images are currently relative to the image folder inside the dataset and you want to make them relative to the dataset folder:
+    change_coco_image_base_dir(coco_dataset, "", "images/")
+    """
+
+    for image in coco_dataset.images:
+        image.file_name = image.file_name.removeprefix(current_prefix)
+        image.file_name = target_prefix + image.file_name
+
+    return coco_dataset
+
+
+def change_coco_json_image_prefix(
+    coco_json_file: str, current_prefix: str, target_prefix: str, target_json_file: Optional[str]
+) -> None:
+    with open(coco_json_file, "r") as f:
+        coco_dataset = CocoInstancesDataset(**json.load(f))
+
+    coco_dataset = change_coco_dataset_image_prefix(coco_dataset, current_prefix, target_prefix)
+    if target_json_file is None:
+        target_json_file = coco_json_file
+    with open(target_json_file, "w") as f:
+        json.dump(coco_dataset.model_dump(), f)

airo_dataset_tools-2025.4.0/airo_dataset_tools/coco_tools/coco_instances_to_yolo.py

@@ -0,0 +1,143 @@
+import json
+import pathlib
+from collections import defaultdict
+
+import cv2
+import tqdm
+from airo_dataset_tools.data_parsers.coco import CocoInstancesDataset
+from airo_dataset_tools.segmentation_mask_converter import BinarySegmentationMask
+
+
+def create_yolo_dataset_from_coco_instances_dataset(
+    coco_dataset_json_path: str, target_directory: str, use_segmentation: bool = False
+) -> None:
+    """Converts a coco dataset to a yolo dataset, either segmentations or detections
+
+
+    coco dataset is expected:
+    dataset/
+        images/
+        <>.json # containing paths relative to /dataset
+
+    yolo dataset format will be
+
+    dataset/
+        images/
+        labels/
+        <>.txt # paths as in coco dataset /images subdir.
+
+        # ultralytics uses an additional <>.json, this has to be created manually..
+        # we do export a .names file, as in the darknet yolo format.
+        # each line contains a class name, the line number is the class id
+        obj.names
+    Args:
+        coco_dataset_json_path: _description_
+        target_directory: _description_
+        use_segmentation: create a segmentation dataset instead of a detection dataset (requires segmentation annotations)
+
+
+    quiz question: how many people would have written a similar converter in the past?...
+
+    """
+    _coco_dataset_json_path = pathlib.Path(coco_dataset_json_path)
+    _target_directory = pathlib.Path(target_directory)
+
+    # load the json file into the coco dataset object
+    annotations = json.load(open(_coco_dataset_json_path, "r"))
+    coco_dataset = CocoInstancesDataset(**annotations)
+
+    target_dir = pathlib.Path(_target_directory)
+    target_dir.mkdir(parents=True, exist_ok=True)
+    image_dir = target_dir / "images"
+    label_dir = target_dir / "labels"
+    image_dir.mkdir(parents=True, exist_ok=True)
+    label_dir.mkdir(parents=True, exist_ok=True)
+
+    # create the yolo dataset
+    annotations = coco_dataset.annotations
+    image_id_to_image = {image.id: image for image in coco_dataset.images}
+    image_id_to_annotations = defaultdict(list)
+    for annotation in annotations:
+        image_id_to_annotations[annotation.image_id].append(annotation)
+
+    category_id_to_category = {category.id: category for category in coco_dataset.categories}
+    # sort by original COCO ID, but whereas COCO does not care about ID ranges, YOLO expects them to be in the range 0..N-1
+    # so we sort them by ID and then use the index as the new ID
+    yolo_category_index = list(sorted(category_id_to_category.values(), key=lambda category: category.id))
+
+    for image_id, annotations in tqdm.tqdm(image_id_to_annotations.items()):
+        coco_image = image_id_to_image[image_id]
+        image_path = pathlib.Path(coco_image.file_name)
+        if not image_path.is_absolute():
+            image_path = _coco_dataset_json_path.parent / image_path
+
+        relative_image_path = image_path.relative_to(_coco_dataset_json_path.parent)
+
+        # ultralytics parser finds the 'latest' occurance of 'images' in the dataset path,
+        #  so we need to replace any occurance of 'image' with 'img'
+        #  https://github.com/ultralytics/ultralytics/issues/3581
+
+        relative_image_path_str = str(relative_image_path).replace("image", "img")
+        relative_image_path = pathlib.Path(relative_image_path_str)
+
+        image = cv2.imread(str(image_path))
+        height, width, _ = image.shape
+
+        label_path = label_dir / f"{relative_image_path.with_suffix('')}.txt"
+        label_path.parent.mkdir(parents=True, exist_ok=True)
+
+        with open(label_path, "w") as file:
+            for annotation in annotations:
+                category = category_id_to_category[annotation.category_id]
+                yolo_id = yolo_category_index.index(category)
+                if use_segmentation:
+                    segmentation = annotation.segmentation
+                    # convert to **single** polygon
+                    segmentation = BinarySegmentationMask.from_coco_segmentation_mask(segmentation, width, height)
+                    segmentation = segmentation.as_single_polygon
+
+                    if segmentation is None:
+                        # should actually never happen as each annotation is assumed to have a segmentation if you pass use_segmentation=True
+                        # but we filter it for convenience to deal with edge cases
+                        print(f"skipping annotation for image {image_path}, as it has no segmentation")
+                        continue
+
+                    file.write(f"{yolo_id}")
+                    for (x, y) in zip(segmentation[0::2], segmentation[1::2]):
+                        file.write(f" {x/width} {y/height}")
+                    file.write("\n")
+
+                else:
+                    x, y, w, h = annotation.bbox
+                    x_center = x + w / 2
+                    y_center = y + h / 2
+                    x_center /= width
+                    y_center /= height
+                    w /= width
+                    h /= height
+                    file.write(f"{yolo_id} {x_center} {y_center} {w} {h}\n")
+
+        image_target_path = image_dir / relative_image_path
+        image_target_path.parent.mkdir(parents=True, exist_ok=True)
+        cv2.imwrite(str(image_target_path), image)
+
+    # create the obj.names file
+    with open(target_dir / "obj.names", "w") as file:
+        # sort categories by id
+        for category in yolo_category_index:
+            file.write(f"{category.name}\n")
+
+
+if __name__ == "__main__":
+    """example usage"""
+    import os
+
+    path = pathlib.Path(__file__).parents[1] / "cvat_labeling" / "example" / "coco.json"
+
+    coco_json_path = str(path)
+    coco_dir = os.path.dirname(coco_json_path)
+    coco_file_name = os.path.basename(coco_json_path)
+    coco_target_dir = os.path.join(os.path.dirname(coco_dir), "yolo_dataset")
+    os.makedirs(coco_target_dir, exist_ok=True)
+
+    create_yolo_dataset_from_coco_instances_dataset(coco_json_path, coco_target_dir, use_segmentation=True)

airo_dataset_tools-2025.4.0/airo_dataset_tools/coco_tools/fiftyone_viewer.py

@@ -0,0 +1,29 @@
+import os
+from typing import List, Optional
+
+
+def view_coco_dataset(
+    labels_json_path: str, dataset_dir: Optional[str] = None, label_types: Optional[List[str]] = None
+) -> None:
+    """visualize a coco dataset in fiftyone"""
+    # lazy import because fiftyone is slow to import and makes the CLI slow.
+    import fiftyone as fo
+
+    if dataset_dir is None:
+        dataset_dir = os.path.dirname(labels_json_path)
+    if label_types is None or not label_types:
+        label_types = ["detections", "segmentations", "keypoints"]
+    else:
+        assert all([label_type in ["detections", "segmentations", "keypoints"] for label_type in label_types])
+
+    labels_json_path = os.path.realpath(labels_json_path)
+
+    dataset = fo.Dataset.from_dir(
+        dataset_type=fo.types.COCODetectionDataset,
+        label_types=label_types,
+        data_path=dataset_dir,
+        labels_path=labels_json_path,
+    )
+
+    session = fo.launch_app(dataset)
+    session.wait()

airo_dataset_tools-2025.4.0/airo_dataset_tools/coco_tools/merge_datasets.py

@@ -0,0 +1,132 @@
+"""merge 2 coco datasets into one"""
+
+import json
+import pathlib
+import shutil
+
+import tqdm
+from airo_dataset_tools.data_parsers.coco import CocoInstancesDataset
+
+
+def merge_coco_annotations(dataset1: CocoInstancesDataset, dataset2: CocoInstancesDataset) -> CocoInstancesDataset:
+    """merge 2 coco annotations schemas. Categories and Annotations are assumed to be unique. Images can be present in both datasets.
+
+    Images will be checked for duplicates based on their file name.
+
+    Annotation IDs will be changed to avoid conflicts and their image IDs will be updated if needed."""
+    categories_1 = dataset1.categories
+    categories_2 = dataset2.categories
+
+    merged_categories = list(categories_1)
+    for category in categories_2:
+        if category.id not in [category.id for category in merged_categories]:
+            merged_categories.append(category)
+        else:
+            same_id_category = merged_categories[[category.id for category in merged_categories].index(category.id)]
+            if category != same_id_category:
+                raise ValueError(
+                    "Categories with the same ID are not equal: " + str(category) + " " + str(same_id_category)
+                )
+
+    merged_images = dataset1.images
+    max_dataset_image_id = max([image.id for image in merged_images])
+    merged_image_paths = [image.file_name for image in merged_images]
+
+    dataset_2_image_id_mapping = {}
+    for image in dataset2.images:
+        if image.file_name not in merged_image_paths:
+            # create a new ID to avoid collisions
+            old_id = image.id
+            image.id = max_dataset_image_id + 1
+            max_dataset_image_id += 1
+            dataset_2_image_id_mapping[old_id] = image.id
+            merged_images.append(image)
+            merged_image_paths.append(image.file_name)
+        else:
+            # dataset already contained this image, so we need to remap the annotations
+            dataset_2_image_id_mapping[image.id] = merged_images[merged_image_paths.index(image.file_name)].id
+
+    merged_annotations = list(dataset1.annotations)
+    max_dataset_annotation_id = max([annotation.id for annotation in merged_annotations])
+    for annotation in dataset2.annotations:
+        if annotation.image_id in dataset_2_image_id_mapping.keys():
+            annotation.image_id = dataset_2_image_id_mapping[annotation.image_id]
+        annotation.id = max_dataset_annotation_id + 1
+        max_dataset_annotation_id += 1
+        merged_annotations.append(annotation)
+
+    merged_dataset = CocoInstancesDataset(
+        categories=merged_categories, images=merged_images, annotations=merged_annotations
+    )
+    return merged_dataset
+
+
+def merge_coco_image_folders(dataset1_base_dir: str, dataset2_base_dir: str, target_dir: str) -> None:
+    """merge 2 image folders into one. Images will be copied to the target folder.
+
+    all base_dirs have following setup:
+
+    base_dir
+    --- images
+    ------ image1.<>
+    ------ image2.<>
+    <>.json // annotations with path relative to base_dir
+    """
+
+    dataset1_base_dir_path = pathlib.Path(dataset1_base_dir)
+    dataset2_base_dir_path = pathlib.Path(dataset2_base_dir)
+    target_dir_path = pathlib.Path(target_dir)
+
+    dataset1_image_paths = [image_path for image_path in dataset1_base_dir_path.iterdir()]
+    dataset2_image_paths = [image_path for image_path in dataset2_base_dir_path.iterdir()]
+
+    target_image_dir = target_dir_path / "images"
+    target_image_dir.mkdir(parents=True, exist_ok=True)
+
+    for image_path in tqdm.tqdm(
+        dataset1_image_paths, desc=f"copying images from {dataset1_base_dir_path.name} to {target_dir_path.name}"
+    ):
+        shutil.copy(image_path, target_image_dir / image_path.name)
+
+    for image_path in tqdm.tqdm(
+        dataset2_image_paths, desc=f"copying images from {dataset2_base_dir_path.name} to {target_dir_path.name}"
+    ):
+        if not (target_image_dir / image_path.name).exists():
+            shutil.copy(image_path, target_image_dir / image_path.name)
+
+
+def merge_coco_datasets(json_path_1: str, json_path_2: str, target_json_path: str) -> None:
+    """merge 2 coco datasets into one. Categories and Annotations are assumed to be unique. Images can be present in both datasets.
+
+    Images will be checked for duplicates based on their file name.
+
+    Annotation IDs will be changed to avoid conflicts and their image IDs will be updated if needed."""
+
+    image_path_1 = pathlib.Path(json_path_1).parent / "images"
+    image_path_2 = pathlib.Path(json_path_2).parent / "images"
+
+    merge_coco_image_folders(str(image_path_1), str(image_path_2), str(pathlib.Path(target_json_path).parent))
+
+    dataset1 = None
+    with open(json_path_1, "r") as f:
+        dataset1 = CocoInstancesDataset(**json.load(f))
+
+    dataset2 = None
+    with open(json_path_2, "r") as f:
+        dataset2 = CocoInstancesDataset(**json.load(f))
+
+    merged_dataset = merge_coco_annotations(dataset1, dataset2)
+
+    with open(target_json_path, "w") as f:
+        json.dump(merged_dataset.model_dump(exclude_none=True), f)
+
+
+if __name__ == "__main__":
+    json1 = (
+        "/home/tlips/Documents/airo-mono/airo-dataset-tools/test/test_data/lego-battery-resized/annotations_train.json"
+    )
+    json2 = (
+        "/home/tlips/Documents/airo-mono/airo-dataset-tools/test/test_data/lego-battery-resized/annotations_val.json"
+    )
+    target_json = "/home/tlips/Documents/airo-mono/airo-dataset-tools/test/test_data/lego-battery-resize-merged/annotations_merged.json"
+    merge_coco_datasets(json1, json2, target_json)

airo_dataset_tools-2025.4.0/airo_dataset_tools/coco_tools/split_dataset.py

@@ -0,0 +1,112 @@
+""" Split a COCO dataset into subsets"""
+import json
+import pathlib
+import random
+from typing import List, Optional
+
+from airo_dataset_tools.data_parsers.coco import CocoImage, CocoInstanceAnnotation, CocoInstancesDataset
+from pydantic import ValidationError
+
+
+def split_coco_dataset(
+    coco_dataset: CocoInstancesDataset,
+    split_ratios: List[float],
+    shuffle_before_splitting: bool = True,
+    shuffle_seed: int = 42,
+) -> List[CocoInstancesDataset]:
+    """Split a COCO dataset into subsets by splitting the images according to the specified relative ratios.
+    All annotations for an image will be placed in the same subset as the image.
+
+    Note that this does not guarantee the ratio of the annotations OR an equal class balance in each subset.
+
+    Ratios must sum to 1.0.
+    """
+    ratio_sum = sum(split_ratios)
+    if abs(ratio_sum - 1.0) > 2e-2:
+        raise ValueError(f"Ratios must sum to 1.0. Ratios sum to {ratio_sum}.")
+
+    # split the images into 2 subsets (random or ordered)
+    images = coco_dataset.images
+
+    if shuffle_before_splitting:
+        # make shuffling reproducible
+        # use local random instance to not influence
+        # the global random state
+        rng = random.Random(shuffle_seed)
+        rng.shuffle(images)
+
+    image_splits: List[List[CocoImage]] = []
+    split_sizes = [round(ratio * len(images)) for ratio in split_ratios]
+    split_sizes[-1] = len(images) - sum(split_sizes[:-1])  # make sure the total number of images is correct
+    print(f"Split sizes: {split_sizes}")
+
+    for size in split_sizes:
+        image_splits.append(images[:size])
+        images = images[size:]
+
+    image_id_to_split_id = {}
+    for split_id, image_split in enumerate(image_splits):
+        for image in image_split:
+            image_id_to_split_id[image.id] = split_id
+
+    # gather the annotations for each subset
+    annotation_splits: List[List[CocoInstanceAnnotation]] = [[] for _ in range(len(split_ratios))]
+    for annotation in coco_dataset.annotations:
+        image_id = annotation.image_id
+        split_id = image_id_to_split_id[image_id]
+        # keep original image_ids and annotation_ids so that you could still reference the original dataset
+        annotation_splits[split_id].append(annotation)
+
+    # check if none of the annotation splits are empty:
+    for split_id, annotation_split in enumerate(annotation_splits):
+        if len(annotation_split) == 0:
+            raise ValueError(
+                f"Split {split_id} is empty, which is not allowed. Please use a larger dataset or a smaller split ratio."
+            )
+    # create a new COCO dataset for each subset
+    coco_dataset_splits: List[CocoInstancesDataset] = []
+    for annotation_split, image_split in zip(annotation_splits, image_splits):
+        coco_dataset_split = CocoInstancesDataset(
+            categories=coco_dataset.categories, images=image_split, annotations=annotation_split
+        )
+        coco_dataset_splits.append(coco_dataset_split)
+
+    return coco_dataset_splits
+
+
+def split_and_save_coco_dataset(
+    coco_json_path: str, split_ratios: List[float], shuffle_before_splitting: bool = True
+) -> None:
+    """Split a COCO dataset into subsets according to the specified relative ratios and save them to disk.
+    Images are split with their corresponding annotations. No guarantees on class balance or annotation ratios.
+
+    Ratios must sum to 1.0.
+
+    If two ratios are specified, the dataset will be split into two subsets. these will be called train/val by default.
+    If three ratios are specified, the dataset will be split into three subsets. these will be called train/val/test by default.
+    """
+    split_names = ["train", "val", "test"]
+    if len(split_ratios) > len(split_names):
+        raise ValueError(f"Only {len(split_names)} splits are supported. {len(split_ratios)} splits were specified.")
+
+    coco_dataset: Optional[CocoInstancesDataset] = None
+    with open(coco_json_path, "r") as f:
+        try:
+            coco_dataset = CocoInstancesDataset(**json.load(f))
+        except ValidationError as e:
+            print(e)
+            raise ValueError("Could not load CocoInstancesDataset")
+
+    coco_dataset_splits = split_coco_dataset(coco_dataset, split_ratios, shuffle_before_splitting)
+
+    for split_id, coco_dataset_split in enumerate(coco_dataset_splits):
+
+        file_name = coco_json_path.replace(".json", f"_{split_names[split_id]}.json")
+        with open(file_name, "w") as f:
+            json.dump(coco_dataset_split.model_dump(exclude_none=True), f)
+
+
+if __name__ == "__main__":
+    json_path = pathlib.Path(__file__).parents[2] / "test" / "test_data" / "person_keypoints_val2017_small.json"
+    print(json_path)
+    split_and_save_coco_dataset(str(json_path), [0.5, 0.5])