150 lines
7.3 KiB
Python
Executable File
150 lines
7.3 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Pair current eight-case Ref2VA summaries and compute decoded-video SSIM."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
import statistics
|
|
import subprocess
|
|
import tempfile
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
def percentile(values: list[float], q: float) -> float:
|
|
ordered = sorted(values)
|
|
position = (len(ordered) - 1) * q
|
|
lower, upper = math.floor(position), math.ceil(position)
|
|
if lower == upper:
|
|
return ordered[lower]
|
|
return ordered[lower] * (upper - position) + ordered[upper] * (position - lower)
|
|
|
|
|
|
def run(command: list[str]) -> subprocess.CompletedProcess[str]:
|
|
result = subprocess.run(command, text=True, capture_output=True, check=False)
|
|
if result.returncode:
|
|
raise RuntimeError(f"command failed ({result.returncode}): {' '.join(command)}\n{result.stderr[-4000:]}")
|
|
return result
|
|
|
|
|
|
def load_cases(path: Path) -> dict[str, dict[str, Any]]:
|
|
payload = json.loads(path.read_text())
|
|
cases = payload.get("cases") or []
|
|
rows = {str(row["record_id"]): row for row in cases if row.get("success")}
|
|
if len(cases) != 8 or len(rows) != 8:
|
|
raise ValueError(f"expected 8 successful unique cases in {path}, got cases={len(cases)} successful={len(rows)}")
|
|
return rows
|
|
|
|
|
|
def reference_contract(row: dict[str, Any]) -> list[tuple[str, str, str]]:
|
|
return sorted((str(r["label"]), str(r["type"]), str(r["sha256"])) for r in row["references"])
|
|
|
|
|
|
def check_pair(
|
|
reference: dict[str, Any],
|
|
candidate: dict[str, Any],
|
|
*,
|
|
allow_step_mismatch: bool = False,
|
|
) -> None:
|
|
fields = ["record_id", "seed", "short_edge", "aspect_ratio", "duration_seconds"]
|
|
if not allow_step_mismatch:
|
|
fields.append("num_inference_steps")
|
|
mismatches = [f"{field}: {reference.get(field)!r} != {candidate.get(field)!r}" for field in fields if reference.get(field) != candidate.get(field)]
|
|
if reference_contract(reference) != reference_contract(candidate):
|
|
mismatches.append("reference attachment labels/types/hashes differ")
|
|
if mismatches:
|
|
raise ValueError("pair contract mismatch: " + "; ".join(mismatches))
|
|
|
|
|
|
def probe(ffprobe: str, path: Path) -> dict[str, Any]:
|
|
result = run([ffprobe, "-v", "error", "-select_streams", "v:0", "-count_frames",
|
|
"-show_entries", "stream=width,height,r_frame_rate,nb_frames,nb_read_frames", "-of", "json", str(path)])
|
|
stream = json.loads(result.stdout)["streams"][0]
|
|
frames = stream.get("nb_read_frames") or stream.get("nb_frames")
|
|
return {"width": int(stream["width"]), "height": int(stream["height"]),
|
|
"fps": stream["r_frame_rate"], "frames": int(frames)}
|
|
|
|
|
|
def parse_stats(path: Path) -> list[float]:
|
|
values = []
|
|
for line in path.read_text().splitlines():
|
|
tokens = dict(token.split(":", 1) for token in line.split() if ":" in token)
|
|
if "All" in tokens:
|
|
values.append(float(tokens["All"]))
|
|
if not values:
|
|
raise ValueError(f"no SSIM frame metrics in {path}")
|
|
return values
|
|
|
|
|
|
def compare(ffmpeg: str, ffprobe: str, reference: Path, candidate: Path, stats: Path) -> tuple[dict[str, Any], list[float]]:
|
|
ref_probe, cand_probe = probe(ffprobe, reference), probe(ffprobe, candidate)
|
|
if ref_probe != cand_probe:
|
|
raise ValueError(f"decoded video contract mismatch: reference={ref_probe}, candidate={cand_probe}")
|
|
graph = ("[0:v]setpts=PTS-STARTPTS,format=yuv420p[reference];"
|
|
"[1:v]setpts=PTS-STARTPTS,format=yuv420p[candidate];"
|
|
f"[reference][candidate]ssim=stats_file={stats}")
|
|
run([ffmpeg, "-hide_banner", "-nostdin", "-loglevel", "error", "-i", str(reference),
|
|
"-i", str(candidate), "-filter_complex", graph, "-an", "-f", "null", "-"])
|
|
values = parse_stats(stats)
|
|
if len(values) != ref_probe["frames"]:
|
|
raise ValueError(f"frame metric count {len(values)} != decoded frames {ref_probe['frames']}")
|
|
return ref_probe, values
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--reference-summary", type=Path, required=True)
|
|
parser.add_argument("--candidate-summary", type=Path, required=True)
|
|
parser.add_argument("--output-dir", type=Path, required=True)
|
|
parser.add_argument("--threshold", type=float, default=0.90)
|
|
parser.add_argument("--ffmpeg", required=True)
|
|
parser.add_argument("--ffprobe", required=True)
|
|
parser.add_argument(
|
|
"--allow-step-mismatch",
|
|
action="store_true",
|
|
help="Allow the candidate to use a different denoising-step count; all other pairing fields remain strict.",
|
|
)
|
|
args = parser.parse_args()
|
|
reference = load_cases(args.reference_summary)
|
|
candidate = load_cases(args.candidate_summary)
|
|
if set(reference) != set(candidate):
|
|
raise SystemExit("reference and candidate record_id sets differ")
|
|
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
pairs = []
|
|
with tempfile.TemporaryDirectory(prefix="ref2va-ssim-") as temp:
|
|
for index, record_id in enumerate(sorted(reference), 1):
|
|
ref, cand = reference[record_id], candidate[record_id]
|
|
check_pair(ref, cand, allow_step_mismatch=args.allow_step_mismatch)
|
|
ref_path, cand_path = Path(ref["file_path"]), Path(cand["file_path"])
|
|
if not ref_path.is_file() or not cand_path.is_file():
|
|
raise ValueError(f"missing paired video: {ref_path} or {cand_path}")
|
|
video, values = compare(args.ffmpeg, args.ffprobe, ref_path, cand_path, Path(temp) / f"{index}.stats")
|
|
mean = statistics.fmean(values)
|
|
row = {"case_number": cand["case_number"], "record_id": record_id,
|
|
"reference_file": str(ref_path), "candidate_file": str(cand_path), **video,
|
|
"ssim_all_mean": mean, "ssim_all_p10": percentile(values, 0.10),
|
|
"ssim_all_min": min(values), "threshold": args.threshold, "passed": mean >= args.threshold}
|
|
pairs.append(row)
|
|
print(f"[{index}/8] case={cand['case_number']} SSIM={mean:.6f}", flush=True)
|
|
means = [row["ssim_all_mean"] for row in pairs]
|
|
summary = {"metric": "FFmpeg decoded YUV420 SSIM All", "threshold": args.threshold,
|
|
"allow_step_mismatch": args.allow_step_mismatch,
|
|
"videos": len(pairs), "mean_video_ssim": statistics.fmean(means),
|
|
"median_video_ssim": statistics.median(means), "p10_video_ssim": percentile(means, 0.10),
|
|
"min_video_ssim": min(means), "videos_at_or_above_threshold": sum(v >= args.threshold for v in means),
|
|
"pass_rate": sum(v >= args.threshold for v in means) / len(means), "pairs": pairs}
|
|
(args.output_dir / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2) + "\n")
|
|
with (args.output_dir / "per_case.tsv").open("w") as handle:
|
|
columns = ("case_number", "record_id", "ssim_all_mean", "ssim_all_p10", "ssim_all_min", "passed", "candidate_file")
|
|
handle.write("\t".join(columns) + "\n")
|
|
for row in sorted(pairs, key=lambda item: item["case_number"]):
|
|
handle.write("\t".join(str(row[column]) for column in columns) + "\n")
|
|
print(json.dumps({key: value for key, value in summary.items() if key != "pairs"}, ensure_ascii=False, indent=2))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|