#!/usr/bin/env python3 """Strict paired IQA scorer matching the released RAESR evaluation metric set.""" from __future__ import annotations import argparse import csv import json import math import os import sys from pathlib import Path DELIVERY_ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(DELIVERY_ROOT)) import numpy as np import torch from PIL import Image from raesr.artifacts import file_record, runtime_environment, write_json IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp"} METRIC_SPEC = { "psnr": ("psnr", {"test_y_channel": True, "color_space": "ycbcr"}, True), "ssim": ("ssim", {"test_y_channel": True, "color_space": "ycbcr"}, True), "lpips": ("lpips", {}, True), "dists": ("dists", {}, True), "niqe": ("niqe", {}, False), "musiq": ("musiq", {}, False), "maniqa_pipal": ("maniqa-pipal", {}, False), "topiq_nr": ("topiq_nr", {}, False), "clipiqa": ("clipiqa", {}, False), } DEFAULT_METRICS = list(METRIC_SPEC) METRIC_CACHE_FILES = { "lpips": [ "hub/checkpoints/alexnet-owt-7be5be79.pth", "hub/pyiqa/LPIPS_v0.1_alex-df73285e.pth", ], "dists": [ "hub/checkpoints/vgg16-397923af.pth", "hub/pyiqa/DISTS_weights-f5e65c96.pth", ], "musiq": ["hub/pyiqa/musiq_koniq_ckpt-e95806b9.pth"], "maniqa_pipal": [ "hub/pyiqa/MANIQA_PIPAL-ae6d356b.pth", "models--timm--vit_base_patch8_224.augreg2_in21k_ft_in1k", ], "topiq_nr": [ "hub/pyiqa/cfanet_nr_koniq_res50-9a73138b.pth", "models--timm--resnet50.a1_in1k", ], "clipiqa": ["hub/clip/RN50.pt"], } def configure_metric_cache(path: Path, metrics: list[str]) -> Path: cache = path.expanduser().resolve() if not cache.is_dir(): raise FileNotFoundError(f"metric cache directory does not exist: {cache}") missing = sorted({ relative for metric in metrics for relative in METRIC_CACHE_FILES.get(metric, []) if not (cache / relative).exists() }) if missing: raise FileNotFoundError( f"metric cache is missing required assets for {metrics}: {missing} ({cache})" ) os.environ["TORCH_HOME"] = str(cache) os.environ["HUGGINGFACE_HUB_CACHE"] = str(cache) return cache def index_images(directory: Path) -> dict[str, Path]: if not directory.is_dir(): raise FileNotFoundError(f"image directory does not exist: {directory}") indexed: dict[str, Path] = {} for path in sorted(directory.iterdir()): if path.suffix.lower() not in IMAGE_SUFFIXES: continue if path.stem in indexed: raise ValueError(f"ambiguous duplicate stem {path.stem!r} in {directory}") try: with Image.open(path) as image: image.convert("RGB").load() except (OSError, ValueError) as error: raise ValueError(f"unreadable image: {path}") from error indexed[path.stem] = path if not indexed: raise ValueError(f"no supported images in {directory}") return indexed def validate_pairs(sr_dir: Path, gt_dir: Path, expected_count: int | None) -> list[tuple[str, Path, Path]]: sr = index_images(sr_dir) gt = index_images(gt_dir) if set(sr) != set(gt): missing = sorted(set(gt) - set(sr))[:10] extra = sorted(set(sr) - set(gt))[:10] raise ValueError(f"SR/GT stems differ: missing_sr={missing}, extra_sr={extra}") if expected_count is not None and len(sr) != expected_count: raise ValueError(f"pair count {len(sr)} != expected {expected_count}") pairs = [] for stem in sorted(sr): with Image.open(sr[stem]) as output, Image.open(gt[stem]) as reference: if output.size != reference.size: raise ValueError( f"dimension mismatch for {stem}: SR={output.size}, GT={reference.size}" ) pairs.append((stem, sr[stem], gt[stem])) return pairs def load_rgb01(path: Path, device: str) -> torch.Tensor: array = np.asarray(Image.open(path).convert("RGB"), dtype=np.float32) / 255.0 return torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0).contiguous().to(device) @torch.no_grad() def score_pairs(pairs: list[tuple[str, Path, Path]], metrics: list[str], device: str) -> tuple[list[dict[str, float | str]], dict[str, float]]: try: import pyiqa except ImportError as error: raise RuntimeError("Install pyiqa from environment/requirements-metrics.txt") from error if not metrics: raise ValueError("at least one metric is required") instances = { name: pyiqa.create_metric(METRIC_SPEC[name][0], device=device, **METRIC_SPEC[name][1]).eval() for name in metrics } rows: list[dict[str, float | str]] = [] values = {name: [] for name in metrics} for index, (stem, sr_path, gt_path) in enumerate(pairs, start=1): sr = load_rgb01(sr_path, device) gt = load_rgb01(gt_path, device) row: dict[str, float | str] = {"image": stem} for name in metrics: _, _, full_reference = METRIC_SPEC[name] score = instances[name](sr.clone(), gt.clone()) if full_reference else instances[name](sr.clone()) value = float(score.item() if torch.is_tensor(score) else score) if not math.isfinite(value): raise ValueError(f"non-finite {name} for {stem}: {value}") row[name] = value values[name].append(value) rows.append(row) if index % 20 == 0 or index == len(pairs): print(f"{index}/{len(pairs)}") return rows, {name: float(np.mean(values[name])) for name in metrics} def write_csv(path: Path, rows: list[dict[str, float | str]], means: dict[str, float], metrics: list[str]) -> None: path.parent.mkdir(parents=True, exist_ok=True) with path.open("w", newline="", encoding="utf-8") as handle: writer = csv.DictWriter(handle, fieldnames=["image", *metrics]) writer.writeheader() writer.writerows(rows) writer.writerow({"image": "MEAN", **means}) def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--sr-dir", type=Path, required=True) parser.add_argument("--gt-dir", type=Path, required=True) parser.add_argument("--out-csv", type=Path, required=True) parser.add_argument("--metrics", default=",".join(DEFAULT_METRICS)) parser.add_argument("--expected-count", type=int) parser.add_argument("--device", default="cuda") parser.add_argument("--metric-cache", type=Path, help="cache root containing Torch and Hugging Face metric weights") parser.add_argument("--offline", action="store_true") args = parser.parse_args() metrics = [name.strip() for name in args.metrics.split(",") if name.strip()] unknown = sorted(set(metrics) - set(METRIC_SPEC)) if unknown: parser.error(f"unknown metrics: {unknown}") if len(set(metrics)) != len(metrics): parser.error("metrics must not contain duplicates") if args.expected_count is not None and args.expected_count < 1: parser.error("--expected-count must be positive") metric_cache = configure_metric_cache(args.metric_cache, metrics) if args.metric_cache else None if args.offline: os.environ["HF_HUB_OFFLINE"] = "1" os.environ["TRANSFORMERS_OFFLINE"] = "1" if args.device.startswith("cuda") and not torch.cuda.is_available(): raise RuntimeError(f"requested {args.device}, but CUDA is not available") pairs = validate_pairs(args.sr_dir.resolve(), args.gt_dir.resolve(), args.expected_count) rows, means = score_pairs(pairs, metrics, args.device) write_csv(args.out_csv, rows, means, metrics) provenance = { "status": "COMPLETE", "sr_dir": str(args.sr_dir.resolve()), "gt_dir": str(args.gt_dir.resolve()), "pair_count": len(pairs), "metrics": metrics, "metric_cache": str(metric_cache) if metric_cache else None, "offline": args.offline, "means": means, "environment": runtime_environment(), "output": file_record(args.out_csv), } write_json(args.out_csv.with_suffix(args.out_csv.suffix + ".provenance.json"), provenance) print(json.dumps({"means": means, "out_csv": str(args.out_csv)}, sort_keys=True)) if __name__ == "__main__": main()