scripts/eval_deblur_comparison.py

#!/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/<dataset>_<scene>.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/<scene>/

  Deblurred outputs:
    data/vddiff_deblurred/scannet/<scene>/     ← VD-Diff output
    data/bags_deblurred/scannet/<scene>/       ← BAGS output (when available)
    data/scannet_blur_proto/vddiff/test/blur/<scene>/  ← 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()