| |
| """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() |
|
|
