92 lines
3.6 KiB
Python
92 lines
3.6 KiB
Python
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
from fastvideo.eval.types import MetricResult
|
|
|
|
|
|
class BaseMetric:
|
|
"""Abstract base class for all eval metrics.
|
|
|
|
Two execution shapes:
|
|
|
|
* **Per-sample** (``is_set_metric=False``, default) — implement
|
|
:meth:`compute`. The Evaluator calls it once per input sample and
|
|
returns one :class:`MetricResult` per sample.
|
|
|
|
* **Set-vs-set** (``is_set_metric=True``) — implement
|
|
:meth:`accumulate` (called once per sample to buffer features)
|
|
and :meth:`finalize` (called once after all samples to compute
|
|
the corpus-level result). Use :meth:`reset` to clear buffers and
|
|
:meth:`merge_from` to fold multi-GPU per-worker state together.
|
|
|
|
Optionally override :meth:`setup` to eagerly load models. Metrics
|
|
that chunk along the time dim for memory hardcode their own chunk
|
|
size in ``__init__`` (see ``optical_flow`` for the canonical
|
|
example). Eval always processes one video per
|
|
:meth:`Evaluator.evaluate` call; ``compute`` / ``accumulate``
|
|
receive a single sample, not a batch.
|
|
"""
|
|
|
|
name: str = ""
|
|
requires_reference: bool = True
|
|
higher_is_better: bool = True
|
|
dependencies: list[str] = []
|
|
needs_gpu: bool = False
|
|
backbone: str | None = None
|
|
is_set_metric: bool = False
|
|
|
|
# Default time-dim chunk size for metrics that batch internally over
|
|
# frames or frame-pairs. Override in subclass __init__ if needed
|
|
# (see ``optical_flow``, ``motion_smoothness``, ``dynamic_degree``).
|
|
_chunk_size: int | None = None
|
|
|
|
def __init__(self) -> None:
|
|
self._device: torch.device = torch.device("cpu")
|
|
|
|
@property
|
|
def device(self) -> torch.device:
|
|
return self._device
|
|
|
|
def to(self, device: str | torch.device) -> BaseMetric:
|
|
"""Move metric (and its internal models) to *device*."""
|
|
self._device = torch.device(device)
|
|
return self
|
|
|
|
def setup(self) -> None: # noqa: B027 - intentionally optional override
|
|
"""Eagerly load models. Called once by :class:`EvalWorker`.
|
|
|
|
Default is a no-op; metrics with no eager state (pixel math,
|
|
closed-form ops) inherit this. Override only if your metric
|
|
needs to load weights.
|
|
"""
|
|
|
|
def _skip(self, sample: dict, reason: str) -> MetricResult:
|
|
"""Return a skipped result (``score=None`` + reason in details)."""
|
|
return MetricResult(name=self.name, score=None, details={"skipped": reason})
|
|
|
|
def compute(self, sample: dict) -> MetricResult:
|
|
"""Per-sample metrics: compute the score for one sample.
|
|
|
|
``sample["video"]`` is ``(T, C, H, W)`` float in ``[0, 1]``.
|
|
``sample["reference"]`` (if used) has the same shape. Return
|
|
``self._skip(sample, reason)`` for missing inputs.
|
|
"""
|
|
raise NotImplementedError(f"{type(self).__name__}.compute is not implemented")
|
|
|
|
# --- set-vs-set protocol (only invoked when is_set_metric=True) ---
|
|
|
|
def reset(self) -> None: # noqa: B027 - intentionally optional override
|
|
"""Clear accumulator state at the start of each evaluate() call."""
|
|
|
|
def accumulate(self, sample: dict) -> None:
|
|
"""Buffer per-sample features for a corpus-level metric."""
|
|
raise NotImplementedError(f"{type(self).__name__}.accumulate is not implemented")
|
|
|
|
def finalize(self) -> MetricResult:
|
|
"""Compute the corpus-level result from buffered state."""
|
|
raise NotImplementedError(f"{type(self).__name__}.finalize is not implemented")
|
|
|
|
def merge_from(self, other: BaseMetric) -> None: # noqa: B027 - intentionally optional override
|
|
"""Multi-GPU: fold another worker's accumulator state into this one."""
|