Perceptual Graphics Quality Evaluation Metrics 02 (including VDP-Based)

import os
import csv
import time
import traceback
from pathlib import Path

import numpy as np
from PIL import Image
import torch
import pyfvvdp


# =========================================================
# Configuration
# =========================================================
DISPLAY_NAME = "standard_4k"
FOVEATED = True
HEATMAP_MODE = "supra-threshold"  # "none", "raw", "threshold", "supra-threshold"
SAVE_HEATMAPS = True

OUTPUT_DIR = Path("out_TargetMethod") 
HEATMAP_DIR = OUTPUT_DIR / "heatmaps"
CSV_PATH = OUTPUT_DIR / "fovvideovdp_TargetMethod.csv"

pairs = [
    ("1AON", "refs/reference01.png", "target/target01.png", "TargetMethod"),
    ("1GRL", "refs/reference02.png", "target/target02.png", "TargetMethod"),
]


# =========================================================
# Helper: convert PyTorch scalar to Python float
# =========================================================
def tensor_to_float(x):
    if isinstance(x, torch.Tensor):
        return float(x.detach().cpu().item())
    return float(x)


# =========================================================
# Helper: save FovVideoVDP heatmap as PNG
# =========================================================
def save_heatmap_png(heatmap_tensor, out_path):
    """
    FovVideoVDP heatmap is usually shaped as:
        [1, C, F, H, W]

    where:
        C = 3 for threshold / supra-threshold heatmaps
        C = 1 for raw heatmap
        F = 1 for a still image
    """
    hm = heatmap_tensor.detach().cpu()

    if hm.ndim != 5:
        raise ValueError(f"Unexpected heatmap shape: {hm.shape}")

    # Take first batch and first frame:
    # [1, C, F, H, W] -> [C, H, W]
    hm = hm[0, :, 0, :, :].to(torch.float32)

    # [C, H, W] -> [H, W, C]
    hm = hm.permute(1, 2, 0).numpy()

    # If raw heatmap has one channel, convert to RGB for PNG saving
    if hm.shape[2] == 1:
        hm = np.repeat(hm, 3, axis=2)

    hm = np.nan_to_num(hm, nan=0.0, posinf=1.0, neginf=0.0)
    hm = np.clip(hm, 0.0, 1.0)

    hm_uint8 = (hm * 255.0).astype(np.uint8)
    Image.fromarray(hm_uint8).save(out_path)


# =========================================================
# Main
# =========================================================
def main():
    print("Starting FovVideoVDP batch evaluation...", flush=True)

    device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
    print(f"Device       : {device}", flush=True)
    print(f"Display model: {DISPLAY_NAME}", flush=True)
    print(f"Foveated     : {FOVEATED}", flush=True)
    print(f"Heatmap mode : {HEATMAP_MODE}", flush=True)

    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    if SAVE_HEATMAPS:
        HEATMAP_DIR.mkdir(parents=True, exist_ok=True)

    print("\nCreating FovVideoVDP metric...", flush=True)

    fv = pyfvvdp.fvvdp(
        display_name=DISPLAY_NAME,
        foveated=FOVEATED,
        heatmap=HEATMAP_MODE,
        device=device,
        quiet=False
    )

    metric_info = fv.get_info_string()
    print(f"Metric info: {metric_info}", flush=True)

    results = []

    for molecule, ref_path, test_path, setting in pairs:
        print(f"\nProcessing {molecule} - {setting}", flush=True)
        print(f"Reference: {ref_path}", flush=True)
        print(f"Test     : {test_path}", flush=True)

        if not os.path.exists(ref_path):
            raise FileNotFoundError(f"Reference image not found: {ref_path}")

        if not os.path.exists(test_path):
            raise FileNotFoundError(f"Test image not found: {test_path}")

        ref = pyfvvdp.load_image_as_array(ref_path)
        test = pyfvvdp.load_image_as_array(test_path)

        if ref.shape[0] != test.shape[0] or ref.shape[1] != test.shape[1]:
            raise ValueError(
                f"Image resolution mismatch for {molecule}: "
                f"ref shape={ref.shape}, test shape={test.shape}"
            )

        start = time.time()

        with torch.no_grad():
            # IMPORTANT:
            # FovVideoVDP uses test first, reference second.
            q_jod, stats = fv.predict(test, ref, dim_order="HWC")

        elapsed = time.time() - start
        q_jod_value = tensor_to_float(q_jod)

        print(f"FovVideoVDP score: {q_jod_value:.6f} JOD", flush=True)
        print(f"Time taken        : {elapsed:.3f} s", flush=True)

        heatmap_path = ""

        if SAVE_HEATMAPS and isinstance(stats, dict) and "heatmap" in stats:
            heatmap_path = HEATMAP_DIR / f"{molecule}_{setting.lower()}_heatmap.png"
            save_heatmap_png(stats["heatmap"], heatmap_path)
            heatmap_path = str(heatmap_path)
            print(f"Saved heatmap     : {heatmap_path}", flush=True)

        results.append({
            "molecule": molecule,
            "setting": setting,
            "fovvideovdp_jod": q_jod_value,
            "display_name": DISPLAY_NAME,
            "foveated": FOVEATED,
            "heatmap_mode": HEATMAP_MODE,
            "metric_info": metric_info,
            "heatmap_path": heatmap_path
        })

    with open(CSV_PATH, "w", newline="") as f:
        writer = csv.DictWriter(
            f,
            fieldnames=[
                "molecule",
                "setting",
                "fovvideovdp_jod",
                "display_name",
                "foveated",
                "heatmap_mode",
                "metric_info",
                "heatmap_path"
            ]
        )
        writer.writeheader()
        writer.writerows(results)

    print(f"\nSaved CSV: {CSV_PATH}", flush=True)
    print("Done.", flush=True)


if __name__ == "__main__":
    try:
        main()
    except Exception:
        print("\nERROR occurred:", flush=True)
        traceback.print_exc()