| |
| """Zero-shot (NO finetune) TokenGS dl3dv_base PSNR on a few TRAIN clips: |
| sample 2 clips from av2 train + 2 from waymo train, eval novel-view PSNR/SSIM. |
| Same 448x768 / n_in=8 / s=s_max protocol as the reported warm-start number.""" |
| import argparse, random, copy |
| import torch |
| from safetensors.torch import load_file |
|
|
| from tokengs import options |
| from mapgs.config import load_config |
| from mapgs.data import UnifiedClipDataset |
| from maptokengs import MapTokenGS |
|
|
| import scripts.finetune_maptokengs as F |
|
|
| AV2 = "/mnt/william/data/unified/av2" |
| WAYMO = "/mnt/william/data/unified/waymo" |
| CKPT = "/mnt/william/tokengs_ckpts/checkpoints/dl3dv_base.safetensors" |
|
|
|
|
| def pick(ds, k, seed): |
| rng = random.Random(seed) |
| idx = sorted(rng.sample(range(len(ds)), k)) |
| sub = copy.copy(ds) |
| sub.clips = [ds.clips[i] for i in idx] |
| return sub, [ds.clips[i].split("/")[-1] for i in idx] |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--height", type=int, default=448) |
| ap.add_argument("--width", type=int, default=768) |
| ap.add_argument("--n-in", type=int, default=8) |
| ap.add_argument("--k", type=int, default=2) |
| ap.add_argument("--seed", type=int, default=0) |
| args = ap.parse_args() |
| dev = "cuda" |
| H, W = args.height, args.width |
|
|
| opt = options.config_defaults["train_dl3dv_base"]; opt = opt() if callable(opt) else opt |
| opt.img_size = (H, W) |
| model = MapTokenGS(opt, s0=0.5, s_max=2.0, max_instances=8, dyn_per_instance=32).to(dev) |
| miss, unexp = model.load_state_dict(load_file(CKPT), strict=False) |
| model.eval() |
| print(f"warm-started dl3dv_base (ZERO-SHOT, no finetune) | fresh {len(miss)} unexpected {len(unexp)} " |
| f"| res {H}x{W} n_in {args.n_in}", flush=True) |
|
|
| cfg = load_config(overrides=["data.name=unified", f"data.root={AV2},{WAYMO}", |
| f"data.height={H}", f"data.width={W}", |
| "model.tokens.n_map=2048", "data.max_instances=8"]) |
|
|
| ds_av2 = UnifiedClipDataset(cfg, roots=[AV2], split="train", n_sup_views=4) |
| ds_way = UnifiedClipDataset(cfg, roots=[WAYMO], split="train", n_sup_views=4) |
| sub_av2, names_av2 = pick(ds_av2, args.k, args.seed) |
| sub_way, names_way = pick(ds_way, args.k, args.seed) |
| print("av2 clips: ", names_av2, flush=True) |
| print("waymo clips:", names_way, flush=True) |
|
|
| with torch.no_grad(): |
| pa, sa, da, na = F.held_out_psnr(model, sub_av2, args.k, args.n_in, 1.0, dev) |
| pw, sw, dw, nw = F.held_out_psnr(model, sub_way, args.k, args.n_in, 1.0, dev) |
|
|
| |
| ps, ss = [], [] |
| for sub in (sub_av2, sub_way): |
| for i in range(len(sub.clips)): |
| d = F.prep(sub[i], args.n_in, 1.0, dev) |
| g = model.forward_reconstruction_mapped(d["mi"], d["anchors"], d["atype"], d["anormal"]) |
| gg = model.gaussian_group()[0]; ig = model.gaussian_instance()[0] |
| model.cur_s = model.s_max |
| rgb, _, _ = F.render_dyn(model, g, gg, ig, d["mi"].decoder.cam_view, |
| d["mi"].decoder.intrinsics, d["sup_frame"], d["dyn"]) |
| rgb = rgb.clamp(0, 1) |
| p, s = float(F.psnr(rgb, d["gt"])), float(F.ssim(rgb, d["gt"])) |
| if p == p and abs(p) != float("inf"): |
| ps.append(p); ss.append(s) |
|
|
| print(f"\nAV2 (n={na}): PSNR {pa:.2f}+-{da:.2f} SSIM {sa:.3f}") |
| print(f"Waymo (n={nw}): PSNR {pw:.2f}+-{dw:.2f} SSIM {sw:.3f}") |
| mp = sum(ps)/len(ps); sd = (sum((x-mp)**2 for x in ps)/len(ps))**0.5 |
| print(f"ALL (n={len(ps)}): PSNR {mp:.2f}+-{sd:.2f} SSIM {sum(ss)/len(ss):.3f}") |
| print("per-clip PSNR:", [round(x,2) for x in ps]) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
