Dataset Utilities

A generation run produces COCO shards. These utilities consolidate them into one dataset, split it for training, and let you inspect the result before you train on it.

Import

from telekinesis.illusion.dataset import CocoShardMerger, DatasetConverter

---

CocoShardMerger

Merges the per-shard COCO files of a dataset into a single annotation file. BinPickingWorker.merge_shards() invokes it automatically.

merger = CocoShardMerger(dataset_dir)

The merger discovers every sub-directory containing a coco_annotations.json, validates that the categories (id and name pairs) match across shards, re-indexes image and annotation ids to be globally unique, rewrites each file_name to <shard_dirname>/images/<frame>.png, and writes merged_coco_annotations.json at the dataset root.

The per-shard files are left in place, so a merge never destroys the source data. Image paths inside the merged file resolve relative to the dataset root.

Categories must match

Shards generated from different specs, or from a spec whose category ids changed mid-run, cannot be merged. Keep id and category_name consistent across every dataset you intend to combine.

---

DatasetConverter

Reads merged_coco_annotations.json and writes a train, valid, and test split in COCO or YOLO format.

DatasetConverter(
    merged_dataset_dir,        # directory containing the merged annotations
    output_dir,                # destination for the split dataset
    dataset_format,            # "coco" | "yolo"
    ratios=(0.7, 0.2, 0.1),
    stratify=True,
    seed=42,
    merged_filename="merged_coco_annotations.json",
).split()

Parameters

Parameter	Type	Default	Description
merged_dataset_dir	str | Path	required	Directory containing the merged annotation file.
output_dir	str | Path	required	Destination for the split dataset. Created if missing.
dataset_format	str	required	"coco" or "yolo".
ratios	(float, float, float)	(0.7, 0.2, 0.1)	Train, valid, and test fractions. Must be non-negative and sum to 1.0.
stratify	bool	True	Stratify by each image's dominant annotation category. False shuffles randomly.
seed	int	42	RNG seed, for reproducible splits.
merged_filename	str	"merged_coco_annotations.json"	Override when the merged file uses a different name.

split() returns the output directory path.

When a stratum is too small to split at the requested ratio, the converter logs a warning and falls back to a random shuffle for that step, rather than failing the run.

COCO Output

RF-DETR-compatible, and readable by any COCO trainer:

<output_dir>/
├── train/
│   ├── _annotations.coco.json
│   ├── 000001.png
│   └── ...
├── valid/
│   ├── _annotations.coco.json
│   └── ...
└── test/
    ├── _annotations.coco.json
    └── ...

Images are renamed to {img_id:06d}.<ext>. Every split keeps the source info, licenses, and the full categories list, so all three share identical category definitions.

YOLO Output

Ultralytics YOLO-seg layout:

<output_dir>/
├── data.yaml
├── train/
│   ├── images/000001.png ...
│   └── labels/000001.txt ...
├── valid/
│   ├── images/ ...
│   └── labels/ ...
└── test/
    ├── images/ ...
    └── labels/ ...

Each label line is class x1 y1 x2 y2 ... xn yn, with coordinates normalized to [0, 1]. RLE segmentations are decoded and the largest external contour is kept; multi-polygon annotations also keep only their largest polygon.

data.yaml follows the Ultralytics convention (path, train, val, test, nc, names). Note that the key is val while the directory on disk is valid/.

Usage

from telekinesis.illusion.dataset import DatasetConverter

DatasetConverter(
    merged_dataset_dir="output/example_bin_picking_gearwheel_2",
    output_dir="output/example_bin_picking_gearwheel_2_split",
    dataset_format="yolo",
    ratios=(0.7, 0.2, 0.1),
    stratify=True,
    seed=42,
).split()

Setting output.dataset_format in a spec YAML runs exactly this step as part of merge_shards().

---

Viewer

Inspect a finished dataset in FiftyOne with the bundled script:

python examples/view_dataset.py

Set DATASET_DIR (and optionally DATASET_NAME) at the top of the script first. The script detects the layout, imports the dataset, prints per-split sample counts and a class-distribution table, then opens the FiftyOne app. FiftyOne is a project dependency, so no extra install is required.

Layout	Detected by	DATASET_DIR should be
merged	merged_coco_annotations.json at the root	The dataset root produced by merge_shards().
shard	coco_annotations.json plus images/	A single shard_* directory.
coco	_annotations.coco.json at the root or in train/valid/test	The DatasetConverter COCO output root, or one split.
yolo	data.yaml at the root	The DatasetConverter YOLO output root.

Detection precedence is merged, then shard, then yolo, then coco. The console table reports, per class, the count and share within each split, the total instances, the number of distinct images containing the class, and that class's share of the grand total. For interactive charts, add a Histograms panel next to the Samples tab in the app.

Quick previews

The lightweight tkinter viewer used by the quickstart examples is a preview tool only. Use FiftyOne for final dataset inspection.

Generate a dataset with a Worker

Run a spec-driven bin-picking dataset end to end: generate shards, merge them, and split into train, valid, and test.

Open tutorial →