udacity_dataset.py

"""
Real dataset loader for the SullyChen "driving_dataset" (the common public
stand-in for Udacity-style steering-angle behavior cloning; the actual
Udacity releases are multi-GB ROS bags requiring rosbag tooling, which is
impractical for this exercise).

Data layout expected (as downloaded/extracted):

    <root>/data.txt          # lines of: "<frame_idx>.jpg <angle_degrees>"
    <root>/<frame_idx>.jpg   # 455x256 RGB dashcam frames, sequential in time

Two honest limitations vs. the original MAVNet paper/head design, both
handled explicitly rather than silently faked:

1. No "halt" ground truth exists in a moving-car dashcam dataset (there's no
   hover/stationary state), so movement labels are restricted to
   {forward, yaw_left, yaw_right}. The 4th class ("halt") will simply never
   appear as a target here - this is a real limitation of this dataset,
   not swept under the rug.
2. No junction ground truth exists either (that came from Cityscapes in the
   original paper, not Udacity). Rather than inventing a fake constant-zero
   junction label and silently training on it (which would trivially predict
   "no junction" and pollute reported metrics with a meaningless number),
   this dataset is flagged with `has_junction_labels = False` so the
   training loop can skip the junction loss/metric entirely when using it.
"""

import os
from dataclasses import dataclass
from typing import List, Tuple

import numpy as np
import torch
from PIL import Image
from torch.utils.data import Dataset

from dataset import route_based_split_three_way
from preprocessing import preprocess_frame

MOVEMENT_CLASSES = ["forward", "yaw_left", "yaw_right", "halt"]
FRAME_SIZE = (100, 100)


@dataclass
class UdacitySample:
    route_id: int
    image_path: str
    angle_deg: float
    movement_label: int


def discretize_angle(angle_deg: float, threshold_deg: float = 5.0) -> int:
    """Map a continuous steering-wheel angle (degrees) to a movement class.

    Positive angle = right turn, negative = left turn (SullyChen dataset
    convention). |angle| <= threshold is treated as "forward". Never
    returns the "halt" class (index 3) - see module docstring.
    """
    if angle_deg > threshold_deg:
        return MOVEMENT_CLASSES.index("yaw_right")
    if angle_deg < -threshold_deg:
        return MOVEMENT_CLASSES.index("yaw_left")
    return MOVEMENT_CLASSES.index("forward")


def flip_movement_label(movement_label: int) -> int:
    """Mirrors a movement label for horizontal-flip augmentation.

    yaw_left <-> yaw_right; forward/halt are symmetric and unchanged.
    """
    if movement_label == MOVEMENT_CLASSES.index("yaw_left"):
        return MOVEMENT_CLASSES.index("yaw_right")
    if movement_label == MOVEMENT_CLASSES.index("yaw_right"):
        return MOVEMENT_CLASSES.index("yaw_left")
    return movement_label


def parse_data_txt(root: str) -> List[Tuple[str, float]]:
    """Parses `data.txt`, returning [(image_path, angle_deg), ...] in the
    original (time-sequential) file order.
    """
    data_txt_path = os.path.join(root, "data.txt")
    if not os.path.isfile(data_txt_path):
        raise FileNotFoundError(
            f"Expected a data.txt file at {data_txt_path}. "
            "Download/extract the SullyChen driving_dataset there first."
        )

    entries = []
    with open(data_txt_path) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            # maxsplit=1: some releases append "angle,timestamp with spaces"
            # after the filename, so a plain .split() would over-split.
            filename, rest = line.split(maxsplit=1)
            angle_deg = float(rest.split(",")[0])
            entries.append((os.path.join(root, filename), angle_deg))
    return entries


