first commit

data/transforms.py

@@ -0,0 +1,141 @@
+"""Phase 1 transform pipeline: letterbox resize, normalize, horizontal flip.
+
+Camera-generalization augmentation (homography warps, photometric jitter) is
+explicitly deferred to Phase 2 per the project plan — this pipeline is the minimal
+set needed to prove the core model trains.
+
+All lane coordinates are carried as plain (x, y) pixel-space lists until the final
+`ToSampledTargets` step, which normalizes to [0,1] and resamples onto a fixed y-grid
+so batches of variable lane-count/length can be collated into fixed-size tensors.
+"""
+from __future__ import annotations
+
+import random
+
+import cv2
+import numpy as np
+import torch
+
+from utils.curve import compute_letterbox, sample_lane_at_ys
+
+IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+
+
+class Compose:
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, sample: dict) -> dict:
+        for t in self.transforms:
+            sample = t(sample)
+        return sample
+
+
+class LetterboxResize:
+    def __init__(self, out_w: int, out_h: int):
+        self.out_w = out_w
+        self.out_h = out_h
+
+    def __call__(self, sample: dict) -> dict:
+        image = sample["image"]
+        src_h, src_w = image.shape[:2]
+        lb = compute_letterbox(src_w, src_h, self.out_w, self.out_h)
+
+        new_w, new_h = int(round(src_w * lb.scale)), int(round(src_h * lb.scale))
+        resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
+
+        canvas = np.zeros((self.out_h, self.out_w, 3), dtype=image.dtype)
+        px, py = int(round(lb.pad_x)), int(round(lb.pad_y))
+        canvas[py:py + new_h, px:px + new_w] = resized
+
+        lanes = []
+        for lane in sample["lanes"]:
+            pts = np.array(lane, dtype=np.float32)
+            pts = lb.apply_points(pts)
+            lanes.append([(float(x), float(y)) for x, y in pts])
+
+        sample["image"] = canvas
+        sample["lanes"] = lanes
+        sample["letterbox"] = lb
+        return sample
+
+
+class RandomHorizontalFlip:
+    def __init__(self, p: float = 0.5):
+        self.p = p
+
+    def __call__(self, sample: dict) -> dict:
+        if random.random() >= self.p:
+            return sample
+        image = sample["image"]
+        w = image.shape[1]
+        sample["image"] = np.ascontiguousarray(image[:, ::-1, :])
+        sample["lanes"] = [[(w - 1 - x, y) for (x, y) in lane] for lane in sample["lanes"]]
+        return sample
+
+
+class Normalize:
+    """Uint8 HWC image -> normalized float32 CHW torch tensor."""
+
+    def __call__(self, sample: dict) -> dict:
+        image = sample["image"].astype(np.float32) / 255.0
+        image = (image - IMAGENET_MEAN) / IMAGENET_STD
+        sample["image"] = torch.from_numpy(image.transpose(2, 0, 1)).float()
+        return sample
+
+
+class ToSampledTargets:
+    """Resample variable-length lane polylines onto a fixed-size training target.
+
+    Produces, for a fixed grid of `num_sample_ys` normalized y-values shared by every
+    sample in a batch:
+      - xs:        (max_lanes, num_sample_ys) normalized x at each sample_y
+      - valid_mask:(max_lanes, num_sample_ys) bool, True where the lane is defined
+      - lane_valid:(max_lanes,) bool, True for real (non-padding) lane slots
+      - endpoints: (max_lanes, 2) normalized (y_start, y_end)
+    """
+
+    def __init__(self, out_w: int, out_h: int, max_lanes: int, num_sample_ys: int):
+        self.out_w = out_w
+        self.out_h = out_h
+        self.max_lanes = max_lanes
+        self.sample_ys = np.linspace(0.0, 1.0, num_sample_ys, dtype=np.float32)
+
+    def __call__(self, sample: dict) -> dict:
+        lanes = sample["lanes"][: self.max_lanes]
+        n = len(self.sample_ys)
+
+        xs = np.zeros((self.max_lanes, n), dtype=np.float32)
+        valid_mask = np.zeros((self.max_lanes, n), dtype=bool)
+        lane_valid = np.zeros((self.max_lanes,), dtype=bool)
+        endpoints = np.zeros((self.max_lanes, 2), dtype=np.float32)
+
+        for i, lane in enumerate(lanes):
+            if len(lane) < 2:
+                continue
+            norm_lane = [(x / self.out_w, y / self.out_h) for (x, y) in lane]
+            lane_xs, lane_valid_mask = sample_lane_at_ys(norm_lane, self.sample_ys)
+            if not lane_valid_mask.any():
+                continue
+            xs[i] = lane_xs
+            valid_mask[i] = lane_valid_mask
+            lane_valid[i] = True
+            ys_in_lane = self.sample_ys[lane_valid_mask]
+            endpoints[i] = [ys_in_lane.min(), ys_in_lane.max()]
+
+        sample["target_xs"] = torch.from_numpy(xs)
+        sample["target_valid_mask"] = torch.from_numpy(valid_mask)
+        sample["target_lane_valid"] = torch.from_numpy(lane_valid)
+        sample["target_endpoints"] = torch.from_numpy(endpoints)
+        sample["sample_ys"] = torch.from_numpy(self.sample_ys)
+        return sample
+
+
+def build_transforms(out_w: int, out_h: int, max_lanes: int, num_sample_ys: int, train: bool) -> Compose:
+    steps = [LetterboxResize(out_w, out_h)]
+    if train:
+        steps.append(RandomHorizontalFlip(p=0.5))
+    steps.append(ToSampledTargets(out_w, out_h, max_lanes, num_sample_ys))
+    steps.append(Normalize())
+    return Compose(steps)

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,
+    }