first commit

configs/default.yaml

@@ -0,0 +1,40 @@
+seed: 42
+
+data:
+  dataset: tusimple          # tusimple | culane | both
+  tusimple_root: /home/suman/Downloads/american_lanes/tusimple_preprocessed
+  val_fraction: 0.1          # carved from the single training/ split (no official val/test present)
+  input_height: 360
+  input_width: 640
+  max_lanes: 8               # N_lanes decoder queries
+  num_sample_ys: 48          # number of y-rows sampled per lane for loss/matching
+  batch_size: 8
+  num_workers: 4
+
+model:
+  backbone: resnet34
+  pretrained: true
+  fusion_channels: 128
+  encoder_layers: 4
+  decoder_layers: 1
+  attn_heads: 8
+  ffn_dim: 512
+  dropout: 0.1
+
+loss:
+  cls_weight: 2.0
+  reg_weight: 5.0
+  endpoint_weight: 1.0
+  bg_class_weight: 0.2       # down-weight the dominant "no lane" class
+
+train:
+  epochs: 50
+  lr_backbone: 1.0e-5
+  lr_new: 1.0e-4
+  weight_decay: 1.0e-4
+  warmup_steps: 300
+  grad_clip_norm: 1.0
+  amp: true
+  log_every: 20
+  checkpoint_dir: checkpoints
+  log_dir: runs

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

engine/evaluate.py