class UdacityDrivingDataset(Dataset):
    """Real-image dataset built from the SullyChen driving_dataset."""

    has_junction_labels = False  # see module docstring

    def __init__(
        self,
        samples: List[UdacitySample],
        preprocessing_mode: str = "raw",
        augment: bool = False,
        seed: int = 0,
    ):
        self.samples = samples
        self.preprocessing_mode = preprocessing_mode
        self.augment = augment
        self._rng = np.random.default_rng(seed)

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        sample = self.samples[idx]
        with Image.open(sample.image_path) as img:
            img = img.convert("RGB").resize(FRAME_SIZE)
            frame = np.asarray(img)

        movement_label = sample.movement_label
        if self.augment and self._rng.random() < 0.5:
            # Random horizontal flip: mirrors the road, so the label must be
            # mirrored too (left turn <-> right turn). Only applied to the
            # training split (never val/test), and only for this dataset -
            # it's a cheap way to both double the effective data and reduce
            # the yaw_left/yaw_right class imbalance, since it's an artifact
            # of this particular drive's route, not a real asymmetry.
            frame = np.fliplr(frame).copy()
            movement_label = flip_movement_label(movement_label)

        processed = preprocess_frame(frame, mode=self.preprocessing_mode)
        frame_tensor = torch.from_numpy(processed).unsqueeze(0)  # (1, H, W)
        return {
            "frame": frame_tensor,
            "movement_label": torch.tensor(movement_label, dtype=torch.long),
            # Constant placeholder - never used in loss/metrics because
            # `has_junction_labels = False` tells the training loop to skip it.
            "junction_label": torch.tensor(0.0, dtype=torch.float32),
        }


def build_udacity_datasets(
    root: str,
    preprocessing_mode: str = "raw",
    angle_threshold_deg: float = 5.0,
    route_chunk_size: int = 500,
    val_fraction: float = 0.1,
    test_fraction: float = 0.2,
    augment_train: bool = True,
    seed: int = 0,
) -> Tuple[UdacityDrivingDataset, UdacityDrivingDataset, UdacityDrivingDataset]:
    """Builds train/val/test UdacityDrivingDataset splits.

    Since this dataset is one continuous, time-ordered drive (no separate
    recorded routes), we chunk it into contiguous blocks of
    `route_chunk_size` frames and treat each block as a pseudo-route. This
    lets us reuse `route_based_split_three_way` to hold out whole contiguous
    stretches of driving - never shuffling individual frames before
    splitting, which was the leakage bug in the original train.py.

    Only the training split gets flip augmentation (`augment_train`); val
    and test always reflect the real, unaugmented distribution.
    """
    entries = parse_data_txt(root)

    all_samples = []
    for frame_idx, (image_path, angle_deg) in enumerate(entries):
        route_id = frame_idx // route_chunk_size
        movement_label = discretize_angle(angle_deg, angle_threshold_deg)
        all_samples.append(UdacitySample(route_id, image_path, angle_deg, movement_label))

    route_ids = [s.route_id for s in all_samples]
    train_routes, val_routes, test_routes = route_based_split_three_way(
        route_ids, val_fraction, test_fraction, seed
    )
    train_routes, val_routes, test_routes = set(train_routes), set(val_routes), set(test_routes)

    train_samples = [s for s in all_samples if s.route_id in train_routes]
    val_samples = [s for s in all_samples if s.route_id in val_routes]
    test_samples = [s for s in all_samples if s.route_id in test_routes]

    train_ds = UdacityDrivingDataset(train_samples, preprocessing_mode, augment=augment_train, seed=seed)
    val_ds = UdacityDrivingDataset(val_samples, preprocessing_mode, augment=False)
    test_ds = UdacityDrivingDataset(test_samples, preprocessing_mode, augment=False)
    return train_ds, val_ds, test_ds


if __name__ == "__main__":
    import sys

    root = sys.argv[1] if len(sys.argv) > 1 else "../driving_dataset/driving_dataset"
    train_ds, val_ds, test_ds = build_udacity_datasets(root)
    print(f"train samples: {len(train_ds)}, val samples: {len(val_ds)}, test samples: {len(test_ds)}")
    item = train_ds[0]
    print("frame shape:", item["frame"].shape)
    print("movement_label:", item["movement_label"])