| """Validate complete benchmark coverage and pool saved frame metrics/FID features.""" |
|
|
| import argparse |
| import json |
| from pathlib import Path |
| import sys |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parent / "metric_vendor")) |
|
|
| import numpy as np |
| import torch |
| import torchmetrics |
| from torchmetrics.image.fid import _compute_fid |
|
|
| from protocol import case_key, case_seed, select_cases |
| from score_case import metric_ranges |
|
|
|
|
| def feature_statistics(case_dirs, name, start, end, device): |
| count, total, product = 0, None, None |
| for directory in case_dirs: |
| feature = torch.from_numpy(np.load(directory / f"{name}_features.npy")[start:end]).to(device, torch.float64) |
| if total is None: |
| total = torch.zeros(feature.shape[1], dtype=torch.float64, device=device) |
| product = torch.zeros((feature.shape[1], feature.shape[1]), dtype=torch.float64, device=device) |
| total += feature.sum(0) |
| product += feature.T @ feature |
| count += len(feature) |
| mean = total / count |
| covariance = (product - total[:, None] @ total[None, :] / count) / (count - 1) |
| return mean, covariance, count |
|
|
|
|
| def pooled_fid(case_dirs, start, end, device): |
| real_mean, real_cov, count = feature_statistics(case_dirs, "real", start, end, device) |
| fake_mean, fake_cov, _ = feature_statistics(case_dirs, "fake", start, end, device) |
| |
| |
| value = _compute_fid(real_mean, real_cov, fake_mean, fake_cov).float() |
| return float(value.cpu()), count |
|
|
|
|
| def validate_coverage(output, rows, dataset, seed, frames): |
| directories = [Path(output) / "cases" / case_key(row) for row in rows] |
| actual = {path.name for path in (Path(output) / "cases").iterdir() if path.is_dir()} |
| expected = {directory.name for directory in directories} |
| errors, summaries = [], [] |
| if len(expected) != len(rows): |
| errors.append("Duplicate case identities in the selected manifest") |
| if actual != expected: |
| errors.append(f"Case coverage mismatch: missing={sorted(expected - actual)}, extra={sorted(actual - expected)}") |
| if torchmetrics.__version__ != "0.11.4": |
| errors.append(f"Expected torchmetrics0.11.4, found {torchmetrics.__version__}") |
| for row, directory in zip(rows, directories): |
| required = [directory / name for name in ("metrics.json", "frame_metrics.npz", "real_features.npy", "fake_features.npy", "prediction.mp4", "reference.mp4")] |
| missing = [str(path) for path in required if not path.is_file()] |
| if missing: |
| errors.append(f"Incomplete case {directory.name}: {missing}") |
| continue |
| summary = json.loads((directory / "metrics.json").read_text()) |
| summaries.append(summary) |
| wanted = (dict(row), dataset, case_seed(row, seed), frames) |
| found = (summary["case"], summary["dataset"], summary["sample_seed"], summary["frames"]) |
| if found != wanted: |
| errors.append(f"Case identity/dataset/seed/horizon mismatch: {directory.name}") |
| with np.load(directory / "frame_metrics.npz") as metrics: |
| names = {"mse", "psnr", "lpips", "ssim"} if dataset == "re10k" else {"mse", "psnr", "lpips"} |
| if set(metrics.files) != names or any(metrics[name].shape != (frames,) or not np.isfinite(metrics[name]).all() for name in metrics.files): |
| errors.append(f"Invalid per-frame metric population: {directory.name}") |
| for name in ("real", "fake"): |
| feature = np.load(directory / f"{name}_features.npy", mmap_mode="r") |
| if feature.shape != (frames, 2048) or not np.isfinite(feature).all(): |
| errors.append(f"Invalid {name} FID feature population: {directory.name}") |
| contracts = [{key: summary[key] for key in ("metric_resolution", "torchmetrics", "lpips", "psnr_epsilon", "fid", "video")} for summary in summaries] |
| if contracts and any(contract != contracts[0] for contract in contracts): |
| errors.append("Metric contracts differ between cases") |
| if errors: |
| raise ValueError("\n".join(errors)) |
| return directories, contracts[0] |
|
|
|
|
| def aggregate(output, manifest, dataset, seed=0, scope=None, fid_manifest=None, |
| expected_frames=None, limit=None, device="cuda"): |
| frames = expected_frames or (300 if dataset == "re10k" else 500) |
| rows = select_cases(manifest, scope, limit=limit) |
| directories, contract = validate_coverage(output, rows, dataset, seed, frames) |
| ranges = metric_ranges(dataset, frames) |
| report = {"dataset": dataset, "seed": seed, "manifest": str(manifest), "scope": scope, |
| "cases": len(rows), "generated_frames_per_case": frames, |
| "complete_full_horizon": frames == (300 if dataset == "re10k" else 500), |
| "limited_selection": limit is not None, "metric_contract": contract, |
| "frame_metrics": {}, "fid": {}} |
| all_values = {} |
| for directory in directories: |
| with np.load(directory / "frame_metrics.npz") as values: |
| for name in values.files: |
| all_values.setdefault(name, []).append(values[name]) |
| all_values = {name: np.stack(parts).astype(np.float64) for name, parts in all_values.items()} |
| for name, start, end in ranges: |
| report["frame_metrics"][name] = {metric: float(value[:, start:end].mean()) for metric, value in all_values.items()} |
| populations = {"all_cases": directories} |
| if fid_manifest is not None: |
| subset = select_cases(fid_manifest) |
| requested = {case_key(row) for row in subset} |
| available = {directory.name for directory in directories} |
| if not requested <= available: |
| raise ValueError(f"FID manifest cases absent from evaluated population: {sorted(requested - available)}") |
| populations["official_fid50"] = [Path(output) / "cases" / case_key(row) for row in subset] |
| for population, selected in populations.items(): |
| report["fid"][population] = {} |
| for name, start, end in ranges: |
| value, count = pooled_fid(selected, start, end, torch.device(device)) |
| report["fid"][population][name] = {"fid": value, "frames": count, "cases": len(selected)} |
| Path(output, "aggregate.json").write_text(json.dumps(report, indent=2) + "\n") |
| return report |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument("--manifest", type=Path, required=True) |
| parser.add_argument("--dataset", choices=["minecraft", "re10k"], required=True) |
| parser.add_argument("--scope") |
| parser.add_argument("--seed", type=int, default=0) |
| parser.add_argument("--fid-manifest", type=Path) |
| parser.add_argument("--expected-frames", type=int) |
| parser.add_argument("--limit", type=int) |
| parser.add_argument("--device", default="cuda") |
| args = parser.parse_args() |
| result = aggregate(args.output, args.manifest, args.dataset, args.seed, args.scope, |
| args.fid_manifest, args.expected_frames, args.limit, args.device) |
| print(json.dumps({"dataset": result["dataset"], "cases": result["cases"], |
| "overall": result["frame_metrics"]["overall"]}), flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
