mapgs/data/unified.py

"""Unified on-disk clip format — the single representation the network trains on.

Every source dataset (Waymo, Argoverse 2, synthetic) is converted into this
schema so the model never sees dataset-specific quirks and Waymo + AV2 can be
trained jointly (just point :class:`UnifiedClipDataset` at multiple roots).

On-disk layout (one directory per 2 s clip)::

    <root>/<clip_id>/
        meta.pt                     # poses, intrinsics, map grid, polylines, anchors, boxes
        images/<f:03d>_<v>.png      # native-resolution RGB, frame f, camera v
        depth/<f:03d>_<v>.npy       # optional metric depth (LiDAR), float16

``meta.pt`` keys (tensors unless noted):
    version:str source:str fps:int num_frames:F num_cameras:V cameras:[str]
    K:[V,3,3]  cam2world:[F,V,4,4]   image_hw:(H0,W0)            # native res, world frame
    ground_heights:[Hy,Wx] ground_meta:(x0,y0,dx,dy)            # GridGroundField
    lanes:[ [Li,3] ]  boundaries:[ [Bi,3] ]
    anchor_pos:[Na,3] anchor_type:[Na] anchor_normal:[Na,3]     # precomputed MAGT anchors
    box_centers:[I,F,3] box_rots:[I,F,3,3] box_size:[I,3] canon_idx:[I]
    has_depth:bool  scene_scale:float                           # TokenGS mean-depth rescale

Images are stored at native resolution; the loader resizes to ``cfg.data.{height,
width}`` with intrinsic scaling, so the *same* unified data serves both training
stages (§3.2 multi-res). ``scene_scale`` is stored (not pre-applied) and divided
out at load time.
"""

from __future__ import annotations

import json
import os
from typing import List, Optional, Union

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.geometry.cameras import resize_with_intrinsics
from mapgs.hdmap.anchors import sample_anchors, AnchorSet
from mapgs.hdmap.ground_field import GridGroundField
from mapgs.hdmap.hdmap import HDMap

UNIFIED_VERSION = "1.0"


