| """Region-level PSNR (edge/flat by Sobel quartile of GT) → Table 7 expansion.""" |
| import os, argparse, glob |
| import numpy as np |
| from PIL import Image |
| import torch |
| import torch.nn.functional as F |
|
|
|
|
| def sobel_mag(img_chw): |
| luma = (0.299*img_chw[0] + 0.587*img_chw[1] + 0.114*img_chw[2]).unsqueeze(0).unsqueeze(0) |
| wx = torch.tensor([[1,0,-1],[2,0,-2],[1,0,-1]], dtype=torch.float32).view(1,1,3,3) |
| wy = torch.tensor([[1,2,1],[0,0,0],[-1,-2,-1]], dtype=torch.float32).view(1,1,3,3) |
| gx = F.conv2d(luma, wx, padding=1) |
| gy = F.conv2d(luma, wy, padding=1) |
| return torch.sqrt(gx**2 + gy**2 + 1e-8).squeeze() |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--methods", nargs="+", default=["vanilla_logged", "vanillasgf_logged", |
| "vanillasgf_nosg", "vanillasgf_hook"]) |
| ap.add_argument("--scenes", nargs="+", default=["francis","m60","panther","bonsai","bicycle","garden"]) |
| ap.add_argument("--outputs_root", default="/root/autodl-tmp/SplatAtlas/outputs") |
| ap.add_argument("--iter", type=int, default=30000) |
| args = ap.parse_args() |
|
|
| print(f"{'Method':<22} {'Scene':<10} {'Edge PSNR':>10} {'Flat PSNR':>10} {'Full PSNR':>10}") |
| print("-"*68) |
| for method in args.methods: |
| all_e, all_f, all_t = [], [], [] |
| for scene in args.scenes: |
| rd = os.path.join(args.outputs_root, f"{method}_{scene}", f"renders_test_{args.iter}") |
| gd = os.path.join(args.outputs_root, f"{method}_{scene}", f"gt_test_{args.iter}") |
| if not (os.path.exists(rd) and os.path.exists(gd)): |
| print(f"{method:<22} {scene:<10} (renders/gt missing)"); continue |
| files = sorted(glob.glob(os.path.join(rd, "*.png")) + glob.glob(os.path.join(rd, "*.jpg"))) |
| se, sf, st = [], [], [] |
| for pf in files: |
| gf = os.path.join(gd, os.path.basename(pf)) |
| if not os.path.exists(gf): continue |
| pred = torch.from_numpy(np.array(Image.open(pf).convert("RGB")).astype(np.float32)/255.).permute(2,0,1) |
| gt = torch.from_numpy(np.array(Image.open(gf).convert("RGB")).astype(np.float32)/255.).permute(2,0,1) |
| if pred.shape != gt.shape: continue |
| sob = sobel_mag(gt) |
| q75, q25 = float(sob.quantile(0.75)), float(sob.quantile(0.25)) |
| em, fm = (sob >= q75), (sob <= q25) |
| d = ((pred-gt)**2).mean(0) |
| se.append(float(d[em].mean()) if em.any() else 0) |
| sf.append(float(d[fm].mean()) if fm.any() else 0) |
| st.append(float(d.mean())) |
| if not st: |
| print(f"{method:<22} {scene:<10} (no images)"); continue |
| ep = -10*np.log10(np.mean(se)+1e-12) |
| fp = -10*np.log10(np.mean(sf)+1e-12) |
| tp = -10*np.log10(np.mean(st)+1e-12) |
| print(f"{method:<22} {scene:<10} {ep:>10.2f} {fp:>10.2f} {tp:>10.2f}") |
| all_e.extend(se); all_f.extend(sf); all_t.extend(st) |
| if all_t: |
| print(f"{method:<22} {'AVERAGE':<10} " |
| f"{-10*np.log10(np.mean(all_e)+1e-12):>10.2f} " |
| f"{-10*np.log10(np.mean(all_f)+1e-12):>10.2f} " |
| f"{-10*np.log10(np.mean(all_t)+1e-12):>10.2f}") |
| print("-"*68) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
