first commit

data/transforms.py

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