first commit

data/tusimple.py

@@ -0,0 +1,97 @@
+"""TuSimple (mask-based) dataset loader.
+
+The data on disk (`<root>/training/{frames,lane-masks}/`) is a segmentation-style
+export rather than the official `label_data_*.json` point annotations, and has no
+official val/test split (see project plan §3.2). This loader:
+  1. Pairs each frame with its mask by filename.
+  2. Extracts per-lane polylines from the mask via `utils.mask_to_lanes`.
+  3. Applies the letterbox/normalize/(flip) transform pipeline to produce
+     fixed-size training targets.
+  4. Carves a seeded train/val split from the single available folder.
+"""
+from __future__ import annotations
+
+import os
+import random
+from pathlib import Path
+
+import cv2
+import torch
+from torch.utils.data import Dataset
+
+from utils.mask_to_lanes import mask_to_lanes
+from data.transforms import build_transforms
+
+
+class TuSimpleMaskDataset(Dataset):
+    def __init__(
+        self,
+        root: str,
+        split: str,
+        out_w: int,
+        out_h: int,
+        max_lanes: int,
+        num_sample_ys: int,
+        val_fraction: float = 0.1,
+        seed: int = 42,
+    ):
+        assert split in ("train", "val")
+        self.root = Path(root)
+        frames_dir = self.root / "training" / "frames"
+        masks_dir = self.root / "training" / "lane-masks"
+
+        frame_files = sorted(f for f in os.listdir(frames_dir) if f.lower().endswith((".jpg", ".jpeg", ".png")))
+        pairs = [f for f in frame_files if (masks_dir / f).exists()]
+        if not pairs:
+            raise RuntimeError(f"No matching frame/mask pairs found under {self.root}")
+
+        rng = random.Random(seed)
+        shuffled = pairs[:]
+        rng.shuffle(shuffled)
+        n_val = max(1, int(len(shuffled) * val_fraction))
+        val_set = set(shuffled[:n_val])
+
+        if split == "val":
+            self.files = [f for f in pairs if f in val_set]
+        else:
+            self.files = [f for f in pairs if f not in val_set]
+
+        self.frames_dir = frames_dir
+        self.masks_dir = masks_dir
+        self.transforms = build_transforms(out_w, out_h, max_lanes, num_sample_ys, train=(split == "train"))
+
+    def __len__(self) -> int:
+        return len(self.files)
+
+    def __getitem__(self, idx: int) -> dict:
+        fname = self.files[idx]
+        image_bgr = cv2.imread(str(self.frames_dir / fname))
+        mask = cv2.imread(str(self.masks_dir / fname))
+        if image_bgr is None or mask is None:
+            raise RuntimeError(f"Failed to read {fname} from {self.frames_dir} / {self.masks_dir}")
+
+        image = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
+        lanes = mask_to_lanes(mask)
+
+        sample = {"image": image, "lanes": lanes, "name": fname}
+        sample = self.transforms(sample)
+        return sample
+
+
+def collate_fn(batch: list[dict]) -> dict:
+    images = torch.stack([b["image"] for b in batch], dim=0)
+    target_xs = torch.stack([b["target_xs"] for b in batch], dim=0)
+    target_valid_mask = torch.stack([b["target_valid_mask"] for b in batch], dim=0)
+    target_lane_valid = torch.stack([b["target_lane_valid"] for b in batch], dim=0)
+    target_endpoints = torch.stack([b["target_endpoints"] for b in batch], dim=0)
+    sample_ys = batch[0]["sample_ys"]
+    names = [b["name"] for b in batch]
+    return {
+        "images": images,
+        "target_xs": target_xs,
+        "target_valid_mask": target_valid_mask,
+        "target_lane_valid": target_lane_valid,
+        "target_endpoints": target_endpoints,
+        "sample_ys": sample_ys,
+        "names": names,
+    }