scripts/oracle_mapvggt.py

#!/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()