137 lines
4.7 KiB
Python
137 lines
4.7 KiB
Python
"""End-to-end tests for single-replica eval through the public API.
|
|
|
|
Runs the lightweight pixel-space metrics — ``common.psnr`` and
|
|
``common.ssim`` — under both shapes that real callers use:
|
|
|
|
* one-shot ``evaluate(video=..., reference=...)`` (the helper in
|
|
``fastvideo.eval.api``);
|
|
* a long-lived ``Evaluator``, called once per sample;
|
|
* a long-lived ``Evaluator``, called with a list of sample dicts to
|
|
fan out (``samples=[...]``).
|
|
|
|
GPU-only metrics live in separate test modules / classes; everything
|
|
here runs on CPU so the suite stays cheap to invoke.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from fastvideo.eval import MetricResult, create_evaluator, evaluate
|
|
|
|
|
|
# Tight resolution + frame count so CPU SSIM stays under a couple of seconds.
|
|
_T, _C, _H, _W = 6, 3, 32, 32
|
|
|
|
|
|
@pytest.fixture
|
|
def gen_ref():
|
|
"""Reproducible (gen, ref) pair shaped (T, C, H, W)."""
|
|
torch.manual_seed(0)
|
|
gen = torch.rand(_T, _C, _H, _W)
|
|
ref = torch.rand(_T, _C, _H, _W)
|
|
return gen, ref
|
|
|
|
|
|
@pytest.fixture
|
|
def evaluator():
|
|
ev = create_evaluator(
|
|
metrics=["common.psnr", "common.ssim"],
|
|
device="cpu",
|
|
)
|
|
yield ev
|
|
ev.shutdown()
|
|
|
|
|
|
def _assert_well_formed(result: MetricResult, name: str) -> None:
|
|
assert isinstance(result, MetricResult)
|
|
assert result.name == name
|
|
assert result.score is not None
|
|
assert isinstance(result.score, float)
|
|
# PSNR / SSIM both populate per-frame details.
|
|
assert "per_frame" in result.details
|
|
assert len(result.details["per_frame"]) == _T
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# One-shot helper
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_evaluate_one_shot_returns_dict_of_metric_results(gen_ref):
|
|
gen, ref = gen_ref
|
|
out = evaluate(generated=gen, reference=ref,
|
|
metrics=["common.psnr", "common.ssim"], device="cpu")
|
|
assert isinstance(out, dict)
|
|
assert set(out.keys()) == {"common.psnr", "common.ssim"}
|
|
_assert_well_formed(out["common.psnr"], "common.psnr")
|
|
_assert_well_formed(out["common.ssim"], "common.ssim")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Long-lived Evaluator, single-sample form
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_evaluator_single_sample_returns_dict(evaluator, gen_ref):
|
|
gen, ref = gen_ref
|
|
out = evaluator.evaluate(video=gen, reference=ref)
|
|
assert isinstance(out, dict)
|
|
assert set(out.keys()) == {"common.psnr", "common.ssim"}
|
|
for name, mr in out.items():
|
|
_assert_well_formed(mr, name)
|
|
|
|
|
|
def test_evaluator_accepts_legacy_5d_input(evaluator, gen_ref):
|
|
"""Callers that still pass ``(1, T, C, H, W)`` should get unwrapped."""
|
|
gen, ref = gen_ref
|
|
out = evaluator.evaluate(video=gen.unsqueeze(0), reference=ref.unsqueeze(0))
|
|
_assert_well_formed(out["common.psnr"], "common.psnr")
|
|
|
|
|
|
def test_evaluator_score_is_deterministic(evaluator, gen_ref):
|
|
gen, ref = gen_ref
|
|
a = evaluator.evaluate(video=gen, reference=ref)
|
|
b = evaluator.evaluate(video=gen.clone(), reference=ref.clone())
|
|
assert a["common.psnr"].score == pytest.approx(b["common.psnr"].score)
|
|
assert a["common.ssim"].score == pytest.approx(b["common.ssim"].score)
|
|
|
|
|
|
def test_evaluator_psnr_identical_videos_is_high(evaluator, gen_ref):
|
|
"""PSNR(x, x) is unbounded above; with our clamp it caps near 100 dB."""
|
|
gen, _ = gen_ref
|
|
out = evaluator.evaluate(video=gen, reference=gen)
|
|
assert out["common.psnr"].score > 50.0
|
|
assert out["common.ssim"].score == pytest.approx(1.0, abs=1e-5)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Long-lived Evaluator, list (fan-out) form
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_evaluator_samples_list_preserves_input_order(evaluator):
|
|
"""When ``samples=[...]`` is passed, results must come back per sample."""
|
|
torch.manual_seed(1)
|
|
samples = []
|
|
for i in range(4):
|
|
# Vary the reference enough that scores differ across rows.
|
|
gen = torch.rand(_T, _C, _H, _W)
|
|
ref = gen + 0.01 * (i + 1) * torch.rand_like(gen)
|
|
samples.append({"video": gen, "reference": ref})
|
|
|
|
out = evaluator.evaluate(samples=samples)
|
|
assert isinstance(out, list)
|
|
assert len(out) == len(samples)
|
|
for row in out:
|
|
assert set(row.keys()) == {"common.psnr", "common.ssim"}
|
|
_assert_well_formed(row["common.psnr"], "common.psnr")
|
|
|
|
# Re-running should give bit-identical scores (no nondeterministic
|
|
# scheduling effects under single-GPU dispatch).
|
|
out2 = evaluator.evaluate(samples=samples)
|
|
for a, b in zip(out, out2):
|
|
assert a["common.psnr"].score == pytest.approx(b["common.psnr"].score)
|
|
|
|
|