98 lines
3.4 KiB
Python
98 lines
3.4 KiB
Python
"""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,
|
|
}
|