| |
| """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() |
|
|
