commit a0aa03e9b7bd1b163cf5a60cecac7660bc940c45 Author: Suman Date: Tue Aug 18 18:50:32 2026 +0530 first commit diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/checkpoints/best.pt b/checkpoints/best.pt new file mode 100644 index 0000000..9764501 Binary files /dev/null and b/checkpoints/best.pt differ diff --git a/checkpoints/last.pt b/checkpoints/last.pt new file mode 100644 index 0000000..dcf2482 Binary files /dev/null and b/checkpoints/last.pt differ diff --git a/configs/default.yaml b/configs/default.yaml new file mode 100644 index 0000000..7c6ed74 --- /dev/null +++ b/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 diff --git a/data/__init__.py b/data/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/data/__pycache__/__init__.cpython-311.pyc b/data/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000..41301b8 Binary files /dev/null and b/data/__pycache__/__init__.cpython-311.pyc differ diff --git a/data/__pycache__/__init__.cpython-313.pyc b/data/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000..ba5cf7d Binary files /dev/null and b/data/__pycache__/__init__.cpython-313.pyc differ diff --git a/data/__pycache__/transforms.cpython-311.pyc b/data/__pycache__/transforms.cpython-311.pyc new file mode 100644 index 0000000..68eda0c Binary files /dev/null and b/data/__pycache__/transforms.cpython-311.pyc differ diff --git a/data/__pycache__/transforms.cpython-313.pyc b/data/__pycache__/transforms.cpython-313.pyc new file mode 100644 index 0000000..314b5fc Binary files /dev/null and b/data/__pycache__/transforms.cpython-313.pyc differ diff --git a/data/__pycache__/tusimple.cpython-311.pyc b/data/__pycache__/tusimple.cpython-311.pyc new file mode 100644 index 0000000..12f66aa Binary files /dev/null and b/data/__pycache__/tusimple.cpython-311.pyc differ diff --git a/data/__pycache__/tusimple.cpython-313.pyc b/data/__pycache__/tusimple.cpython-313.pyc new file mode 100644 index 0000000..cbddafd Binary files /dev/null and b/data/__pycache__/tusimple.cpython-313.pyc differ diff --git a/data/transforms.py b/data/transforms.py new file mode 100644 index 0000000..8cfe901 --- /dev/null +++ b/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) diff --git a/data/tusimple.py b/data/tusimple.py new file mode 100644 index 0000000..a36e41b --- /dev/null +++ b/data/tusimple.py @@ -0,0 +1,97 @@ +"""TuSimple (mask-based) dataset loader. + +The data on disk (`/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, + } diff --git a/engine/__init__.py b/engine/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/engine/__pycache__/__init__.cpython-311.pyc b/engine/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000..5bb51c4 Binary files /dev/null and b/engine/__pycache__/__init__.cpython-311.pyc differ diff --git a/engine/__pycache__/__init__.cpython-313.pyc b/engine/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000..22fac43 Binary files /dev/null and b/engine/__pycache__/__init__.cpython-313.pyc differ diff --git a/engine/__pycache__/evaluate.cpython-311.pyc b/engine/__pycache__/evaluate.cpython-311.pyc new file mode 100644 index 0000000..7504169 Binary files /dev/null and b/engine/__pycache__/evaluate.cpython-311.pyc differ diff --git a/engine/__pycache__/evaluate.cpython-313.pyc b/engine/__pycache__/evaluate.cpython-313.pyc new file mode 100644 index 0000000..cc9c0ac Binary files /dev/null and b/engine/__pycache__/evaluate.cpython-313.pyc differ diff --git a/engine/__pycache__/train.cpython-311.pyc b/engine/__pycache__/train.cpython-311.pyc new file mode 100644 index 0000000..02c3647 Binary files /dev/null and b/engine/__pycache__/train.cpython-311.pyc differ diff --git a/engine/__pycache__/train.cpython-313.pyc b/engine/__pycache__/train.cpython-313.pyc new file mode 100644 index 0000000..28488db Binary files /dev/null and b/engine/__pycache__/train.cpython-313.pyc differ diff --git a/engine/evaluate.py b/engine/evaluate.py new file mode 100644 index 0000000..3d085b3 --- /dev/null +++ b/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} diff --git a/engine/train.py b/engine/train.py new file mode 100644 index 0000000..f8d189f --- /dev/null +++ b/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) diff --git a/models/__init__.py b/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/models/__pycache__/__init__.cpython-311.pyc b/models/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000..908024b Binary files /dev/null and b/models/__pycache__/__init__.cpython-311.pyc differ diff --git a/models/__pycache__/__init__.cpython-313.pyc b/models/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000..6900053 Binary files /dev/null and b/models/__pycache__/__init__.cpython-313.pyc differ diff --git a/models/__pycache__/backbone.cpython-311.pyc b/models/__pycache__/backbone.cpython-311.pyc new file mode 100644 index 0000000..62d7abf Binary files /dev/null and b/models/__pycache__/backbone.cpython-311.pyc differ diff --git a/models/__pycache__/backbone.cpython-313.pyc b/models/__pycache__/backbone.cpython-313.pyc new file mode 100644 index 0000000..3952dba Binary files /dev/null and b/models/__pycache__/backbone.cpython-313.pyc differ diff --git a/models/__pycache__/head.cpython-311.pyc b/models/__pycache__/head.cpython-311.pyc new file mode 100644 index 0000000..e59891d Binary files /dev/null and b/models/__pycache__/head.cpython-311.pyc differ diff --git a/models/__pycache__/head.cpython-313.pyc b/models/__pycache__/head.cpython-313.pyc new file mode 100644 index 0000000..8e09eb3 Binary files /dev/null and b/models/__pycache__/head.cpython-313.pyc differ diff --git a/models/__pycache__/laneformer.cpython-311.pyc b/models/__pycache__/laneformer.cpython-311.pyc new file mode 100644 index 0000000..717bb85 Binary files /dev/null and b/models/__pycache__/laneformer.cpython-311.pyc differ diff --git a/models/__pycache__/laneformer.cpython-313.pyc b/models/__pycache__/laneformer.cpython-313.pyc new file mode 100644 index 0000000..30e16e3 Binary files /dev/null and b/models/__pycache__/laneformer.cpython-313.pyc differ diff --git a/models/__pycache__/losses.cpython-311.pyc b/models/__pycache__/losses.cpython-311.pyc new file mode 100644 index 0000000..df25a61 Binary files /dev/null and b/models/__pycache__/losses.cpython-311.pyc differ diff --git a/models/__pycache__/losses.cpython-313.pyc b/models/__pycache__/losses.cpython-313.pyc new file mode 100644 index 0000000..62d5536 Binary files /dev/null and b/models/__pycache__/losses.cpython-313.pyc differ diff --git a/models/__pycache__/matcher.cpython-311.pyc b/models/__pycache__/matcher.cpython-311.pyc new file mode 100644 index 0000000..abea0a7 Binary files /dev/null and b/models/__pycache__/matcher.cpython-311.pyc differ diff --git a/models/__pycache__/matcher.cpython-313.pyc b/models/__pycache__/matcher.cpython-313.pyc new file mode 100644 index 0000000..586752b Binary files /dev/null and b/models/__pycache__/matcher.cpython-313.pyc differ diff --git a/models/__pycache__/positional_encoding.cpython-311.pyc b/models/__pycache__/positional_encoding.cpython-311.pyc new file mode 100644 index 0000000..c6b32e1 Binary files /dev/null and b/models/__pycache__/positional_encoding.cpython-311.pyc differ diff --git a/models/__pycache__/positional_encoding.cpython-313.pyc b/models/__pycache__/positional_encoding.cpython-313.pyc new file mode 100644 index 0000000..9a186f5 Binary files /dev/null and b/models/__pycache__/positional_encoding.cpython-313.pyc differ diff --git a/models/__pycache__/transformer.cpython-311.pyc b/models/__pycache__/transformer.cpython-311.pyc new file mode 100644 index 0000000..c94c7b6 Binary files /dev/null and b/models/__pycache__/transformer.cpython-311.pyc differ diff --git a/models/__pycache__/transformer.cpython-313.pyc b/models/__pycache__/transformer.cpython-313.pyc new file mode 100644 index 0000000..b01a322 Binary files /dev/null and b/models/__pycache__/transformer.cpython-313.pyc differ diff --git a/models/backbone.py b/models/backbone.py new file mode 100644 index 0000000..642024d --- /dev/null +++ b/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) diff --git a/models/head.py b/models/head.py new file mode 100644 index 0000000..3480432 --- /dev/null +++ b/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 diff --git a/models/laneformer.py b/models/laneformer.py new file mode 100644 index 0000000..876b54e --- /dev/null +++ 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) diff --git a/models/losses.py b/models/losses.py new file mode 100644 index 0000000..99359e8 --- /dev/null +++ b/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), + } diff --git a/models/matcher.py b/models/matcher.py new file mode 100644 index 0000000..8b49d8c --- /dev/null +++ b/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 diff --git a/models/positional_encoding.py b/models/positional_encoding.py new file mode 100644 index 0000000..a936a5b --- /dev/null +++ b/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) diff --git a/models/transformer.py b/models/transformer.py new file mode 100644 index 0000000..4930bce --- /dev/null +++ b/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 diff --git a/scripts/train.py b/scripts/train.py new file mode 100644 index 0000000..07c8a97 --- /dev/null +++ b/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() diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_curve_utils.py b/tests/test_curve_utils.py new file mode 100644 index 0000000..7c887da --- /dev/null +++ b/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 == [] diff --git a/tests/test_matcher.py b/tests/test_matcher.py new file mode 100644 index 0000000..66f5ec3 --- /dev/null +++ b/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 diff --git a/tests/test_model_forward.py b/tests/test_model_forward.py new file mode 100644 index 0000000..812c7c6 --- /dev/null +++ b/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) diff --git a/train_run1.log b/train_run1.log new file mode 100644 index 0000000..531b3e4 --- /dev/null +++ b/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) diff --git a/train_run2_gpu.log b/train_run2_gpu.log new file mode 100644 index 0000000..096eedc --- /dev/null +++ b/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) diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/utils/__pycache__/__init__.cpython-311.pyc b/utils/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000..5d0ba94 Binary files /dev/null and b/utils/__pycache__/__init__.cpython-311.pyc differ diff --git a/utils/__pycache__/__init__.cpython-313.pyc b/utils/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000..94679ad Binary files /dev/null and b/utils/__pycache__/__init__.cpython-313.pyc differ diff --git a/utils/__pycache__/curve.cpython-311.pyc b/utils/__pycache__/curve.cpython-311.pyc new file mode 100644 index 0000000..f20344e Binary files /dev/null and b/utils/__pycache__/curve.cpython-311.pyc differ diff --git a/utils/__pycache__/curve.cpython-313.pyc b/utils/__pycache__/curve.cpython-313.pyc new file mode 100644 index 0000000..45c53a4 Binary files /dev/null and b/utils/__pycache__/curve.cpython-313.pyc differ diff --git a/utils/__pycache__/mask_to_lanes.cpython-311.pyc b/utils/__pycache__/mask_to_lanes.cpython-311.pyc new file mode 100644 index 0000000..832d6c6 Binary files /dev/null and b/utils/__pycache__/mask_to_lanes.cpython-311.pyc differ diff --git a/utils/__pycache__/mask_to_lanes.cpython-313.pyc b/utils/__pycache__/mask_to_lanes.cpython-313.pyc new file mode 100644 index 0000000..fa3fbaa Binary files /dev/null and b/utils/__pycache__/mask_to_lanes.cpython-313.pyc differ diff --git a/utils/curve.py b/utils/curve.py new file mode 100644 index 0000000..ea12f13 --- /dev/null +++ b/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) diff --git a/utils/mask_to_lanes.py b/utils/mask_to_lanes.py new file mode 100644 index 0000000..6cefc7e --- /dev/null +++ b/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