| """Held-out novel-view evaluation for pose-free feed-forward splatting. |
| context = even frames, targets = odd frames. Target poses come from a joint run on all frames, |
| aligned to the context-only reconstruction with a similarity transform (Umeyama) on camera centres.""" |
| import sys, os, time, json, glob, subprocess |
| sys.path.insert(0, os.path.expanduser('~/polymorphic/vendor/AnySplat')) |
| import numpy as np, torch, lpips |
| from pathlib import Path |
| from PIL import Image |
| import torchvision.transforms.functional as TF |
| from skimage.metrics import structural_similarity |
| from src.model.model.anysplat import AnySplat |
|
|
| def umeyama(src, dst): |
| mu_s, mu_d = src.mean(0), dst.mean(0); xs, xd = src-mu_s, dst-mu_d |
| cov = xd.T @ xs / len(src); U,S,Vt = np.linalg.svd(cov) |
| D = np.eye(3); D[2,2] = np.sign(np.linalg.det(U)*np.linalg.det(Vt)) |
| R = U @ D @ Vt; s = (S*np.diag(D)).sum() / xs.var(0).sum(); t = mu_d - s*R@mu_s |
| return s, R, t |
|
|
| def load(f, mode, long_edge): |
| im = Image.open(f).convert('RGB'); w,h = im.size |
| if mode == 'square': |
| s = long_edge/ min(w,h); im = im.resize((round(w*s), round(h*s)), Image.BICUBIC); w,h = im.size |
| l,t = (w-long_edge)//2, (h-long_edge)//2; im = im.crop((l,t,l+long_edge,t+long_edge)) |
| else: |
| s = long_edge/max(w,h); im = im.resize((max(14,round(w*s/14)*14), max(14,round(h*s/14)*14)), Image.BICUBIC) |
| return TF.to_tensor(im) |
|
|
| def main(video, total, mode, long_edge, outdir): |
| out = Path(outdir); fr = out/'frames'; fr.mkdir(parents=True, exist_ok=True) |
| for f in fr.glob('*.jpg'): f.unlink() |
| dur = float(subprocess.check_output(['ffprobe','-v','error','-show_entries','format=duration','-of','csv=p=0',video]).decode()) |
| subprocess.run(['ffmpeg','-hide_banner','-loglevel','error','-i',video,'-vf',f'fps={total/dur}','-q:v','2','-frames:v',str(total),str(fr/'%04d.jpg'),'-y'],check=True) |
| files = sorted(glob.glob(str(fr/'*.jpg'))); ctx_i = list(range(0,len(files),2)); tgt_i = list(range(1,len(files),2)) |
| imgs = torch.stack([load(f, mode, long_edge) for f in files]).cuda() |
| model = AnySplat.from_pretrained('lhjiang/anysplat').cuda().eval() |
| for p in model.parameters(): p.requires_grad=False |
| with torch.no_grad(): |
| gA, pA = model.inference(imgs[ctx_i].unsqueeze(0)); gB, pB = model.inference(imgs.unsqueeze(0)) |
| eA, eB, iB = pA['extrinsic'][0].double().cpu().numpy(), pB['extrinsic'][0].double().cpu().numpy(), pB['intrinsic'][0].float() |
| s, R, t = umeyama(eB[ctx_i,:3,3], eA[:,:3,3]) |
| resid = np.linalg.norm((s*(R@eB[ctx_i,:3,3].T)).T + t - eA[:,:3,3], axis=1) |
| ext_t = [] |
| for i in tgt_i: |
| E = np.eye(4); E[:3,:3] = R @ eB[i,:3,:3]; E[:3,3] = s*R@eB[i,:3,3] + t; ext_t.append(E) |
| ext_t = torch.tensor(np.stack(ext_t), dtype=torch.float32, device='cuda').unsqueeze(0); intr_t = iB[tgt_i].unsqueeze(0).cuda() |
| H,W = imgs.shape[-2:] |
| with torch.no_grad(): |
| r = model.decoder.forward(gA, ext_t, intr_t, torch.ones(1,len(tgt_i),device='cuda')*.1, torch.ones(1,len(tgt_i),device='cuda')*100,(H,W)).color[0].clamp(0,1) |
| gt = imgs[tgt_i] |
| |
| m = min(H,W); t0,l0 = (H-m)//2,(W-m)//2 |
| rc, gc = r[:,:,t0:t0+m,l0:l0+m], gt[:,:,t0:t0+m,l0:l0+m] |
| psnr = (-10*torch.log10(((rc-gc)**2).mean(dim=(1,2,3)))).cpu().numpy() |
| ssim = np.array([structural_similarity(a.permute(1,2,0).cpu().numpy(), b.permute(1,2,0).cpu().numpy(), channel_axis=2, data_range=1.0) for a,b in zip(rc,gc)]) |
| lp = lpips.LPIPS(net='vgg').cuda().eval() |
| with torch.no_grad(): lpv = lp(rc*2-1, gc*2-1).flatten().cpu().numpy() |
| res = dict(video=Path(video).name, mode=mode, long_edge=long_edge, input_hw=[H,W], n_ctx=len(ctx_i), n_tgt=len(tgt_i), gaussians=int(gA.means.shape[1]), |
| align_resid_mean=float(resid.mean()), PSNR=float(psnr.mean()), SSIM=float(ssim.mean()), LPIPS=float(lpv.mean())) |
| print(json.dumps(res)); (out/f'eval_{mode}_{long_edge}.json').write_text(json.dumps(res, indent=1)) |
| for k in (0, len(tgt_i)//2): |
| Image.fromarray((torch.cat([gc[k],rc[k]],-1).permute(1,2,0).cpu().numpy()*255).astype(np.uint8)).save(out/f'nvs_{mode}_{long_edge}_{k}.png') |
|
|
| if __name__ == '__main__': |
| main(sys.argv[1], int(sys.argv[2]), sys.argv[3], int(sys.argv[4]), sys.argv[5]) |
|
|
