File size: 3,162 Bytes
8cf92b3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | #!/usr/bin/env python3
"""Canonical loader + evaluator for a released MapVGGT (+RefineUNet) checkpoint.
Reconstructs the model, loads the model.*/unet.* state dict, and reports held-out-SCENE
PSNR/SSIM on a segment-disjoint Waymo val split. This is the reference inference path for
the released weights (verifies they round-trip).
Env: VGGT_OMEGA_REPO (vggt-omega clone), MAPVGGT_VGGT_CKPT (base VGGT-Omega weights).
"""
import argparse, os, copy, statistics as st
import torch
from safetensors.torch import load_file
from mapgs.config import load_config
from mapgs.data import UnifiedClipDataset
from mapgs.eval.metrics import psnr, ssim
from mapvggt import MapVGGT
from mapvggt.refine import RefineUNet
from scripts.train_mapvggt_refine import render_rda
from scripts.train_mapvggt_full import prep
DEV = "cuda"
def load_mapvggt_refine(ckpt_path):
"""Reconstruct MapVGGT + RefineUNet from a refine checkpoint (model.*/unet.* keys)."""
model = MapVGGT(with_map=False, with_dyn=False, finetune_backbone=True).to(DEV).eval()
unet = RefineUNet().to(DEV).eval()
sd = load_file(ckpt_path)
mmiss, _ = model.load_state_dict({k[6:]: v for k, v in sd.items() if k.startswith("model.")}, strict=False)
umiss, _ = unet.load_state_dict({k[5:]: v for k, v in sd.items() if k.startswith("unet.")}, strict=False)
return model, unet
@torch.no_grad()
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--ckpt", default="/mnt/william/runs/mapvggt_refine_best.safetensors")
ap.add_argument("--roots", default="/mnt/william/data/unified/waymo")
ap.add_argument("--val-segs", type=int, default=40)
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)
args = ap.parse_args()
H, W = args.height, args.width
model, unet = load_mapvggt_refine(args.ckpt)
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])
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
ps, ss = [], []
for i in range(len(vds.clips)):
d = prep(vds[i], args.n_in, DEV)
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"])))
mp = st.mean(ps); sd = st.pstdev(ps)
print(f"loaded {os.path.basename(args.ckpt)} | held-out-SCENE val (n={len(ps)}): "
f"PSNR {mp:.2f}±{sd:.2f} SSIM {st.mean(ss):.3f}")
if __name__ == "__main__":
main()
|