# --------------------------------------------------------------------------- #
# Writer (used by every converter)
# --------------------------------------------------------------------------- #
def write_unified_clip(
    root: str,
    clip_id: str,
    source: str,
    fps: int,
    images: torch.Tensor,        # [F, V, 3, H0, W0] in [0,1]
    K: torch.Tensor,             # [V, 3, 3]  (native res)
    cam2world: torch.Tensor,     # [F, V, 4, 4]  (world/city frame)
    ground: GridGroundField,
    lanes: List[torch.Tensor],
    boundaries: List[torch.Tensor],
    anchors: AnchorSet,
    boxes: Optional[dict] = None,     # {box_centers,box_rots,box_size,canon_idx}
    depth: Optional[torch.Tensor] = None,   # [F, V, H0, W0] metric, <=0 unknown
    scene_scale: float = 1.0,
    cameras: Optional[List[str]] = None,
    image_ext: str = "png",                 # "png" (lossless) or "jpg" (q95, ~3.7x smaller)
) -> str:
    import imageio.v2 as imageio

    F, V = images.shape[:2]
    H0, W0 = images.shape[-2:]
    clip_dir = os.path.join(root, clip_id)
    os.makedirs(os.path.join(clip_dir, "images"), exist_ok=True)
    has_depth = depth is not None
    if has_depth:
        os.makedirs(os.path.join(clip_dir, "depth"), exist_ok=True)

    jpg_kw = {"quality": 95} if image_ext == "jpg" else {}
    for f in range(F):
        for v in range(V):
            arr = (images[f, v].clamp(0, 1).permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
            imageio.imwrite(os.path.join(clip_dir, "images", f"{f:03d}_{v}.{image_ext}"), arr, **jpg_kw)
            if has_depth:
                np.save(os.path.join(clip_dir, "depth", f"{f:03d}_{v}.npy"),
                        depth[f, v].cpu().numpy().astype(np.float16))

    boxes = boxes or {}
    meta = {
        "version": UNIFIED_VERSION, "source": source, "clip_id": clip_id,
        "fps": int(fps), "num_frames": int(F), "num_cameras": int(V),
        "cameras": cameras or [f"cam{v}" for v in range(V)],
        "K": K.float().cpu(), "cam2world": cam2world.float().cpu(), "image_hw": (int(H0), int(W0)),
        "ground_heights": ground.heights.float().cpu(),
        "ground_meta": (float(ground.x0), float(ground.y0), float(ground.dx), float(ground.dy)),
        "lanes": [l.float().cpu() for l in lanes],
        "boundaries": [b.float().cpu() for b in boundaries],
        "anchor_pos": anchors.positions.float().cpu(),
        "anchor_type": anchors.types.cpu(),
        "anchor_normal": anchors.normals.float().cpu(),
        "box_centers": boxes.get("box_centers", torch.zeros(0, F, 3)).float().cpu(),
        "box_rots": boxes.get("box_rots", torch.zeros(0, F, 3, 3)).float().cpu(),
        "box_size": boxes.get("box_size", torch.zeros(0, 3)).float().cpu(),
        "canon_idx": boxes.get("canon_idx", torch.zeros(0, dtype=torch.long)).cpu(),
        "has_depth": bool(has_depth), "scene_scale": float(scene_scale),
    }
    torch.save(meta, os.path.join(clip_dir, "meta.pt"))
    return clip_dir


# --------------------------------------------------------------------------- #
# Dataset
# --------------------------------------------------------------------------- #
def _discover_clips(roots: List[str]) -> List[str]:
    clips: List[str] = []
    for root in roots:
        if not os.path.isdir(root):
            continue
        for name in sorted(os.listdir(root)):
            d = os.path.join(root, name)
            if os.path.isfile(os.path.join(d, "meta.pt")):
                clips.append(d)
    return clips


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


class UnifiedClipDataset(Dataset):
    """Reads unified clips from one or more roots (mixed sources -> joint training)."""

    def __init__(self, cfg: MapGSConfig, roots: Union[str, List[str]], split: str = "train",
                 n_sup_views: int = 6):
        self.cfg = cfg
        self.H, self.W = cfg.data.height, cfg.data.width
        self.n_sup = n_sup_views
        if isinstance(roots, str):
            roots = [roots]
        # allow either <root> or <root>/<split>
        expanded = []
        for r in roots:
            rs = os.path.join(r, split)
            expanded.append(rs if os.path.isdir(rs) else r)
        self.clips = _discover_clips(expanded)
        if not self.clips:
            raise FileNotFoundError(f"no unified clips under {expanded}")
        self._mcfg = MapConfig(**{**cfg.map.__dict__, "n_anchors": cfg.model.tokens.n_map})

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

    def _load_image(self, clip_dir, f, v, K):
        import imageio.v2 as imageio
        base = os.path.join(clip_dir, "images", f"{f:03d}_{v}")
        path = base + ".png" if os.path.exists(base + ".png") else base + ".jpg"
        arr = imageio.imread(path)
        img = torch.from_numpy(np.asarray(arr)).float().permute(2, 0, 1) / 255.0
        if getattr(self, "letterbox", False):
            from mapgs.geometry.cameras import letterbox_with_intrinsics
            return letterbox_with_intrinsics(img, K, self.H, self.W)
        return resize_with_intrinsics(img, K, self.H, self.W)

    def _load_depth(self, clip_dir, f, v):
        path = os.path.join(clip_dir, "depth", f"{f:03d}_{v}.npy")
        if not os.path.exists(path):
            return torch.full((self.H, self.W), -1.0)
        d = torch.from_numpy(np.load(path).astype(np.float32))[None, None]
        d = torch.nn.functional.interpolate(d, size=(self.H, self.W), mode="nearest")[0, 0]
        return d

    def __getitem__(self, idx: int) -> MapGSSample:
        clip_dir = self.clips[idx]
        m = torch.load(os.path.join(clip_dir, "meta.pt"), weights_only=False)
        F, V = m["num_frames"], m["num_cameras"]
        fps = m["fps"]
        s = m["scene_scale"]
        inv = 1.0 / s

        ctx = _context_frames(fps, self.cfg.data.context_times, 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)

        # local-frame origin: mean of context camera centers. Real datasets use a
        # global/city frame with huge translations; centering keeps Plucker moments
        # and signed-exp positions in a numerically sane range (avoids bf16 NaNs).
        ci = torch.tensor([f for f, v in ctx_pairs]); cv = torch.tensor([v for f, v in ctx_pairs])
        origin = m["cam2world"][ci, cv][:, :3, 3].mean(0)  # [3]

        def gather(pairs):
            imgs, Ks, c2ws, deps = [], [], [], []
            for (f, v) in pairs:
                img, Ksc = self._load_image(clip_dir, f, v, m["K"][v])
                imgs.append(img); Ks.append(Ksc)
                c2w = m["cam2world"][f, v].clone()
                c2w[:3, 3] = (c2w[:3, 3] - origin) * inv
                c2ws.append(c2w)
                deps.append(self._load_depth(clip_dir, f, v) * (inv if m["has_depth"] else 1.0))
            return (torch.stack(imgs), torch.stack(Ks), torch.stack(c2ws), torch.stack(deps))

        ctx_images, ctx_K, ctx_c2w, _ = gather(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 * 9973 + 1)
        perm = torch.randperm(len(all_pairs), generator=g)[: self.n_sup]
        sup_pairs = [all_pairs[i] for i in perm.tolist()]
        sup_images, sup_K, sup_c2w, sup_depth = gather(sup_pairs)
        sup_frame = torch.tensor([f for f, v in sup_pairs], dtype=torch.long)
        sup_mono = torch.where(sup_depth > 0, sup_depth, torch.full_like(sup_depth, -1.0))

        x0, y0, dx, dy = m["ground_meta"]
        ground = GridGroundField(m["ground_heights"], x0, y0, dx, dy).transformed(origin, inv)
        lanes = [(l - origin) * inv for l in m["lanes"]]
        boundaries = [(b - origin) * inv for b in m["boundaries"]]

        # anchors: use precomputed if they match the model token budget, else resample
        if m["anchor_pos"].shape[0] == self.cfg.model.tokens.n_map:
            apos, atype, anorm = (m["anchor_pos"] - origin) * inv, m["anchor_type"], m["anchor_normal"]
        else:
            A = sample_anchors(HDMap(ground=ground, lanes=lanes, boundaries=boundaries),
                               self._mcfg, seed=idx)
            apos, atype, anorm = A.positions, A.types, A.normals

        # cap dynamic instances (real clips can have 30+ actors; uncapped this blows up
        # dynamic-token memory = instances * n_dyn_per_instance * N_G). Keep the first
        # data.max_instances; their per-frame boxes are already centered+scaled below.
        I = min(m["box_centers"].shape[0], self.cfg.data.max_instances)
        if I > 0:
            raw_c = m["box_centers"][:I]                  # world frame; untracked frames are (0,0,0)
            bv = raw_c.norm(dim=-1) > 1.0                 # [I,F] per-frame validity (before centering)
            bc = (raw_c - origin) * inv
            br = m["box_rots"][:I]
            bs = m["box_size"][:I] * inv
            bk = m["canon_idx"][:I]
        else:
            bc, br, bs, bk = m["box_centers"], m["box_rots"], m["box_size"], m["canon_idx"]
            bv = torch.zeros(0, F, dtype=torch.bool)

        return MapGSSample(
            ctx_images=ctx_images, ctx_K=ctx_K, ctx_c2w=ctx_c2w, ctx_tids=ctx_tids,
            sup_images=sup_images, sup_K=sup_K, sup_c2w=sup_c2w, sup_depth=sup_depth,
            sup_mono=sup_mono, sup_frame=sup_frame,
            anchor_pos=apos, anchor_type=atype, anchor_normal=anorm,
            ground=ground, lanes=lanes, boundaries=boundaries,
            box_centers=bc, box_rots=br, box_size=bs, canon_idx=bk, box_valid=bv,
            scene_id=idx, scene_scale=s,
        )