@@ -0,0 +1,96 @@
+"""TuSimple-style accuracy/FP/FN evaluation (paper Eq. 10).
+
+Simplification flagged explicitly: the official TuSimple metric uses a 25px
+tolerance defined in original-resolution pixel space with per-clip point sampling.
+We don't retain clip structure (our data is single frames, not clips) or the
+per-sample letterbox scale at collation time, so we evaluate in normalized [0,1]
+canonical (640-wide) space with an equivalent threshold (25/640 ~= 0.039). This is a
+consistent proxy metric for comparing our own checkpoints, not a pixel-exact
+reproduction of the official script -- consistent with the val-split caveat already
+flagged in the project plan (no official test_label.json available).
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from scipy.optimize import linear_sum_assignment
+
+from models.head import eval_curve
+
+X_THRESHOLD_NORM = 25.0 / 640.0
+CLS_PROB_THRESHOLD = 0.5
+
+
+@torch.no_grad()
+def evaluate(model, dataloader, device, sample_ys: torch.Tensor) -> dict:
+    model.eval()
+    total_correct_points = 0
+    total_gt_points = 0
+    total_fp_lanes = 0
+    total_fn_lanes = 0
+    total_pred_lanes = 0
+    total_gt_lanes = 0
+
+    for batch in dataloader:
+        images = batch["images"].to(device)
+        target_xs = batch["target_xs"].to(device)
+        target_valid_mask = batch["target_valid_mask"].to(device)
+        target_lane_valid = batch["target_lane_valid"].to(device)
+        sys_dev = sample_ys.to(device)
+
+        out = model(images)
+        probs = F.softmax(out["cls_logits"], dim=-1)[..., 1]  # (B, N)
+        pred_xs_all = eval_curve(out["curve_coeffs"], sys_dev.view(1, 1, -1).expand(*probs.shape, -1))  # (B, N, S)
+
+        B = images.shape[0]
+        for b in range(B):
+            pos_idx = (probs[b] > CLS_PROB_THRESHOLD).nonzero(as_tuple=True)[0]
+            gt_idx = target_lane_valid[b].nonzero(as_tuple=True)[0]
+
+            total_pred_lanes += pos_idx.numel()
+            total_gt_lanes += gt_idx.numel()
+
+            if gt_idx.numel() == 0:
+                total_fp_lanes += pos_idx.numel()
+                continue
+            if pos_idx.numel() == 0:
+                total_fn_lanes += gt_idx.numel()
+                gt_mask = target_valid_mask[b, gt_idx]
+                total_gt_points += int(gt_mask.sum().item())
+                continue
+
+            pred_xs = pred_xs_all[b, pos_idx]          # (P, S)
+            gt_xs = target_xs[b, gt_idx]                # (M, S)
+            gt_mask = target_valid_mask[b, gt_idx]      # (M, S)
+
+            diff = (pred_xs.unsqueeze(1) - gt_xs.unsqueeze(0)).abs()  # (P, M, S)
+            mask = gt_mask.unsqueeze(0).float()
+            denom = mask.sum(dim=-1).clamp(min=1.0)
+            cost = (diff * mask).sum(dim=-1) / denom    # (P, M)
+
+            pred_local, gt_local = linear_sum_assignment(cost.cpu().numpy())
+
+            matched_pred = set(pred_local.tolist())
+            matched_gt = set(gt_local.tolist())
+            total_fp_lanes += pos_idx.numel() - len(matched_pred)
+            total_fn_lanes += gt_idx.numel() - len(matched_gt)
+
+            for p_local, g_local in zip(pred_local, gt_local):
+                m = gt_mask[g_local]
+                if not m.any():
+                    continue
+                correct = (diff[p_local, g_local][m] < X_THRESHOLD_NORM).sum().item()
+                total_correct_points += correct
+                total_gt_points += int(m.sum().item())
+
+            unmatched_gt = set(range(gt_idx.numel())) - matched_gt
+            for g_local in unmatched_gt:
+                m = gt_mask[g_local]
+                total_gt_points += int(m.sum().item())
+
+    accuracy = total_correct_points / total_gt_points if total_gt_points > 0 else 0.0
+    fp = total_fp_lanes / total_pred_lanes if total_pred_lanes > 0 else 0.0
+    fn = total_fn_lanes / total_gt_lanes if total_gt_lanes > 0 else 0.0
+
+    model.train()
+    return {"accuracy": accuracy, "fp": fp, "fn": fn}

engine/train.py

@@ -0,0 +1,141 @@
+"""Phase 1 training loop: AdamW (discriminative LR), warmup+cosine, grad clipping,
+checkpointing, TuSimple accuracy/FP/FN eval per epoch. AMP/EMA are skipped for now
+since this runs on CPU (no usable CUDA on this machine -- see project notes); both
+are one-line additions once GPU training is available.
+"""
+from __future__ import annotations
+
+import math
+import os
+import time
+
+import torch
+from torch.utils.data import DataLoader
+
+from data.tusimple import TuSimpleMaskDataset, collate_fn
+from models.laneformer import LaneFormer
+from models.losses import compute_loss
+from engine.evaluate import evaluate
+
+
+def build_optimizer(model: LaneFormer, lr_backbone: float, lr_new: float, weight_decay: float) -> torch.optim.Optimizer:
+    backbone_params = list(model.backbone.parameters())
+    backbone_ids = {id(p) for p in backbone_params}
+    new_params = [p for p in model.parameters() if id(p) not in backbone_ids]
+    return torch.optim.AdamW(
+        [
+            {"params": backbone_params, "lr": lr_backbone},
+            {"params": new_params, "lr": lr_new},
+        ],
+        weight_decay=weight_decay,
+    )
+
+
+def build_scheduler(optimizer: torch.optim.Optimizer, warmup_steps: int, total_steps: int):
+    def lr_lambda(step: int) -> float:
+        if step < warmup_steps:
+            return step / max(1, warmup_steps)
+        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
+        return 0.5 * (1.0 + math.cos(math.pi * min(progress, 1.0)))
+
+    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+
+
+def train(cfg: dict) -> None:
+    torch.manual_seed(cfg["seed"])
+    torch.set_num_threads(cfg.get("num_threads", 12))
+    device = torch.device("cuda" if (cfg["train"].get("use_cuda") and torch.cuda.is_available()) else "cpu")
+    print(f"Using device: {device}", flush=True)
+
+    data_cfg = cfg["data"]
+    train_ds = TuSimpleMaskDataset(
+        root=data_cfg["tusimple_root"], split="train",
+        out_w=data_cfg["input_width"], out_h=data_cfg["input_height"],
+        max_lanes=data_cfg["max_lanes"], num_sample_ys=data_cfg["num_sample_ys"],
+        val_fraction=data_cfg["val_fraction"], seed=cfg["seed"],
+    )
+    val_ds = TuSimpleMaskDataset(
+        root=data_cfg["tusimple_root"], split="val",
+        out_w=data_cfg["input_width"], out_h=data_cfg["input_height"],
+        max_lanes=data_cfg["max_lanes"], num_sample_ys=data_cfg["num_sample_ys"],
+        val_fraction=data_cfg["val_fraction"], seed=cfg["seed"],
+    )
+    train_loader = DataLoader(
+        train_ds, batch_size=data_cfg["batch_size"], shuffle=True,
+        collate_fn=collate_fn, num_workers=data_cfg["num_workers"],
+    )
+    val_loader = DataLoader(
+        val_ds, batch_size=data_cfg["batch_size"], shuffle=False,
+        collate_fn=collate_fn, num_workers=data_cfg["num_workers"],
+    )
+    print(f"train={len(train_ds)} val={len(val_ds)} samples", flush=True)
+
+    model_cfg = cfg["model"]
+    model = LaneFormer(
+        backbone_name=model_cfg["backbone"], pretrained=model_cfg["pretrained"],
+        d_model=model_cfg["fusion_channels"], max_lanes=data_cfg["max_lanes"],
+        encoder_layers=model_cfg["encoder_layers"], decoder_layers=model_cfg["decoder_layers"],
+        nhead=model_cfg["attn_heads"], ffn_dim=model_cfg["ffn_dim"], dropout=model_cfg["dropout"],
+    ).to(device)
+
+    train_cfg = cfg["train"]
+    optimizer = build_optimizer(model, train_cfg["lr_backbone"], train_cfg["lr_new"], train_cfg["weight_decay"])
+    total_steps = train_cfg["epochs"] * len(train_loader)
+    scheduler = build_scheduler(optimizer, train_cfg["warmup_steps"], total_steps)
+
+    loss_cfg = cfg["loss"]
+    ckpt_dir = train_cfg["checkpoint_dir"]
+    os.makedirs(ckpt_dir, exist_ok=True)
+
+    best_acc = -1.0
+    global_step = 0
+
+    for epoch in range(train_cfg["epochs"]):
+        model.train()
+        epoch_start = time.time()
+        running = {"total": 0.0, "cls_loss": 0.0, "reg_loss": 0.0, "endpoint_loss": 0.0}
+        n_batches = 0
+
+        for batch in train_loader:
+            images = batch["images"].to(device)
+            targets = {k: batch[k].to(device) for k in
+                       ["target_xs", "target_valid_mask", "target_lane_valid", "target_endpoints"]}
+            sample_ys = batch["sample_ys"].to(device)
+
+            out = model(images)
+            loss_dict = compute_loss(
+                out, targets, sample_ys,
+                cls_weight=loss_cfg["cls_weight"], reg_weight=loss_cfg["reg_weight"],
+                endpoint_weight=loss_cfg["endpoint_weight"], bg_class_weight=loss_cfg["bg_class_weight"],
+            )
+
+            optimizer.zero_grad()
+            loss_dict["total"].backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), train_cfg["grad_clip_norm"])
+            optimizer.step()
+            scheduler.step()
+
+            for k in running:
+                running[k] += float(loss_dict[k].detach())
+            n_batches += 1
+            global_step += 1
+
+            if global_step % train_cfg["log_every"] == 0:
+                avg = {k: v / n_batches for k, v in running.items()}
+                lr = scheduler.get_last_lr()[-1]
+                print(f"epoch {epoch} step {global_step} lr {lr:.2e} "
+                      f"loss {avg['total']:.4f} (cls {avg['cls_loss']:.4f} "
+                      f"reg {avg['reg_loss']:.4f} ep {avg['endpoint_loss']:.4f})", flush=True)
+
+        epoch_time = time.time() - epoch_start
+        metrics = evaluate(model, val_loader, device, val_ds[0]["sample_ys"] if len(val_ds) else train_ds[0]["sample_ys"])
+        print(f"[epoch {epoch}] time={epoch_time/60:.1f}min val_acc={metrics['accuracy']:.4f} "
+              f"fp={metrics['fp']:.4f} fn={metrics['fn']:.4f}", flush=True)
+
+        ckpt_path = os.path.join(ckpt_dir, "last.pt")
+        torch.save({"model": model.state_dict(), "epoch": epoch, "metrics": metrics}, ckpt_path)
+        if metrics["accuracy"] > best_acc:
+            best_acc = metrics["accuracy"]
+            torch.save({"model": model.state_dict(), "epoch": epoch, "metrics": metrics},
+                       os.path.join(ckpt_dir, "best.pt"))
+            print(f"  new best (acc={best_acc:.4f}), saved best.pt", flush=True)

models/backbone.py

@@ -0,0 +1,76 @@
+"""ResNet34 backbone + 3-scale FPN-style fusion.
+
+Reproduces the paper's genuinely distinctive idea (multi-scale down/up-sample fusion
+to help extract thin, elongated lane structure) with a concretely-defined FPN instead
+of the paper's ambiguous "32-group ResNet32/ResNeXt50" description.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torchvision
+
+
+class ConvBNMish(nn.Module):
+    def __init__(self, in_ch: int, out_ch: int, kernel_size: int = 3, padding: int = 1):
+        super().__init__()
+        self.conv = nn.Conv2d(in_ch, out_ch, kernel_size, padding=padding, bias=False)
+        self.bn = nn.BatchNorm2d(out_ch)
+        self.act = nn.Mish(inplace=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.act(self.bn(self.conv(x)))
+
+
+class Backbone(nn.Module):
+    """Fuses stride 8/16/32 features (paper's 3-scale idea), then downsamples the
+    fused map to stride 32 before returning it -- full O(N^2) self-attention over a
+    stride-8 map (e.g. 45x80=3600 tokens) is prohibitively expensive/OOMs on CPU;
+    the paper's own tiny FLOPs count (0.425 GMACs) implies the same heavy
+    downsampling before attention, just left unstated. Stride 32 on a 360x640 input
+    gives an 11x20=220-token sequence, ~260x cheaper attention.
+    """
+
+    def __init__(self, name: str = "resnet34", pretrained: bool = True, out_channels: int = 128):
+        super().__init__()
+        weights = torchvision.models.ResNet34_Weights.IMAGENET1K_V1 if pretrained else None
+        net = torchvision.models.resnet34(weights=weights)
+
+        self.stem = nn.Sequential(net.conv1, net.bn1, net.relu, net.maxpool)
+        self.layer1 = net.layer1  # stride 4
+        self.layer2 = net.layer2  # stride 8,  C=128
+        self.layer3 = net.layer3  # stride 16, C=256
+        self.layer4 = net.layer4  # stride 32, C=512
+
+        # Shared extractor delta: project each scale to a common channel dim.
+        self.reduce_c3 = ConvBNMish(128, out_channels, kernel_size=1, padding=0)
+        self.reduce_c4 = ConvBNMish(256, out_channels, kernel_size=1, padding=0)
+        self.reduce_c5 = ConvBNMish(512, out_channels, kernel_size=1, padding=0)
+
+        self.fuse_c4 = ConvBNMish(out_channels, out_channels)
+        self.fuse_c3 = ConvBNMish(out_channels, out_channels)
+
+        # stride 8 -> stride 32 for the transformer input (see class docstring).
+        self.downsample = nn.Sequential(
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=2, padding=1),
+            nn.BatchNorm2d(out_channels), nn.Mish(inplace=True),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=2, padding=1),
+            nn.BatchNorm2d(out_channels), nn.Mish(inplace=True),
+        )
+
+        self.out_channels = out_channels
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.stem(x)
+        x = self.layer1(x)
+        c3 = self.layer2(x)      # stride 8
+        c4 = self.layer3(c3)     # stride 16
+        c5 = self.layer4(c4)     # stride 32
+
+        p5 = self.reduce_c5(c5)
+        p4 = self.reduce_c4(c4) + nn.functional.interpolate(p5, size=c4.shape[-2:], mode="nearest")
+        p4 = self.fuse_c4(p4)
+        p3 = self.reduce_c3(c3) + nn.functional.interpolate(p4, size=c3.shape[-2:], mode="nearest")
+        p3 = self.fuse_c3(p3)
+
+        return self.downsample(p3)  # (B, out_channels, H/32, W/32)

