"""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):: // meta.pt # poses, intrinsics, map grid, polylines, anchors, boxes images/_.png # native-resolution RGB, frame f, camera v depth/_.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 or / 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, )