#!/usr/bin/env python3 """Oracle upper-bound experiment for MapVGGT. Take a few scene-disjoint val clips. For each: 1. Feed-forward: run the trained model -> gaussians -> render held-out sup views -> PSNR/SSIM. 2. Per-scene optimization on the GAUSSIAN PARAMETERS themselves (free means/scales/quats/ opacities/colors), Adam against the held-out target views. This is the upper bound the *representation* (pixel-aligned, sparse-where-disoccluded, possibly w/ map+dyn gaussians) can reach with perfect fitting. - Phase A: optimize ONLY colors (report dB from color alone). - Phase B: optimize ALL params (the representation ceiling). 3. Disocclusion coverage: fraction of target pixels with NO gaussian projecting near them. Reports feed-forward vs per-scene-optimized PSNR per clip; the gap separates "bad prediction" from "representation ceiling". """ import argparse, os, copy, sys import numpy as np import torch import torch.nn.functional as F from gsplat import rasterization from mapgs.config import load_config from mapgs.data import UnifiedClipDataset from mapgs.eval.metrics import psnr, ssim from mapvggt import MapVGGT from mapvggt.heads import cat_gaussians DEV = "cuda" sys.path.insert(0, "/mnt/william") from scripts.train_mapvggt_full import prep, render_scene # reuse exact prep + render def build_val(args): cfg = load_config(overrides=["data.name=unified", f"data.root={args.roots}", f"data.height={args.height}", f"data.width={args.width}", "model.tokens.n_map=2048"]) full = UnifiedClipDataset(cfg, roots=args.roots.split(","), split="train", n_sup_views=6) def segid(p): return "_".join(os.path.basename(p.rstrip("/")).split("_")[:2]) segs = sorted(set(segid(c) for c in full.clips)) val_segs = set(segs[:args.val_segs]) seen, vclips = set(), [] for c in full.clips: sgi = segid(c) if sgi in val_segs and sgi not in seen: seen.add(sgi); vclips.append(c) vds = copy.copy(full); vds.clips = vclips return vds def assemble_full_gaussians(model, gsm, d): """Replicate render_scene's union (static+map + per-frame dynamics) into ONE set of free gaussian params, plus a list mapping which gaussians belong to which sup frame. Dynamics are frame-dependent, so we bake per-frame copies and tag them. Returns a dict of leaf tensors and a per-gaussian frame-mask (None => valid for all frames).""" canon = model.dyn_canonical() if (model.with_dyn and d["dyn"] is not None) else None means = [gsm["means"]]; scales = [gsm["scales"]]; quats = [gsm["quats"]] opac = [gsm["opacities"]]; colors = [gsm["colors"]] frame_of = [torch.full((gsm["means"].shape[0],), -1, dtype=torch.long, device=DEV)] # -1 = all frames if canon is not None: dn = d["dyn"] for i in range(d["sup_c2w"].shape[0]): fi = int(d["sup_frame"][i]) dg = model.dyn_head.place(canon, fi, dn["c"], dn["R"], dn["canon"], dn["valid"], dn["radius"], gain=1.0) n = dg["means"].shape[0] if n > 0: means.append(dg["means"]); scales.append(dg["scales"]); quats.append(dg["quats"]) opac.append(dg["opacities"]); colors.append(dg["colors"]) frame_of.append(torch.full((n,), i, dtype=torch.long, device=DEV)) g = dict(means=torch.cat(means), scales=torch.cat(scales), quats=torch.cat(quats), opacities=torch.cat(opac), colors=torch.cat(colors), frame_of=torch.cat(frame_of)) return g def render_from_params(g, d, H, W, view_idx): """Render sup view `view_idx` from free param dict g (means in world, scales, quats raw, opacities raw-logit, colors raw-logit). frame_of selects dyn gaussians.""" fo = g["frame_of"] keep = (fo == -1) | (fo == view_idx) out, _, _ = rasterization( means=g["means"][keep], quats=F.normalize(g["quats"][keep], dim=-1), scales=g["scales"][keep].clamp(min=1e-6), opacities=torch.sigmoid(g["opac_logit"][keep]), colors=torch.sigmoid(g["color_logit"][keep]), viewmats=torch.inverse(d["sup_c2w"][view_idx:view_idx + 1]), Ks=d["sup_K"][view_idx:view_idx + 1], width=W, height=H, near_plane=0.01, far_plane=500.0, render_mode="RGB") return out[0, ..., :3].clamp(0, 1).permute(2, 0, 1) def eval_params(g, d, H, W): rgbs = [render_from_params(g, d, H, W, i) for i in range(d["sup_c2w"].shape[0])] rgb = torch.stack(rgbs) return float(psnr(rgb, d["sup_img"])), float(ssim(rgb, d["sup_img"])) def disocclusion_coverage(g, d, H, W, thresh_px=1.5): """Fraction of target pixels with NO gaussian center projecting within thresh_px. Uses gaussian *centers* projected into each sup view (a generous lower bound on holes: a pixel with no nearby center is genuinely uncovered).""" holes_total, px_total = 0, 0 means = g["means"] # world for i in range(d["sup_c2w"].shape[0]): w2c = torch.inverse(d["sup_c2w"][i]) K = d["sup_K"][i] ones = torch.ones(means.shape[0], 1, device=DEV) cam = (w2c @ torch.cat([means, ones], 1).T).T[:, :3] z = cam[:, 2] front = z > 0.05 uv = (K @ cam.T).T u = uv[:, 0] / uv[:, 2]; v = uv[:, 1] / uv[:, 2] valid = front & (u >= 0) & (u < W) & (v >= 0) & (v < H) # bin centers into pixel grid (downsampled by thresh) -> coverage mask ui = (u[valid] / thresh_px).long(); vi = (v[valid] / thresh_px).long() gw = int(np.ceil(W / thresh_px)); gh = int(np.ceil(H / thresh_px)) cov = torch.zeros(gh, gw, dtype=torch.bool, device=DEV) ui = ui.clamp(0, gw - 1); vi = vi.clamp(0, gh - 1) cov[vi, ui] = True holes_total += int((~cov).sum()); px_total += gh * gw return holes_total / max(1, px_total) def optimize(g, d, H, W, steps, lr, params): """Optimize selected params (list of keys among means,scales,quats,opac_logit,color_logit). Returns the optimized g (in place leaves).""" for k in params: g[k] = g[k].detach().clone().requires_grad_(True) opt = torch.optim.Adam([g[k] for k in params], lr=lr) tgt = d["sup_img"] for step in range(steps): rgbs = [render_from_params(g, d, H, W, i) for i in range(d["sup_c2w"].shape[0])] rgb = torch.stack(rgbs) loss = F.l1_loss(rgb, tgt) + 0.1 * (1 - ssim(rgb, tgt)) opt.zero_grad(set_to_none=True) loss.backward() opt.step() for k in params: g[k] = g[k].detach() return g def main(): ap = argparse.ArgumentParser() ap.add_argument("--roots", default="/mnt/william/data/unified/waymo") ap.add_argument("--ckpt", default="/mnt/william/runs/abl_full_best.safetensors") ap.add_argument("--n-in", type=int, default=8) ap.add_argument("--height", type=int, default=256) ap.add_argument("--width", type=int, default=448) ap.add_argument("--val-segs", type=int, default=40) ap.add_argument("--clips", type=int, default=3) # how many val clips to oracle ap.add_argument("--color-steps", type=int, default=400) ap.add_argument("--all-steps", type=int, default=800) ap.add_argument("--lr-color", type=float, default=0.02) ap.add_argument("--lr-all", type=float, default=0.01) args = ap.parse_args() H, W = args.height, args.width model = MapVGGT(with_map=True, with_dyn=True, finetune_backbone=False).to(DEV) from safetensors.torch import load_file sd = load_file(args.ckpt) miss, unexp = model.load_state_dict(sd, strict=False) print(f"loaded {args.ckpt}: {len(sd)} tensors; missing(non-vggt) " f"{[m for m in miss if not m.startswith('vggt.')][:5]}", flush=True) model.eval(); model.cur_s = model.s_max vds = build_val(args) print(f"val clips available: {len(vds.clips)}; oracling first {args.clips}", flush=True) rows = [] for ci in range(min(args.clips, len(vds.clips))): d = prep(vds[ci], args.n_in, DEV) with torch.no_grad(): gsm = model(d["in_img"], d["in_K"], d["in_c2w"], d["ap"], d["at"], d["an"]) # feed-forward via the SAME render path used in training/eval ff_rgb, _ = render_scene(model, gsm, d, H, W, gain=1.0) ff_psnr = float(psnr(ff_rgb, d["sup_img"])); ff_ssim = float(ssim(ff_rgb, d["sup_img"])) # assemble free params with torch.no_grad(): g = assemble_full_gaussians(model, gsm, d) g["opac_logit"] = torch.logit(g["opacities"].clamp(1e-4, 1 - 1e-4)) g["color_logit"] = torch.logit(g["colors"].clamp(1e-4, 1 - 1e-4)) nG = g["means"].shape[0] hole_frac = disocclusion_coverage(g, d, H, W) # sanity: rebuild PSNR through render_from_params (should match ff closely) with torch.no_grad(): rp_psnr, _ = eval_params(g, d, H, W) # Phase A: optimize colors only g = optimize(g, d, H, W, args.color_steps, args.lr_color, ["color_logit"]) with torch.no_grad(): ca_psnr, ca_ssim = eval_params(g, d, H, W) # Phase B: optimize ALL params (continue from color-opt state) g = optimize(g, d, H, W, args.all_steps, args.lr_all, ["means", "scales", "quats", "opac_logit", "color_logit"]) with torch.no_grad(): full_psnr, full_ssim = eval_params(g, d, H, W) clipname = os.path.basename(vds.clips[ci]) print(f"\n[clip {ci}] {clipname} G={nG//1000}k holes~{hole_frac*100:.1f}%", flush=True) print(f" feed-forward PSNR {ff_psnr:6.2f} SSIM {ff_ssim:.3f}", flush=True) print(f" (param-render chk) PSNR {rp_psnr:6.2f}", flush=True) print(f" +color-only opt PSNR {ca_psnr:6.2f} SSIM {ca_ssim:.3f} (Δ color {ca_psnr-ff_psnr:+.2f} dB)", flush=True) print(f" +ALL-param opt PSNR {full_psnr:6.2f} SSIM {full_ssim:.3f} (Δ all {full_psnr-ff_psnr:+.2f} dB)", flush=True) rows.append((clipname, nG, hole_frac, ff_psnr, ca_psnr, full_psnr, ff_ssim, full_ssim)) print("\n================ SUMMARY ================", flush=True) print(f"{'clip':28s} {'G(k)':>5s} {'hole%':>6s} {'FF':>6s} {'+col':>6s} {'+all':>6s} {'gap':>6s}", flush=True) for (cn, nG, hf, ff, ca, fu, _, _) in rows: print(f"{cn[:28]:28s} {nG//1000:5d} {hf*100:6.1f} {ff:6.2f} {ca:6.2f} {fu:6.2f} {fu-ff:6.2f}", flush=True) arr = np.array([(r[3], r[4], r[5]) for r in rows]) print(f"\nMEAN FF {arr[:,0].mean():.2f} | +color {arr[:,1].mean():.2f} " f"(Δ {arr[:,1].mean()-arr[:,0].mean():+.2f}) | +all {arr[:,2].mean():.2f} " f"(Δ {arr[:,2].mean()-arr[:,0].mean():+.2f})", flush=True) print(f"MEAN hole-fraction {np.mean([r[2] for r in rows])*100:.1f}%", flush=True) if __name__ == "__main__": main()