models/head.py

@@ -0,0 +1,52 @@
+"""Per-query prediction heads: lane/background classification + cubic curve regression.
+
+Phase 1 note: curve coefficients are regressed directly (unconstrained) rather than
+with LSTR's full numerical-reparameterization trick -- kept simple for the first
+trainable pass; revisit if training shows y^3-term gradient instability.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+class MLP(nn.Module):
+    def __init__(self, in_dim: int, hidden_dim: int, out_dim: int, num_layers: int = 2):
+        super().__init__()
+        dims = [in_dim] + [hidden_dim] * (num_layers - 1) + [out_dim]
+        layers = []
+        for i in range(len(dims) - 1):
+            layers.append(nn.Linear(dims[i], dims[i + 1]))
+            if i < len(dims) - 2:
+                layers.append(nn.ReLU(inplace=True))
+        self.net = nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+class PredictionHeads(nn.Module):
+    """x = k*y^3 + m*y^2 + n*y + b, plus (y_start, y_end), plus lane/background logits."""
+
+    def __init__(self, d_model: int, hidden_dim: int = 128):
+        super().__init__()
+        self.cls_head = MLP(d_model, hidden_dim, 2, num_layers=2)
+        self.curve_head = MLP(d_model, hidden_dim, 4, num_layers=3)     # k, m, n, b
+        self.endpoint_head = MLP(d_model, hidden_dim, 2, num_layers=2)  # y_start, y_end (pre-sigmoid)
+
+    def forward(self, queries: torch.Tensor) -> dict:
+        """queries: (B, N_lanes, d_model)"""
+        cls_logits = self.cls_head(queries)              # (B, N, 2)
+        curve_coeffs = self.curve_head(queries)           # (B, N, 4)
+        endpoints = torch.sigmoid(self.endpoint_head(queries))  # (B, N, 2) in [0,1]
+        return {"cls_logits": cls_logits, "curve_coeffs": curve_coeffs, "endpoints": endpoints}
+
+
+def eval_curve(curve_coeffs: torch.Tensor, ys: torch.Tensor) -> torch.Tensor:
+    """curve_coeffs: (..., 4) = [k,m,n,b]; ys: (...,S) or (S,) broadcastable normalized y.
+    Returns x(y) with shape (..., S).
+    """
+    k, m, n, b = curve_coeffs.unbind(dim=-1)  # each (...,)
+    k, m, n, b = k.unsqueeze(-1), m.unsqueeze(-1), n.unsqueeze(-1), b.unsqueeze(-1)
+    y = ys
+    return k * y**3 + m * y**2 + n * y + b

models/laneformer.py

@@ -0,0 +1,55 @@
+"""LaneFormer-CUSTOM Phase 1: backbone -> PE -> transformer encoder/decoder -> heads.
+
+Reasoning/verification module (feature-correction + confidence scoring) is Phase 2 --
+this assembly is the minimal architecture needed to prove the core trains.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from models.backbone import Backbone
+from models.positional_encoding import PositionEmbedding2D
+from models.transformer import Encoder, Decoder
+from models.head import PredictionHeads
+
+
+class LaneFormer(nn.Module):
+    def __init__(
+        self,
+        backbone_name: str = "resnet34",
+        pretrained: bool = True,
+        d_model: int = 128,
+        max_lanes: int = 8,
+        encoder_layers: int = 4,
+        decoder_layers: int = 1,
+        nhead: int = 8,
+        ffn_dim: int = 512,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+        self.backbone = Backbone(backbone_name, pretrained=pretrained, out_channels=d_model)
+        self.pos_embed = PositionEmbedding2D(d_model)
+        self.encoder = Encoder(encoder_layers, d_model, nhead, ffn_dim, dropout)
+        self.decoder = Decoder(decoder_layers, d_model, nhead, ffn_dim, dropout)
+        self.query_embed = nn.Embedding(max_lanes, d_model)
+        self.heads = PredictionHeads(d_model)
+        self.max_lanes = max_lanes
+        self.d_model = d_model
+
+    def forward(self, images: torch.Tensor) -> dict:
+        B = images.shape[0]
+        feat = self.backbone(images)                    # (B, C, H, W)
+        _, C, H, W = feat.shape
+
+        src = feat.flatten(2).permute(0, 2, 1)           # (B, H*W, C)
+        pos = self.pos_embed(H, W, images.device)        # (H*W, C)
+        pos = pos.unsqueeze(0).expand(B, -1, -1)          # (B, H*W, C)
+
+        memory = self.encoder(src, pos)                  # (B, H*W, C)
+
+        query_pos = self.query_embed.weight.unsqueeze(0).expand(B, -1, -1)  # (B, max_lanes, C)
+        tgt = torch.zeros_like(query_pos)
+        decoded = self.decoder(tgt, memory, query_pos, pos)  # (B, max_lanes, C)
+
+        return self.heads(decoded)

models/losses.py

@@ -0,0 +1,78 @@
+"""Loss terms: Hungarian-matched classification (cross-entropy) + curve regression
+(Smooth L1) + endpoint regression. Reasoning-module auxiliary loss is Phase 2.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from models.head import eval_curve
+from models.matcher import hungarian_match
+
+
+def compute_loss(
+    outputs: dict,
+    targets: dict,
+    sample_ys: torch.Tensor,
+    cls_weight: float = 2.0,
+    reg_weight: float = 5.0,
+    endpoint_weight: float = 1.0,
+    bg_class_weight: float = 0.2,
+) -> dict:
+    cls_logits = outputs["cls_logits"]        # (B, N, 2)
+    curve_coeffs = outputs["curve_coeffs"]     # (B, N, 4)
+    endpoints = outputs["endpoints"]           # (B, N, 2)
+
+    target_xs = targets["target_xs"]
+    target_valid_mask = targets["target_valid_mask"]
+    target_lane_valid = targets["target_lane_valid"]
+    target_endpoints = targets["target_endpoints"]
+
+    B, N, _ = cls_logits.shape
+    device = cls_logits.device
+
+    matches = hungarian_match(
+        cls_logits, curve_coeffs, target_xs, target_valid_mask, target_lane_valid,
+        sample_ys, cls_weight=cls_weight, reg_weight=reg_weight,
+    )
+
+    cls_targets = torch.zeros(B, N, dtype=torch.long, device=device)
+    reg_losses = []
+    endpoint_losses = []
+
+    for b, (pred_idx, gt_idx) in enumerate(matches):
+        if pred_idx.numel() == 0:
+            continue
+        pred_idx = pred_idx.to(device)
+        gt_idx = gt_idx.to(device)
+        cls_targets[b, pred_idx] = 1
+
+        matched_coeffs = curve_coeffs[b, pred_idx]                # (M, 4)
+        pred_xs = eval_curve(matched_coeffs, sample_ys.unsqueeze(0).expand(pred_idx.numel(), -1))  # (M, S)
+        gt_xs = target_xs[b, gt_idx]                               # (M, S)
+        gt_mask = target_valid_mask[b, gt_idx]                     # (M, S)
+
+        if gt_mask.any():
+            reg_loss = F.smooth_l1_loss(pred_xs[gt_mask], gt_xs[gt_mask], reduction="mean")
+            reg_losses.append(reg_loss)
+
+        pred_endpoints = endpoints[b, pred_idx]                    # (M, 2)
+        gt_endpoints = target_endpoints[b, gt_idx]                 # (M, 2)
+        endpoint_losses.append(F.smooth_l1_loss(pred_endpoints, gt_endpoints, reduction="mean"))
+
+    class_weights = torch.tensor([bg_class_weight, 1.0], device=device)
+    cls_loss = F.cross_entropy(cls_logits.reshape(-1, 2), cls_targets.reshape(-1), weight=class_weights)
+
+    reg_loss = torch.stack(reg_losses).mean() if reg_losses else torch.zeros((), device=device)
+    endpoint_loss = torch.stack(endpoint_losses).mean() if endpoint_losses else torch.zeros((), device=device)
+
+    total = cls_weight * cls_loss + reg_weight * reg_loss + endpoint_weight * endpoint_loss
+
+    return {
+        "total": total,
+        "cls_loss": cls_loss.detach(),
+        "reg_loss": reg_loss.detach(),
+        "endpoint_loss": endpoint_loss.detach(),
+        "num_matched": sum(p.numel() for p, _ in matches),
+    }

models/matcher.py

@@ -0,0 +1,59 @@
+"""Hungarian bipartite matching between predicted queries and GT lanes.
+
+Replaces the paper's vague "distance-based positive/negative" anchor rule. This is
+load-bearing for training (not an optional improvement): DETR-style architectures
+with fixed slot assignment typically fail to train due to the permutation problem.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from scipy.optimize import linear_sum_assignment
+
+from models.head import eval_curve
+
+
+@torch.no_grad()
+def hungarian_match(
+    cls_logits: torch.Tensor,       # (B, N, 2)
+    curve_coeffs: torch.Tensor,     # (B, N, 4)
+    target_xs: torch.Tensor,        # (B, N, S)
+    target_valid_mask: torch.Tensor,  # (B, N, S) bool
+    target_lane_valid: torch.Tensor,  # (B, N) bool
+    sample_ys: torch.Tensor,        # (S,)
+    cls_weight: float = 1.0,
+    reg_weight: float = 5.0,
+) -> list[tuple[torch.Tensor, torch.Tensor]]:
+    """Returns, per batch item, (pred_indices, gt_indices) of the matched pairs."""
+    B, N, _ = cls_logits.shape
+    probs = F.softmax(cls_logits, dim=-1)[..., 1]  # (B, N) prob of "lane"
+
+    pred_xs = eval_curve(curve_coeffs, sample_ys.view(1, 1, -1).expand(B, N, -1))  # (B, N, S)
+
+    results = []
+    for b in range(B):
+        valid_gt = target_lane_valid[b].nonzero(as_tuple=True)[0]  # (M,)
+        if valid_gt.numel() == 0:
+            results.append((torch.empty(0, dtype=torch.long), torch.empty(0, dtype=torch.long)))
+            continue
+
+        gt_xs = target_xs[b, valid_gt]            # (M, S)
+        gt_mask = target_valid_mask[b, valid_gt]  # (M, S)
+
+        # (N, M) cost matrix
+        diff = (pred_xs[b].unsqueeze(1) - gt_xs.unsqueeze(0)).abs()  # (N, M, S)
+        mask = gt_mask.unsqueeze(0).float()                          # (1, M, S)
+        denom = mask.sum(dim=-1).clamp(min=1.0)                      # (N broadcast, M)
+        reg_cost = (diff * mask).sum(dim=-1) / denom                 # (N, M)
+
+        cls_cost = -probs[b].unsqueeze(1).expand(-1, valid_gt.numel())  # (N, M)
+
+        cost = cls_weight * cls_cost + reg_weight * reg_cost
+        cost_np = cost.cpu().numpy()
+
+        pred_idx, gt_idx_local = linear_sum_assignment(cost_np)
+        gt_idx = valid_gt[gt_idx_local]
+
+        results.append((torch.as_tensor(pred_idx, dtype=torch.long),
+                         torch.as_tensor(gt_idx, dtype=torch.long)))
+    return results

models/positional_encoding.py

@@ -0,0 +1,35 @@
+"""Standard 2D sine-cosine positional encoding, normalized by grid size.
+
+Normalizing by grid size (rather than absolute pixel index, as the paper's Eq. 14
+ambiguously implies) keeps PE values independent of input resolution -- one of the
+camera/resolution-agnostic design choices in the project plan.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+class PositionEmbedding2D(nn.Module):
+    def __init__(self, dim: int, temperature: float = 10000.0):
+        super().__init__()
+        assert dim % 4 == 0, "dim must be divisible by 4 for 2D sine-cosine PE"
+        self.dim = dim
+        self.temperature = temperature
+
+    def forward(self, h: int, w: int, device: torch.device) -> torch.Tensor:
+        """Returns (h*w, dim) position embedding."""
+        num_pos_feats = self.dim // 2  # half the channels for x, half for y
+        y_embed = torch.linspace(0, 1, h, device=device).unsqueeze(1).repeat(1, w)   # (h, w)
+        x_embed = torch.linspace(0, 1, w, device=device).unsqueeze(0).repeat(h, 1)   # (h, w)
+
+        dim_t = torch.arange(num_pos_feats, device=device, dtype=torch.float32)
+        dim_t = self.temperature ** (2 * (dim_t // 2) / num_pos_feats)
+
+        pos_x = x_embed[..., None] / dim_t
+        pos_y = y_embed[..., None] / dim_t
+        pos_x = torch.stack((pos_x[..., 0::2].sin(), pos_x[..., 1::2].cos()), dim=-1).flatten(-2)
+        pos_y = torch.stack((pos_y[..., 0::2].sin(), pos_y[..., 1::2].cos()), dim=-1).flatten(-2)
+
+        pos = torch.cat((pos_y, pos_x), dim=-1)  # (h, w, dim)
+        return pos.flatten(0, 1)  # (h*w, dim)

models/transformer.py

@@ -0,0 +1,93 @@
+"""DETR-style transformer encoder/decoder.
+
+Positional encoding is injected into queries/keys at every attention call (not just
+added once at the input), following DETR -- the paper's own ablation found 4 encoder
+layers / 1 decoder layer to be the sweet spot before overfitting, which we reuse here.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+def _with_pos(x: torch.Tensor, pos: torch.Tensor | None) -> torch.Tensor:
+    return x if pos is None else x + pos
+
+
+class EncoderLayer(nn.Module):
+    def __init__(self, d_model: int, nhead: int, dim_feedforward: int, dropout: float):
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
+        self.ffn = nn.Sequential(
+            nn.Linear(d_model, dim_feedforward), nn.ReLU(inplace=True),
+            nn.Dropout(dropout), nn.Linear(dim_feedforward, d_model),
+        )
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.drop1 = nn.Dropout(dropout)
+        self.drop2 = nn.Dropout(dropout)
+
+    def forward(self, src: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
+        q = k = _with_pos(src, pos)
+        attn_out, _ = self.self_attn(q, k, src)
+        src = self.norm1(src + self.drop1(attn_out))
+        ffn_out = self.ffn(src)
+        src = self.norm2(src + self.drop2(ffn_out))
+        return src
+
+
+class Encoder(nn.Module):
+    def __init__(self, num_layers: int, d_model: int, nhead: int, dim_feedforward: int, dropout: float):
+        super().__init__()
+        self.layers = nn.ModuleList([
+            EncoderLayer(d_model, nhead, dim_feedforward, dropout) for _ in range(num_layers)
+        ])
+
+    def forward(self, src: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
+        for layer in self.layers:
+            src = layer(src, pos)
+        return src
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, d_model: int, nhead: int, dim_feedforward: int, dropout: float):
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
+        self.cross_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
+        self.ffn = nn.Sequential(
+            nn.Linear(d_model, dim_feedforward), nn.ReLU(inplace=True),
+            nn.Dropout(dropout), nn.Linear(dim_feedforward, d_model),
+        )
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.drop1 = nn.Dropout(dropout)
+        self.drop2 = nn.Dropout(dropout)
+        self.drop3 = nn.Dropout(dropout)
+
+    def forward(self, tgt: torch.Tensor, memory: torch.Tensor, query_pos: torch.Tensor, memory_pos: torch.Tensor) -> torch.Tensor:
+        q = k = _with_pos(tgt, query_pos)
+        attn_out, _ = self.self_attn(q, k, tgt)
+        tgt = self.norm1(tgt + self.drop1(attn_out))
+
+        attn_out, _ = self.cross_attn(
+            _with_pos(tgt, query_pos), _with_pos(memory, memory_pos), memory,
+        )
+        tgt = self.norm2(tgt + self.drop2(attn_out))
+
+        ffn_out = self.ffn(tgt)
+        tgt = self.norm3(tgt + self.drop3(ffn_out))
+        return tgt
+
+
+class Decoder(nn.Module):
+    def __init__(self, num_layers: int, d_model: int, nhead: int, dim_feedforward: int, dropout: float):
+        super().__init__()
+        self.layers = nn.ModuleList([
+            DecoderLayer(d_model, nhead, dim_feedforward, dropout) for _ in range(num_layers)
+        ])
+
+    def forward(self, tgt: torch.Tensor, memory: torch.Tensor, query_pos: torch.Tensor, memory_pos: torch.Tensor) -> torch.Tensor:
+        for layer in self.layers:
+            tgt = layer(tgt, memory, query_pos, memory_pos)
+        return tgt

scripts/train.py

@@ -0,0 +1,35 @@
+"""CLI entrypoint: python scripts/train.py [--config configs/default.yaml] [--epochs N]"""
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+
+import yaml
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from engine.train import train  # noqa: E402
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="configs/default.yaml")
+    parser.add_argument("--epochs", type=int, default=None, help="override train.epochs")
+    parser.add_argument("--num-threads", type=int, default=12)
+    parser.add_argument("--use-cuda", action="store_true", help="try CUDA if available")
+    args = parser.parse_args()
+
+    with open(args.config) as f:
+        cfg = yaml.safe_load(f)
+
+    if args.epochs is not None:
+        cfg["train"]["epochs"] = args.epochs
+    cfg["num_threads"] = args.num_threads
+    cfg["train"]["use_cuda"] = args.use_cuda
+
+    train(cfg)
+
+
+if __name__ == "__main__":
+    main()

tests/test_curve_utils.py

@@ -0,0 +1,45 @@
+import numpy as np
+
+from utils.curve import compute_letterbox, sample_lane_at_ys
+from utils.mask_to_lanes import mask_to_lanes
+
+
+def test_letterbox_roundtrip():
+    lb = compute_letterbox(src_w=1280, src_h=720, out_w=640, out_h=360)
+    pts = np.array([[100.0, 200.0], [1000.0, 500.0]], dtype=np.float32)
+    warped = lb.apply_points(pts)
+    restored = lb.invert_points(warped)
+    assert np.allclose(pts, restored, atol=1e-3)
+
+
+def test_letterbox_stays_in_bounds():
+    lb = compute_letterbox(src_w=1280, src_h=720, out_w=640, out_h=360)
+    corners = np.array([[0, 0], [1279, 719]], dtype=np.float32)
+    warped = lb.apply_points(corners)
+    assert (warped[:, 0] >= -1).all() and (warped[:, 0] <= 641).all()
+    assert (warped[:, 1] >= -1).all() and (warped[:, 1] <= 361).all()
+
+
+def test_sample_lane_at_ys_interpolates_and_masks():
+    lane = [(0.0, 0.0), (10.0, 1.0)]  # x = 10*y
+    sample_ys = np.array([0.0, 0.5, 1.0, 1.5], dtype=np.float32)
+    xs, valid = sample_lane_at_ys(lane, sample_ys)
+    assert valid.tolist() == [True, True, True, False]
+    assert np.allclose(xs[:3], [0.0, 5.0, 10.0], atol=1e-4)
+
+
+def test_mask_to_lanes_separates_instances():
+    mask = np.zeros((100, 200), dtype=np.uint8)
+    mask[:, 20:24] = 255   # vertical strip -> one lane instance
+    mask[:, 150:154] = 255  # another, far apart -> separate instance
+    lanes = mask_to_lanes(mask)
+    assert len(lanes) == 2
+    for lane in lanes:
+        assert len(lane) >= 5
+
+
+def test_mask_to_lanes_drops_small_noise_blobs():
+    mask = np.zeros((100, 200), dtype=np.uint8)
+    mask[10:12, 10:12] = 255  # tiny 2x2 JPEG-noise-like blob, below MIN_BLOB_AREA
+    lanes = mask_to_lanes(mask)
+    assert lanes == []

tests/test_matcher.py

@@ -0,0 +1,54 @@
+import torch
+
+from models.matcher import hungarian_match
+
+
+def test_matcher_recovers_perfect_assignment():
+    """If predicted curves exactly equal GT curves (in different query order),
+    the matcher should recover the correct pred<->gt pairing."""
+    torch.manual_seed(0)
+    B, N, S = 1, 4, 10
+    sample_ys = torch.linspace(0, 1, S)
+
+    gt_xs = torch.rand(B, N, S)
+    target_valid_mask = torch.ones(B, N, S, dtype=torch.bool)
+    target_lane_valid = torch.tensor([[True, True, False, False]])
+
+    # Build predictions as a *permuted* copy of the (valid) GT curves via exact
+    # per-point fit isn't needed here -- construct predicted curve coeffs that
+    # exactly reproduce gt_xs at the two valid GT rows, placed at query indices [2, 0].
+    curve_coeffs = torch.zeros(B, N, 4)
+    cls_logits = torch.zeros(B, N, 2)
+
+    for q_idx, gt_idx in [(2, 0), (0, 1)]:
+        # fit a cubic through the GT points exactly (deg-3 fit on 10 pts, small residual)
+        ys_np = sample_ys.numpy()
+        xs_np = gt_xs[0, gt_idx].numpy()
+        coeffs = torch.tensor([0.0, 0.0, 0.0, 0.0])
+        import numpy as np
+        fit = np.polyfit(ys_np, xs_np, deg=3)
+        coeffs = torch.tensor(fit, dtype=torch.float32)
+        curve_coeffs[0, q_idx] = coeffs
+        cls_logits[0, q_idx] = torch.tensor([-5.0, 5.0])  # confident "lane"
+
+    matches = hungarian_match(cls_logits, curve_coeffs, gt_xs, target_valid_mask,
+                               target_lane_valid, sample_ys)
+    pred_idx, gt_idx = matches[0]
+    pairs = set(zip(pred_idx.tolist(), gt_idx.tolist()))
+    assert (2, 0) in pairs
+    assert (0, 1) in pairs
+
+
+def test_matcher_handles_no_gt_lanes():
+    B, N, S = 1, 4, 10
+    sample_ys = torch.linspace(0, 1, S)
+    cls_logits = torch.zeros(B, N, 2)
+    curve_coeffs = torch.zeros(B, N, 4)
+    target_xs = torch.zeros(B, N, S)
+    target_valid_mask = torch.zeros(B, N, S, dtype=torch.bool)
+    target_lane_valid = torch.zeros(B, N, dtype=torch.bool)
+
+    matches = hungarian_match(cls_logits, curve_coeffs, target_xs, target_valid_mask,
+                               target_lane_valid, sample_ys)
+    pred_idx, gt_idx = matches[0]
+    assert pred_idx.numel() == 0 and gt_idx.numel() == 0

tests/test_model_forward.py

@@ -0,0 +1,27 @@
+import torch
+
+from models.laneformer import LaneFormer
+
+
+def test_forward_shapes():
+    torch.manual_seed(0)
+    model = LaneFormer(pretrained=False, max_lanes=8)
+    x = torch.randn(2, 3, 360, 640)
+    out = model(x)
+    assert out["cls_logits"].shape == (2, 8, 2)
+    assert out["curve_coeffs"].shape == (2, 8, 4)
+    assert out["endpoints"].shape == (2, 8, 2)
+    assert torch.isfinite(out["cls_logits"]).all()
+    assert torch.isfinite(out["curve_coeffs"]).all()
+    assert (out["endpoints"] >= 0).all() and (out["endpoints"] <= 1).all()
+
+
+def test_backward_runs():
+    torch.manual_seed(0)
+    model = LaneFormer(pretrained=False, max_lanes=8)
+    x = torch.randn(1, 3, 360, 640)
+    out = model(x)
+    loss = out["cls_logits"].sum() + out["curve_coeffs"].sum() + out["endpoints"].sum()
+    loss.backward()
+    grads = [p.grad for p in model.parameters() if p.requires_grad]
+    assert any(g is not None and torch.isfinite(g).all() for g in grads)

train_run1.log

@@ -0,0 +1,24 @@
+Using device: cpu
+train=3264 val=362 samples
+/media/suman/Backup_of_extra_/miniconda3/lib/python3.13/site-packages/torch/cuda/__init__.py:422: UserWarning: Found GPU0 NVIDIA GeForce GTX 1070 which is of compute capability (CC) 6.1.
+The following list shows the CCs this version of PyTorch was built for and the hardware CCs it supports:
+- 7.5 which supports hardware CC >=7.5,<8.0
+- 8.0 which supports hardware CC >=8.0,<9.0 except {8.7}
+- 8.6 which supports hardware CC >=8.6,<9.0 except {8.7}
+- 9.0 which supports hardware CC >=9.0,<10.0
+- 10.0 which supports hardware CC >=10.0,<11.0 except {10.1}
+- 12.0 which supports hardware CC >=12.0,<13.0
+Your installed torch==2.13.0+cu130 does not include kernels for this GPU. Reinstall the same version against a CUDA build that does, e.g.:
+  For CUDA 12.6 use pip install torch==2.13.0 --index-url https://download.pytorch.org/whl/cu126
+  _warn_unsupported_code(d, device_cc, code_ccs)
+/media/suman/Backup_of_extra_/miniconda3/lib/python3.13/site-packages/torch/cuda/__init__.py:540: UserWarning: 
+NVIDIA GeForce GTX 1070 with CUDA capability sm_61 is not compatible with the current PyTorch installation.
+The current PyTorch install supports CUDA capabilities sm_75 sm_80 sm_86 sm_90 sm_100 sm_120.
+If you want to use the NVIDIA GeForce GTX 1070 GPU with PyTorch, please check the instructions at https://pytorch.org/get-started/locally/
+
+  queued_call()
+epoch 0 step 20 lr 6.67e-06 loss 1.7232 (cls 0.7415 reg 0.0426 ep 0.0272)
+epoch 0 step 40 lr 1.33e-05 loss 1.4397 (cls 0.6244 reg 0.0330 ep 0.0261)
+epoch 0 step 60 lr 2.00e-05 loss 1.3305 (cls 0.5836 reg 0.0281 ep 0.0228)
+epoch 0 step 80 lr 2.67e-05 loss 1.2681 (cls 0.5636 reg 0.0242 ep 0.0199)
+epoch 0 step 100 lr 3.33e-05 loss 1.2162 (cls 0.5468 reg 0.0209 ep 0.0180)

train_run2_gpu.log

@@ -0,0 +1,33 @@
+Using device: cuda
+train=3264 val=362 samples
+epoch 0 step 20 lr 6.67e-06 loss 1.7249 (cls 0.7415 reg 0.0430 ep 0.0270)
+epoch 0 step 40 lr 1.33e-05 loss 1.4445 (cls 0.6270 reg 0.0330 ep 0.0257)
+epoch 0 step 60 lr 2.00e-05 loss 1.3320 (cls 0.5855 reg 0.0277 ep 0.0227)
+epoch 0 step 80 lr 2.67e-05 loss 1.2668 (cls 0.5647 reg 0.0235 ep 0.0199)
+epoch 0 step 100 lr 3.33e-05 loss 1.2138 (cls 0.5473 reg 0.0203 ep 0.0179)
+epoch 0 step 120 lr 4.00e-05 loss 1.1753 (cls 0.5346 reg 0.0179 ep 0.0167)
+epoch 0 step 140 lr 4.67e-05 loss 1.1205 (cls 0.5117 reg 0.0162 ep 0.0163)
+epoch 0 step 160 lr 5.33e-05 loss 1.0547 (cls 0.4813 reg 0.0152 ep 0.0161)
+epoch 0 step 180 lr 6.00e-05 loss 0.9954 (cls 0.4542 reg 0.0143 ep 0.0155)
+epoch 0 step 200 lr 6.67e-05 loss 0.9368 (cls 0.4272 reg 0.0135 ep 0.0149)
+epoch 0 step 220 lr 7.33e-05 loss 0.8772 (cls 0.3994 reg 0.0128 ep 0.0143)
+epoch 0 step 240 lr 8.00e-05 loss 0.8338 (cls 0.3789 reg 0.0124 ep 0.0139)
+epoch 0 step 260 lr 8.67e-05 loss 0.7896 (cls 0.3580 reg 0.0120 ep 0.0135)
+epoch 0 step 280 lr 9.33e-05 loss 0.7535 (cls 0.3410 reg 0.0117 ep 0.0131)
+epoch 0 step 300 lr 1.00e-04 loss 0.7211 (cls 0.3261 reg 0.0112 ep 0.0127)
+epoch 0 step 320 lr 1.00e-04 loss 0.6926 (cls 0.3128 reg 0.0109 ep 0.0123)
+epoch 0 step 340 lr 1.00e-04 loss 0.6691 (cls 0.3023 reg 0.0105 ep 0.0119)
+epoch 0 step 360 lr 1.00e-04 loss 0.6456 (cls 0.2914 reg 0.0103 ep 0.0114)
+epoch 0 step 380 lr 1.00e-04 loss 0.6255 (cls 0.2821 reg 0.0100 ep 0.0110)
+epoch 0 step 400 lr 1.00e-04 loss 0.6058 (cls 0.2732 reg 0.0098 ep 0.0106)
+[epoch 0] time=1.6min val_acc=0.3567 fp=0.0996 fn=0.0110
+  new best (acc=0.3567), saved best.pt
+epoch 1 step 420 lr 1.00e-04 loss 0.2317 (cls 0.1039 reg 0.0043 ep 0.0023)
+epoch 1 step 440 lr 1.00e-04 loss 0.2455 (cls 0.1115 reg 0.0041 ep 0.0021)
+epoch 1 step 460 lr 1.00e-04 loss 0.2396 (cls 0.1094 reg 0.0038 ep 0.0019)
+epoch 1 step 480 lr 1.00e-04 loss 0.2309 (cls 0.1051 reg 0.0038 ep 0.0019)
+epoch 1 step 500 lr 1.00e-04 loss 0.2224 (cls 0.1013 reg 0.0036 ep 0.0018)
+epoch 1 step 520 lr 1.00e-04 loss 0.2321 (cls 0.1058 reg 0.0037 ep 0.0017)
+epoch 1 step 540 lr 1.00e-04 loss 0.2246 (cls 0.1024 reg 0.0036 ep 0.0016)
+epoch 1 step 560 lr 1.00e-04 loss 0.2242 (cls 0.1020 reg 0.0037 ep 0.0017)
+epoch 1 step 580 lr 1.00e-04 loss 0.2265 (cls 0.1031 reg 0.0037 ep 0.0016)

utils/curve.py

@@ -0,0 +1,77 @@
+"""Letterbox resize + curve sampling utilities.
+
+Camera-agnostic design note: images of any source resolution/aspect ratio are
+letterboxed (aspect-preserving resize + pad) into a fixed network input size, and all
+lane coordinates are normalized to [0,1] within that canonical frame. `LetterboxTransform`
+carries enough state to map predictions back to original-image pixel coordinates.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+
+@dataclass
+class LetterboxTransform:
+    scale: float
+    pad_x: float
+    pad_y: float
+    out_w: int
+    out_h: int
+    src_w: int
+    src_h: int
+
+    def apply_points(self, points: np.ndarray) -> np.ndarray:
+        """points: (N,2) array of (x,y) in source-image pixel coords -> letterboxed pixel coords."""
+        out = points.copy().astype(np.float32)
+        out[:, 0] = out[:, 0] * self.scale + self.pad_x
+        out[:, 1] = out[:, 1] * self.scale + self.pad_y
+        return out
+
+    def invert_points(self, points: np.ndarray) -> np.ndarray:
+        """Inverse of apply_points: letterboxed pixel coords -> source-image pixel coords."""
+        out = points.copy().astype(np.float32)
+        out[:, 0] = (out[:, 0] - self.pad_x) / self.scale
+        out[:, 1] = (out[:, 1] - self.pad_y) / self.scale
+        return out
+
+
+def compute_letterbox(src_w: int, src_h: int, out_w: int, out_h: int) -> LetterboxTransform:
+    scale = min(out_w / src_w, out_h / src_h)
+    new_w, new_h = src_w * scale, src_h * scale
+    pad_x = (out_w - new_w) / 2.0
+    pad_y = (out_h - new_h) / 2.0
+    return LetterboxTransform(scale=scale, pad_x=pad_x, pad_y=pad_y,
+                               out_w=out_w, out_h=out_h, src_w=src_w, src_h=src_h)
+
+
+def sample_lane_at_ys(lane_points: list[tuple[float, float]], sample_ys: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    """Linearly interpolate a polyline's x(y) at the given y-values.
+
+    Returns (x_values, valid_mask) — valid_mask is False outside the polyline's
+    observed y-range (no extrapolation).
+    """
+    pts = np.array(sorted(lane_points, key=lambda p: p[1]), dtype=np.float32)
+    ys_src = pts[:, 1]
+    xs_src = pts[:, 0]
+
+    y_min, y_max = ys_src.min(), ys_src.max()
+    valid = (sample_ys >= y_min) & (sample_ys <= y_max)
+
+    xs_out = np.zeros_like(sample_ys, dtype=np.float32)
+    if valid.any():
+        xs_out[valid] = np.interp(sample_ys[valid], ys_src, xs_src)
+    return xs_out, valid
+
+
+def polyfit_cubic(lane_points: list[tuple[float, float]]) -> np.ndarray:
+    """Least-squares cubic fit x = k*y^3 + m*y^2 + n*y + b (normalized coords expected).
+
+    Used only as a reference/visualization utility, not inside the training loop
+    (the model regresses these coefficients directly).
+    """
+    pts = np.array(lane_points, dtype=np.float64)
+    ys, xs = pts[:, 1], pts[:, 0]
+    coeffs = np.polyfit(ys, xs, deg=3)  # k, m, n, b
+    return coeffs.astype(np.float32)

utils/mask_to_lanes.py

@@ -0,0 +1,51 @@
+"""Extract per-lane polylines from binary (JPEG-compressed) lane-segmentation masks.
+
+The dataset at hand stores lanes as a single-class binary mask (white=lane,
+black=background) rather than TuSimple's official per-lane (x,y) JSON annotations, so
+lane *instances* have to be recovered from connected components of the thresholded
+mask. Masks are lossy JPEG, so we clean small compression-noise blobs before labeling.
+"""
+from __future__ import annotations
+
+import cv2
+import numpy as np
+
+MIN_BLOB_AREA = 150          # drop connected components smaller than this (JPEG noise)
+BINARY_THRESHOLD = 127
+MORPH_KERNEL = np.ones((3, 3), np.uint8)
+MIN_POINTS_PER_LANE = 5       # discard components too short to be a real lane
+
+
+def mask_to_lanes(mask: np.ndarray) -> list[list[tuple[float, float]]]:
+    """Convert an (H,W) or (H,W,3) binary-ish mask into a list of lane polylines.
+
+    Each polyline is a list of (x, y) pixel points, one point per mask row that the
+    lane instance covers, ordered top-to-bottom (increasing y).
+    """
+    if mask.ndim == 3:
+        mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
+
+    binary = (mask > BINARY_THRESHOLD).astype(np.uint8)
+    binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN, MORPH_KERNEL)
+
+    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(binary, connectivity=8)
+
+    lanes: list[list[tuple[float, float]]] = []
+    for label_id in range(1, num_labels):  # 0 is background
+        area = stats[label_id, cv2.CC_STAT_AREA]
+        if area < MIN_BLOB_AREA:
+            continue
+
+        ys, xs = np.where(labels == label_id)
+        if ys.size == 0:
+            continue
+
+        row_to_xs: dict[int, list[int]] = {}
+        for y, x in zip(ys, xs):
+            row_to_xs.setdefault(int(y), []).append(int(x))
+
+        points = [(float(np.mean(row_to_xs[y])), float(y)) for y in sorted(row_to_xs)]
+        if len(points) >= MIN_POINTS_PER_LANE:
+            lanes.append(points)
+
+    return lanes