#!/usr/bin/env python3 """Compare deblurring methods (VD-Diff, BAGS, baseline) on ScanNet / TUM. For each (dataset, scene, method) triple the script measures PSNR and SSIM between the method's output and the sharp reference, then writes a summary table to outputs/deblur_comparison/_.json and prints a Markdown table to stdout. Expected folder layout ---------------------- Each deblurred folder must contain images that match the sharp reference one-to-one (same count, same sort order). Sharp reference (ScanNet prototype): data/scannet_blur_proto/vddiff/test/sharp// Deblurred outputs: data/vddiff_deblurred/scannet// ← VD-Diff output data/bags_deblurred/scannet// ← BAGS output (when available) data/scannet_blur_proto/vddiff/test/blur// ← blurred baseline Usage ----- python eval_deblur_comparison.py --dataset scannet --scene scene0004_00 python eval_deblur_comparison.py --dataset tum --scene freiburg2_xyz \\ --ref data/tum_sharp_proto/freiburg2_xyz """ from __future__ import annotations import argparse import json import math from pathlib import Path from typing import NamedTuple import cv2 import numpy as np BASE = Path("/home/szha0669/storage/blur_slam_exp") # ── metric helpers ──────────────────────────────────────────────────────────── def _psnr(a: np.ndarray, b: np.ndarray) -> float: mse = np.mean((a.astype(np.float32) - b.astype(np.float32)) ** 2) return float("inf") if mse == 0 else 20 * math.log10(255.0 / math.sqrt(mse)) def _ssim_gray(a: np.ndarray, b: np.ndarray) -> float: a = cv2.cvtColor(a, cv2.COLOR_BGR2GRAY).astype(np.float64) b = cv2.cvtColor(b, cv2.COLOR_BGR2GRAY).astype(np.float64) c1, c2 = (0.01 * 255) ** 2, (0.03 * 255) ** 2 mu_a, mu_b = a.mean(), b.mean() va, vb = a.var(), b.var() cov = ((a - mu_a) * (b - mu_b)).mean() return ((2 * mu_a * mu_b + c1) * (2 * cov + c2)) / \ ((mu_a ** 2 + mu_b ** 2 + c1) * (va + vb + c2)) class PairMetrics(NamedTuple): psnr: float ssim: float def _eval_folder_pair(ref_dir: Path, pred_dir: Path) -> list[PairMetrics]: ref_imgs = sorted([*ref_dir.glob("*.png"), *ref_dir.glob("*.jpg")]) pred_imgs = sorted([*pred_dir.glob("*.png"), *pred_dir.glob("*.jpg")]) if not pred_imgs: return [] if len(ref_imgs) != len(pred_imgs): n = min(len(ref_imgs), len(pred_imgs)) ref_imgs, pred_imgs = ref_imgs[:n], pred_imgs[:n] results = [] for r, p in zip(ref_imgs, pred_imgs): a = cv2.imread(str(r), cv2.IMREAD_COLOR) b = cv2.imread(str(p), cv2.IMREAD_COLOR) if a is None or b is None: continue if a.shape != b.shape: b = cv2.resize(b, (a.shape[1], a.shape[0]), interpolation=cv2.INTER_LINEAR) results.append(PairMetrics(_psnr(a, b), _ssim_gray(a, b))) return results def _summary(metrics: list[PairMetrics]) -> dict: if not metrics: return {"available": False, "mean_psnr": None, "mean_ssim": None, "n": 0} return { "available": True, "mean_psnr": float(np.mean([m.psnr for m in metrics])), "mean_ssim": float(np.mean([m.ssim for m in metrics])), "n": len(metrics), } # ── main ───────────────────────────────────────────────────────────────────── def main() -> None: ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--dataset", default="scannet", choices=["scannet", "tum"]) ap.add_argument("--scene", default="scene0004_00") ap.add_argument("--ref", type=Path, help="Override sharp-reference folder path") ap.add_argument("--vddiff-dir", type=Path, help="Override VD-Diff output folder") ap.add_argument("--bags-dir", type=Path, help="Override BAGS output folder (omit if not yet available)") ap.add_argument("--blur-dir", type=Path, help="Override blurred-input folder (used as 'no-deblur' baseline)") ap.add_argument("--out", type=Path, help="Write JSON results here (default: outputs/deblur_comparison/)") args = ap.parse_args() # ── resolve default paths ───────────────────────────────────────────────── if args.dataset == "scannet": ref_dir = args.ref or BASE / "data/scannet_blur_proto/vddiff/test/sharp" / args.scene blur_dir = args.blur_dir or BASE / "data/scannet_blur_proto/vddiff/test/blur" / args.scene vddiff_dir = args.vddiff_dir or BASE / "data/vddiff_deblurred/scannet" / args.scene bags_dir = args.bags_dir or BASE / "data/bags_deblurred/scannet" / args.scene else: ref_dir = args.ref or BASE / "data/tum_sharp_proto" / args.scene blur_dir = args.blur_dir or BASE / "data/tum_blur_proto" / args.scene vddiff_dir = args.vddiff_dir or BASE / "data/vddiff_deblurred/tum" / args.scene bags_dir = args.bags_dir or BASE / "data/bags_deblurred/tum" / args.scene if not ref_dir.exists(): raise FileNotFoundError(f"Sharp reference not found: {ref_dir}") # ── evaluate each method ────────────────────────────────────────────────── methods: dict[str, Path] = { "blur_baseline": blur_dir, "vd_diff": vddiff_dir, "bags": bags_dir, } results: dict[str, dict] = {} for name, pred_dir in methods.items(): if pred_dir.exists(): metrics = _eval_folder_pair(ref_dir, pred_dir) results[name] = _summary(metrics) results[name]["pred_dir"] = str(pred_dir) else: results[name] = {"available": False, "mean_psnr": None, "mean_ssim": None, "n": 0, "pred_dir": str(pred_dir), "note": "directory not found"} payload = { "dataset": args.dataset, "scene": args.scene, "ref_dir": str(ref_dir), "methods": results, } # ── write JSON ──────────────────────────────────────────────────────────── out_path = args.out or BASE / "outputs/deblur_comparison" / f"{args.dataset}_{args.scene}.json" out_path.parent.mkdir(parents=True, exist_ok=True) out_path.write_text(json.dumps(payload, indent=2) + "\n") # ── print Markdown table ────────────────────────────────────────────────── header = f"## Deblur comparison — {args.dataset} / {args.scene}\n" header += f"Reference: {ref_dir}\n\n" header += f"| Method | PSNR ↑ | SSIM ↑ | n frames |\n" header += f"|-----------------|--------|--------|----------|\n" print(header, end="") for name, r in results.items(): if r["available"]: print(f"| {name:<15} | {r['mean_psnr']:6.2f} | {r['mean_ssim']:6.4f} | {r['n']:8} |") else: note = r.get("note", "not available") print(f"| {name:<15} | N/A | N/A | ({note}) |") print(f"\nResults written to: {out_path}") if __name__ == "__main__": main()