Prepare a Dataset

SUMMARY

Prepare a COCO-format dataset for Iris, then inspect its samples and batches before training. Each dataset split must provide images, annotations, bounding boxes, and category IDs; segmentation tasks also require a polygon or RLE mask for every annotated object.

Dataset Inspection Workflow

Start by pointing COCODataset at a dataset split. The inspection script verifies that Iris can read the images and annotations, and that the data loader produces matching image and target batches.

Resolve missing images, invalid boxes, unexpected labels, or absent masks before proceeding to training.

Create inspect_dataset.py

"""Load and inspect a COCO dataset."""

import argparse
from pathlib import Path

from loguru import logger
from torch.utils.data import DataLoader

from telekinesis.iris.dataset import COCODataset


def main(args: argparse.Namespace) -> None:
    """Load a COCO dataset and inspect individual and batched samples."""
    # ===================== Load Dataset ===================================
    dataset = COCODataset(
        dataset_dir=args.dataset_dir,
    )

    # ===================== Inspect Sample =================================
    logger.info(f"Dataset size: {len(dataset)}")

    image, target = dataset[0]
    logger.info("Single sample:")
    logger.info(f"  image shape: {image}")
    logger.info(f"  boxes: {target['boxes']}")
    logger.info(f"  labels: {target['labels']}")
    logger.info(f"  image_id: {target['image_id']}")

    def collate_fn(batch):
        """Group dataset samples into image and target tuples."""
        return tuple(zip(*batch))

    # ===================== Create Data Loader =============================
    dataloader = DataLoader(
        dataset,
        batch_size=4,
        shuffle=True,
        num_workers=0,
        collate_fn=collate_fn,
    )

    # ===================== Inspect Batch ==================================
    images, targets = next(iter(dataloader))

    logger.info("")
    logger.info("Batch:")
    logger.info(f"  images: {len(images)}")
    logger.info(f"  targets: {len(targets)}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--dataset-dir",
        type=Path,
        required=True,
        help="Path to the COCO dataset directory.",
    )
    main(parser.parse_args())

Run the Inspection

Pass the directory containing your COCO annotations and images:

python inspect_dataset.py --dataset-dir path/to/dataset

Segmentation annotations are required

When loading masks, every object must have a non-empty COCO segmentation. Iris raises an error rather than silently training without a mask.

Run the Example

A runnable version is available in the Telekinesis examples repository:

cd telekinesis-examples
python examples/iris/inspect_training_dataset.py \
  --dataset-dir path/to/dataset

Supported Layouts

Use the standard Iris layout:

dataset/train/
├── images/
│   ├── image-001.jpg
│   └── image-002.jpg
└── annotations.json

Iris also recognizes a Roboflow-style COCO split, where images and _annotations.coco.json are in the same directory:

dataset/train/
├── image-001.jpg
├── image-002.jpg
└── _annotations.coco.json

The annotation file must contain the usual COCO images, annotations, and categories arrays. Bounding boxes use COCO's [x, y, width, height] format in JSON; COCODataset converts them to xyxy tensors.

Dataset Reference

Each dataset item is (image, target). The image is an RGB float32 tensor in CHW layout and the [0, 1] range. The target contains:

Key	Meaning
boxes	Absolute xyxy boxes, shape (N, 4)
labels	COCO category IDs
image_id	COCO image ID
area	Object areas
iscrowd	COCO crowd flags
masks	Optional binary masks, shape (N, H, W)

Next Step

Continue with Train RFDETR or Train SAM3-LoRA.