mapgs/data/real.py

"""Adapters for real driving datasets (§3.1).

These are *structured stubs*: the shared conversion :func:`raw_to_sample`
(build the Delaunay ground field, sample MAGT anchors, assemble the clip, apply
the TokenGS depth rescale) is fully implemented, so plugging in a real dataset
only requires implementing the dataset-specific I/O in ``load_raw`` — reading
images, intrinsics/extrinsics, the HD-map layers, and (optionally) LiDAR/boxes
into a :class:`RawClip`.

Per the design's integrity note, no dataset numbers are produced here; you must
run on real data yourself.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import List, Optional

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

from mapgs.config import MapGSConfig, MapConfig
from mapgs.data.types import MapGSSample
from mapgs.hdmap.ground_field import grid_field_from_points
from mapgs.hdmap.hdmap import HDMap
from mapgs.hdmap.anchors import sample_anchors


@dataclass
class RawClip:
    """What every real adapter must produce for one 2 s clip.

    Frames are indexed ``0..F-1`` at the dataset FPS; ``V`` cameras per frame.
    All poses share the map's world frame (UTM-like). Depth/mono/boxes optional.
    """

    images: torch.Tensor        # [F, V, 3, H, W] in [0,1]
    K: torch.Tensor             # [V, 3, 3]
    cam2world: torch.Tensor     # [F, V, 4, 4]
    ground_points: np.ndarray   # [G, 3] HD-map ground samples (world)
    lanes: List[torch.Tensor]   # lane polylines [Li, 3]
    boundaries: List[torch.Tensor] = field(default_factory=list)
    depth: Optional[torch.Tensor] = None     # [F, V, H, W] (LiDAR), -1 where unknown
    mono_depth: Optional[torch.Tensor] = None  # [F, V, H, W] frozen mono prior
    box_centers: Optional[torch.Tensor] = None  # [I, F, 3]
    box_rots: Optional[torch.Tensor] = None     # [I, F, 3, 3]
    box_size: Optional[torch.Tensor] = None     # [I, 3]
    canon_idx: Optional[torch.Tensor] = None    # [I]


def _context_frames(cfg: MapGSConfig, F: int) -> List[int]:
    return sorted({min(int(round(t * cfg.data.fps)), F - 1) for t in cfg.data.context_times})


def raw_to_sample(raw: RawClip, cfg: MapGSConfig, idx: int, n_sup_views: int = 6,
                  rescale_depth: bool = True) -> MapGSSample:
    """Convert a :class:`RawClip` into a :class:`MapGSSample`."""
    F, V = raw.cam2world.shape[:2]
    H, W = raw.images.shape[-2:]
    device = "cpu"

    # TokenGS: rescale translations so mean scene depth ~= 1 (improves conditioning)
    scale = 1.0
    if rescale_depth and raw.depth is not None:
        valid = raw.depth[raw.depth > 0]
        if valid.numel() > 0:
            scale = float(valid.mean())
    cam2world = raw.cam2world.clone()
    cam2world[..., :3, 3] /= scale
    ground_pts = raw.ground_points.copy()
    ground_pts[:, :] /= scale
    lanes = [l / scale for l in raw.lanes]
    boundaries = [b / scale for b in raw.boundaries]

    ground = grid_field_from_points(ground_pts, spacing=cfg.map.ground_spacing)
    hdmap = HDMap(ground=ground, lanes=lanes, boundaries=boundaries)
    mcfg = MapConfig(**{**cfg.map.__dict__, "n_anchors": cfg.model.tokens.n_map})
    anchors = sample_anchors(hdmap, mcfg, device="cpu", seed=idx)

    ctx = _context_frames(cfg, F)
    ctx_pairs = [(f, v) for f in ctx for v in range(V)]
    ctx_tids = torch.tensor([ti for ti, f in enumerate(ctx) for v in range(V)], dtype=torch.long)
    ci = torch.tensor([f for f, v in ctx_pairs]); cv = torch.tensor([v for f, v in ctx_pairs])

    ctx_set = set(ctx_pairs)
    all_pairs = [(f, v) for f in range(F) for v in range(V) if (f, v) not in ctx_set]
    g = torch.Generator().manual_seed(idx)
    perm = torch.randperm(len(all_pairs), generator=g)[:n_sup_views]
    sup_pairs = [all_pairs[i] for i in perm.tolist()]
    sf = torch.tensor([f for f, v in sup_pairs]); sv = torch.tensor([v for f, v in sup_pairs])

    depth = raw.depth if raw.depth is not None else torch.full((F, V, H, W), -1.0)
    mono = raw.mono_depth if raw.mono_depth is not None else depth.clamp_min(0.0)
    if rescale_depth and scale != 1.0:
        depth = torch.where(depth > 0, depth / scale, depth)
        mono = mono / scale

    has_dyn = raw.box_centers is not None and raw.box_centers.shape[0] > 0
    box_centers = (raw.box_centers / scale) if has_dyn else torch.zeros(0, F, 3)
    box_rots = raw.box_rots if has_dyn else torch.zeros(0, F, 3, 3)
    box_size = (raw.box_size / scale) if has_dyn else torch.zeros(0, 3)
    canon_idx = raw.canon_idx if has_dyn else torch.zeros(0, dtype=torch.long)

    return MapGSSample(
        ctx_images=raw.images[ci, cv], ctx_K=raw.K[cv], ctx_c2w=cam2world[ci, cv], ctx_tids=ctx_tids,
        sup_images=raw.images[sf, sv], sup_K=raw.K[sv], sup_c2w=cam2world[sf, sv],
        sup_depth=depth[sf, sv], sup_mono=mono[sf, sv], sup_frame=sf,
        anchor_pos=anchors.positions, anchor_type=anchors.types, anchor_normal=anchors.normals,
        ground=ground, lanes=lanes, boundaries=boundaries,
        box_centers=box_centers, box_rots=box_rots, box_size=box_size, canon_idx=canon_idx,
        scene_id=idx, scene_scale=scale,
    )


class BaseDrivingDataset(Dataset):
    """Implement :meth:`load_raw` and :attr:`clip_index` for a real dataset."""

    def __init__(self, cfg: MapGSConfig, split: str = "train", n_sup_views: int = 6):
        self.cfg = cfg
        self.split = split
        self.n_sup = n_sup_views
        self.clip_index: List = []  # populate in subclass: one entry per 2 s clip

    def load_raw(self, idx: int) -> RawClip:
        raise NotImplementedError

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

    def __getitem__(self, idx: int) -> MapGSSample:
        return raw_to_sample(self.load_raw(idx), self.cfg, idx, self.n_sup)


class WaymoDataset(BaseDrivingDataset):
    """Waymo Open v2 (primary, train-from-scratch). 3 front cameras.

    load_raw must read:
      * images + intrinsics/extrinsics for the 3 front cameras over a 2 s clip;
      * ground via LiDAR-aggregated / map elevation -> ``ground_points``;
      * lane polylines + road edges from the map proto -> ``lanes`` / ``boundaries``;
      * (optional) LiDAR range image -> projected sparse ``depth``;
      * (optional) labelled 3D boxes -> ``box_*``.
    """

    def load_raw(self, idx: int) -> RawClip:  # pragma: no cover - requires Waymo SDK + data
        raise NotImplementedError(
            "Implement Waymo v2 reading (waymo-open-dataset). See RawClip docstring "
            "for the required fields; 798 train / 202 val clips per §3.1."
        )


class NuScenesDataset(BaseDrivingDataset):
    """nuScenes (finetune-from-Waymo + zero-shot). Map-expansion drivable area + lanes;
    aggregated-LiDAR ground. 750 train / 150 val per §3.1."""

    def load_raw(self, idx: int) -> RawClip:  # pragma: no cover
        raise NotImplementedError(
            "Implement nuScenes reading (nuscenes-devkit + map-expansion)."
        )


class Argoverse2Dataset(BaseDrivingDataset):
    """Argoverse 2 (richest-map benchmark + MapNeRF continuity). Vector map with
    centerlines/boundaries; provided ground raster."""

    def load_raw(self, idx: int) -> RawClip:  # pragma: no cover
        raise NotImplementedError(
            "Implement AV2 reading (av2 devkit); use the ground-height raster directly."
        )