File size: 2,643 Bytes
c75b162 | 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
"""
Reproduce Table 2: Blur | EVSSM | A(noNIMA) | D(global)
- Blur : original ScanNet frame NIMA (no-ref input quality)
- EVSSM : bags_main renders NIMA (BAGS trained on EVSSM-deblurred data, best proxy)
- A(noNIMA): BAGS on bags_orig, no NIMA, no depth (bags_A_4scenes experiment)
- D(global): BAGS on bags_orig, NIMA + depth alpha=0.1 (bags_D_4scenes experiment)
Run first: scripts/run_bags_AD_4scenes.sh
"""
import os, glob, json, numpy as np, torch
import pyiqa
BASE = "/home/szha0669/storage/blur_slam_exp"
ITERS = 15000
SCENES = ["scene0000_00", "scene0000_01", "scene0005_00", "scene0006_00"]
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
nima_fn = pyiqa.create_metric("nima-koniq", device=device)
def nima_dir(d):
imgs = sorted(glob.glob(f"{d}/*.png") + glob.glob(f"{d}/*.jpg"))
if not imgs:
return None
return float(np.mean([float(nima_fn(p).item()) for p in imgs]))
# Load pre-computed Blur + EVSSM from nima_main_results.json
main_results = {}
main_path = f"{BASE}/outputs/eval/nima_main_results.json"
if os.path.exists(main_path):
with open(main_path) as f:
main_results = json.load(f)
print(f"\n{'Scene':<16} {'Blur':>7} {'EVSSM':>7} {'A(noNIMA)':>11} {'D(global)':>11}")
print("-" * 58)
results = {}
for scene in SCENES:
blur_nima = main_results.get(scene, {}).get("orig") # ScanNet original frames
evssm_nima = main_results.get(scene, {}).get("bags") # bags_main renders (proxy for EVSSM)
# A(noNIMA): bags_A_4scenes (BAGS on bags_orig, no NIMA, no depth)
a_dir = f"{BASE}/outputs/bags_A_4scenes/scannet/{scene}/test/ours_{ITERS}/test_preds_1"
a_nima = nima_dir(a_dir) if os.path.isdir(a_dir) else None
# D(global): bags_D_4scenes (BAGS on bags_orig, NIMA + depth alpha=0.1)
d_dir = f"{BASE}/outputs/bags_D_4scenes/scannet/{scene}/test/ours_{ITERS}/test_preds_1"
d_nima = nima_dir(d_dir) if os.path.isdir(d_dir) else None
results[scene] = dict(blur=blur_nima, evssm=evssm_nima, A=a_nima, D=d_nima)
def fmt(v):
return f"{v:.3f}" if v is not None else " N/A"
print(f"{scene:<16} {fmt(blur_nima):>7} {fmt(evssm_nima):>7} {fmt(a_nima):>11} {fmt(d_nima):>11}")
# MEAN
print("-" * 58)
for key, label in [("blur","Blur"),("evssm","EVSSM"),("A","A(noNIMA)"),("D","D(global)")]:
vals = [results[s][key] for s in SCENES if results[s].get(key) is not None]
if vals:
print(f"{'MEAN':<16} {np.mean(vals):>7.3f} ({label})")
out = f"{BASE}/outputs/eval/table2_comparison.json"
with open(out, "w") as f:
json.dump(results, f, indent=2)
print(f"\nSaved: {out}")
|