142 lines
5.1 KiB
Python
142 lines
5.1 KiB
Python
"""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)
|