#!/usr/bin/env python3 """MapVGGT + render-then-refine UNet decoder. Backbone-only static per-pixel gaussians (map/dyn are neutral) → gsplat renders RGB+depth+alpha → UNet refines → photometric loss. Warm-start from abl_base_best; train end-to-end (UNet + head + gentle VGGT) on scene-disjoint Waymo. Reports final-ckpt held-out-SCENE PSNR/SSIM.""" import argparse, time, os, random, copy import numpy as np, torch import torch.nn.functional as F from gsplat import rasterization from safetensors.torch import load_file, save_file from mapgs.config import load_config from mapgs.data import UnifiedClipDataset from mapgs.hdmap.rasterize_map import rasterize_map_depth from mapgs.losses import mapdepth_loss, Tempering from mapgs.eval.metrics import psnr, ssim from mapvggt import MapVGGT from mapvggt.refine import RefineUNet from scripts.train_mapvggt_full import prep DEV = "cuda" def render_rda(g, c2w, K, H, W): """Batched render over views -> rgb [S,3,H,W], depth [S,H,W], alpha [S,H,W].""" out, alpha, _ = rasterization(means=g["means"], quats=g["quats"], scales=g["scales"], opacities=g["opacities"], colors=g["colors"], viewmats=torch.inverse(c2w), Ks=K, width=W, height=H, near_plane=0.01, far_plane=500.0, render_mode="RGB+ED") rgb = out[..., :3].clamp(0, 1).permute(0, 3, 1, 2) depth = out[..., 3] return rgb, depth, alpha[..., 0] @torch.no_grad() def evaluate(model, unet, vds, n_in, H, W, device): model.eval(); unet.eval(); ps, ss = [], [] for i in range(len(vds.clips)): d = prep(vds[i], n_in, device) g = model(d["in_img"], d["in_K"], d["in_c2w"]) rgb, dep, al = render_rda(g, d["sup_c2w"], d["sup_K"], H, W) ref = unet(rgb, dep, al) p = float(psnr(ref, d["sup_img"])) if p == p and abs(p) != float("inf"): ps.append(p); ss.append(float(ssim(ref, d["sup_img"]))) model.train(); unet.train() if not model.finetune_backbone: model.vggt.eval() mp = sum(ps) / max(1, len(ps)); sd = (sum((x - mp) ** 2 for x in ps) / max(1, len(ps))) ** 0.5 return mp, sum(ss) / max(1, len(ss)), sd def main(): ap = argparse.ArgumentParser() ap.add_argument("--roots", default="/mnt/william/data/unified/waymo") ap.add_argument("--init", default="/mnt/william/runs/abl_base_best.safetensors") ap.add_argument("--iters", type=int, default=4000) 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("--lr-unet", type=float, default=3e-4) ap.add_argument("--lr-head", type=float, default=1e-4) ap.add_argument("--lr-vggt", type=float, default=1e-5) ap.add_argument("--lam-md", type=float, default=0.5) ap.add_argument("--val-segs", type=int, default=40) ap.add_argument("--eval-every", type=int, default=500) ap.add_argument("--seed", type=int, default=0) ap.add_argument("--out", default="/mnt/william/runs/mapvggt_refine.safetensors") args = ap.parse_args() H, W = args.height, args.width random.seed(args.seed); np.random.seed(args.seed); torch.manual_seed(args.seed) model = MapVGGT(with_map=False, with_dyn=False, finetune_backbone=True).to(DEV) miss, unexp = model.load_state_dict(load_file(args.init), strict=False) unet = RefineUNet().to(DEV) cfg = load_config(overrides=["data.name=unified", f"data.root={args.roots}", f"data.height={H}", f"data.width={W}", "model.tokens.n_map=2048"]) full = UnifiedClipDataset(cfg, roots=args.roots.split(","), split="train", n_sup_views=6) segid = lambda p: "_".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]) ds = copy.copy(full); ds.clips = [c for c in full.clips if segid(c) not in val_segs] seen = set(); vclips = [c for c in full.clips if segid(c) in val_segs and not (segid(c) in seen or seen.add(segid(c)))] vds = copy.copy(full); vds.clips = vclips print(f"MapVGGT+RefineUNet | init {os.path.basename(args.init)} (miss {len(miss)}) | train {len(ds.clips)} | " f"val {len(vds.clips)} scene-disjoint segs | {H}x{W} n_in {args.n_in}", flush=True) temper = Tempering(cfg.loss, cfg.model.tokens, args.iters) vggt_ids = {id(p) for p in model.vggt.parameters()} opt = torch.optim.AdamW([ {"params": list(unet.parameters()), "lr": args.lr_unet}, {"params": [p for p in model.parameters() if p.requires_grad and id(p) not in vggt_ids], "lr": args.lr_head}, {"params": [p for p in model.vggt.parameters() if p.requires_grad], "lr": args.lr_vggt}, ], betas=(0.9, 0.95)) def tsd(): sd = {f"unet.{k}": v for k, v in unet.state_dict().items()} sd.update({f"model.{k}": v for k, v in model.state_dict().items()}) return sd b_ps, b_ss, b_sd = evaluate(model, unet, vds, args.n_in, H, W, DEV) print(f"BEFORE (UNet zero-init == baseline): val PSNR {b_ps:.2f}±{b_sd:.2f} SSIM {b_ss:.3f}", flush=True) best_path = args.out.replace(".safetensors", "_best.safetensors"); best = b_ps t = time.time() for step in range(args.iters): d = prep(ds[step % len(ds.clips)], args.n_in, DEV) g = model(d["in_img"], d["in_K"], d["in_c2w"]) rgb, depth, al = render_rda(g, d["sup_c2w"], d["sup_K"], H, W) ref = unet(rgb, depth, al) l_rgb = F.l1_loss(ref, d["sup_img"]) + 0.1 * (1 - ssim(ref, d["sup_img"])) with torch.no_grad(): md, mask = rasterize_map_depth(d["ground"], d["sup_K"], d["sup_c2w"], H, W) l_md = mapdepth_loss(depth, md, mask, temper.eps(step), cfg.loss.huber_delta) if mask.any() else depth.sum() * 0 loss = l_rgb + args.lam_md * l_md opt.zero_grad(set_to_none=True) if torch.isfinite(loss): loss.backward() gn = torch.nn.utils.clip_grad_norm_([p for grp in opt.param_groups for p in grp["params"]], 1.0) if torch.isfinite(gn): opt.step() if step % 50 == 0 or step < 4: print(f"it {step:5d} | loss {float(loss):.4f} rgb {float(l_rgb):.4f} md {float(l_md):.4f} " f"| {time.time()-t:.0f}s", flush=True) if step > 0 and step % args.eval_every == 0: e_ps, e_ss, e_sd = evaluate(model, unet, vds, args.n_in, H, W, DEV) tag = "" if e_ps > best: best = e_ps; save_file(tsd(), best_path); tag = " *best" print(f" [eval @ {step}] val PSNR {e_ps:.2f}±{e_sd:.2f} SSIM {e_ss:.3f}{tag} | {time.time()-t:.0f}s", flush=True) a_ps, a_ss, a_sd = evaluate(model, unet, vds, args.n_in, H, W, DEV) if a_ps > best: best = a_ps; save_file(tsd(), best_path) print(f"\nAFTER ({args.iters} it): val PSNR {a_ps:.2f}±{a_sd:.2f} SSIM {a_ss:.3f} | BEFORE {b_ps:.2f} | BEST {best:.2f}", flush=True) if __name__ == "__main__": main()