Applies to v0.1.0.

Labels & training

The masks on disk store class colours. For a loss function you need integer class indices. seg_labels.py converts between them losslessly.

The snippets below import seg_labels and use PyTorch, so run them from examples/collect_dataset/ inside your virtual environment (see Environment).

The class legend

Index = position in the list (road = 0 … background = 9). Colours are the sim's actual render. Note tunnel / lidar are 127, not 128 (a float→byte truncation the legend already accounts for).

0  road         128,  64, 128     5  obstacle     135, 206, 250
1  lane           0, 255,   0     6  adboard       69,  69,  69
2  stop         255,   0,   0     7  tunnel       255, 127,   0
3  lidar        127, 127, 127     8  nondrivable  244,  35, 232
4  vehicle        0,   0, 142     9  background     0,   0,   0

This legend is the single source of truth for colour ↔ index and must match the sim's SceneConfig.segCategories. It lives in both collect_seg_dataset.py (CLASS_LEGEND → classes.json) and seg_labels.py (CLASSES). If you recolour a class in Unity, update both.

Colour → index

rgb_to_index packs each pixel's RGB into a 24-bit key and looks it up in the legend, a fast, vectorised mapping. It returns the index mask and a count of unmapped pixels:

import numpy as np
from PIL import Image
from seg_labels import rgb_to_index, index_to_rgb, NUM_CLASSES

mask_rgb = np.array(Image.open("dataset/mask/frame_000001.png").convert("RGB"))
label, n_unmapped = rgb_to_index(mask_rgb)   # label: HxW uint8 in [0, NUM_CLASSES)
assert n_unmapped == 0                        # 0 on a correct flat / non-AA render

Image.fromarray(index_to_rgb(label)).save("check.png")   # round-trip to colour

n_unmapped > 0 is your drift alarm: a colour in the mask isn't in the legend. A class was added or recoloured in Unity, or the seg camera is antialiasing.

Load it in PyTorch

Convert colour → index lazily inside the Dataset, so the labels on disk stay lossless, human-viewable colour PNGs:

import json, os
import numpy as np
from PIL import Image
from torch.utils.data import Dataset
from seg_labels import rgb_to_index

class NanoSimSegDataset(Dataset):
    def __init__(self, root, transform=None):
        self.root = root
        with open(os.path.join(root, "manifest.jsonl")) as f:
            self.items = [json.loads(line) for line in f]
        self.transform = transform

    def __len__(self):
        return len(self.items)

    def __getitem__(self, i):
        rec  = self.items[i]
        rgb  = np.array(Image.open(os.path.join(self.root, rec["rgb"])).convert("RGB"))
        mask = np.array(Image.open(os.path.join(self.root, rec["mask"])).convert("RGB"))
        label, _ = rgb_to_index(mask)
        if self.transform:
            rgb, label = self.transform(rgb, label)
        return rgb, label

Getting a dataset worth training on

• Decorrelate frames. Consecutive frames on a slow lap are near-identical. Use --stride 3–5 and several --loops.
• Cover the track. The route graph has 12 valid branch combinations; multiple laps let the expert take different branches. Vary --steer-noise for off-centre views the perfect expert never produces on its own.
• Mind class balance. background and road dominate; stop / vehicle are sparse, so weight the loss (inverse-frequency) or oversample rare frames.
• Hold out a validation split by lap or stemrange, so val frames aren't near-duplicates of train frames.
• Validate before training. Run validate_dataset.py on the output first. It catches count mismatches and non-lossless masks without the sim.

Notes & gotchas

• Don't antialias the seg camera.The lossless colour→index mapping needs a flat, non-AA render. AA'd edges blend two class colours into a third the legend doesn't know.
• Manifest stays small — only pose + route metadata go into manifest.jsonl; the images live on disk.