rpx_benchmark.metrics.registry

  1"""Core of the metric plugin system.
  2
  3Contains the :class:`MetricCalculator` base class, the task→calculators
  4registry, the ``@register_metric`` decorator, and a thin
  5:class:`MetricSuite` facade that the runner uses.
  6
  7Most library users should not touch this module directly — instead
  8import from :mod:`rpx_benchmark.metrics`.
  9"""
 10
 11from __future__ import annotations
 12
 13from abc import ABC, abstractmethod
 14from collections import defaultdict
 15from dataclasses import dataclass
 16from typing import Any, Callable, Dict, List, Type
 17
 18import numpy as np
 19
 20from ..api import TaskType
 21from ..exceptions import MetricError
 22from ..logging_utils import get_logger
 23
 24log = get_logger(__name__)
 25
 26
 27# --------------------------------------------------------------------------- #
 28# Abstract base
 29# --------------------------------------------------------------------------- #
 30
 31
 32class MetricCalculator(ABC):
 33    """Abstract base for a family of metrics bound to a single task.
 34
 35    Subclasses are lightweight, stateless objects. The ``compute``
 36    method accepts one prediction and its ground truth and returns a
 37    dict mapping metric name → scalar float.
 38
 39    Subclasses must set ``name`` (human-readable identifier used for
 40    registration) and implement :meth:`compute`. Subclasses should
 41    take their own hyperparameters via ``__init__`` when needed.
 42
 43    Examples
 44    --------
 45    >>> from rpx_benchmark.metrics import MetricCalculator, register_metric
 46    >>> from rpx_benchmark.api import TaskType
 47    >>> @register_metric(TaskType.MONOCULAR_DEPTH)
 48    ... class MeanPred(MetricCalculator):
 49    ...     name = "mean_pred"
 50    ...     def compute(self, prediction, ground_truth):
 51    ...         return {"mean_pred": float(prediction.depth_map.mean())}
 52    """
 53
 54    name: str = ""
 55
 56    @abstractmethod
 57    def compute(self, prediction: Any, ground_truth: Any) -> Dict[str, float]:
 58        """Return a dict of metric name → scalar for one sample.
 59
 60        Parameters
 61        ----------
 62        prediction : Any
 63            Task-specific Prediction dataclass
 64            (e.g. :class:`DepthPrediction`). The concrete type is
 65            whatever the calculator is designed for.
 66        ground_truth : Any
 67            Task-specific GroundTruth dataclass
 68            (e.g. :class:`DepthGroundTruth`).
 69
 70        Returns
 71        -------
 72        dict[str, float]
 73            Metric values. Keys should be short snake_case names;
 74            values must be numeric so the runner can average them.
 75
 76        Raises
 77        ------
 78        MetricError
 79            When inputs have the wrong shape/type/content.
 80        """
 81
 82
 83# --------------------------------------------------------------------------- #
 84# Registry state
 85# --------------------------------------------------------------------------- #
 86
 87_CALCULATORS: Dict[TaskType, List[MetricCalculator]] = defaultdict(list)
 88
 89
 90def register_metric(
 91    task: TaskType,
 92) -> Callable[[Type[MetricCalculator]], Type[MetricCalculator]]:
 93    """Class decorator registering a calculator under a task.
 94
 95    Parameters
 96    ----------
 97    task : TaskType
 98        Task the calculator belongs to. Multiple calculators can be
 99        registered against the same task and will be composed at
100        evaluation time.
101
102    Returns
103    -------
104    Callable
105        The decorator; when applied to a class it instantiates it
106        once with no arguments and appends to the registry.
107
108    Raises
109    ------
110    MetricError
111        If ``task`` is not a :class:`TaskType`, or if the decorated
112        object is not a :class:`MetricCalculator` subclass.
113
114    Examples
115    --------
116    >>> from rpx_benchmark.api import TaskType
117    >>> from rpx_benchmark.metrics import MetricCalculator, register_metric
118    >>> @register_metric(TaskType.MONOCULAR_DEPTH)
119    ... class MyMetric(MetricCalculator):
120    ...     name = "my_metric"
121    ...     def compute(self, prediction, ground_truth):
122    ...         return {"my_metric": 0.0}
123    """
124    if not isinstance(task, TaskType):
125        raise MetricError(
126            f"register_metric expected a TaskType, got {type(task).__name__}",
127            hint="Use one of the members of rpx_benchmark.api.TaskType.",
128        )
129
130    def decorator(cls: Type[MetricCalculator]) -> Type[MetricCalculator]:
131        if not (isinstance(cls, type) and issubclass(cls, MetricCalculator)):
132            raise MetricError(
133                f"@register_metric can only decorate MetricCalculator subclasses; got {cls!r}",
134            )
135        instance = cls()
136        _CALCULATORS[task].append(instance)
137        log.debug("registered metric %s for task %s", instance.name or cls.__name__, task.value)
138        return cls
139
140    return decorator
141
142
143def unregister_metric(task: TaskType, name: str) -> bool:
144    """Remove a previously registered calculator by its ``name`` field.
145
146    Parameters
147    ----------
148    task : TaskType
149    name : str
150        Name attribute of the calculator to remove.
151
152    Returns
153    -------
154    bool
155        True if a calculator was removed; False if no match was found.
156    """
157    before = len(_CALCULATORS.get(task, []))
158    _CALCULATORS[task] = [c for c in _CALCULATORS.get(task, []) if c.name != name]
159    after = len(_CALCULATORS[task])
160    return after < before
161
162
163def clear_registry() -> None:
164    """Remove every registered calculator. Primarily for tests."""
165    _CALCULATORS.clear()
166
167
168def get_calculators(task: TaskType) -> List[MetricCalculator]:
169    """Return the calculators registered for ``task`` (empty list if none)."""
170    return list(_CALCULATORS.get(task, []))
171
172
173def available_metrics() -> Dict[TaskType, List[str]]:
174    """List the registered metric names grouped by task.
175
176    Returns
177    -------
178    dict
179        ``{TaskType: [calculator_name, ...]}``. Only tasks with at
180        least one registered calculator appear.
181    """
182    return {task: [c.name for c in calcs] for task, calcs in _CALCULATORS.items() if calcs}
183
184
185def compute_metrics(
186    task: TaskType,
187    prediction: Any,
188    ground_truth: Any,
189) -> Dict[str, float]:
190    """Run every registered calculator for ``task`` and merge outputs.
191
192    Parameters
193    ----------
194    task : TaskType
195    prediction : Any
196    ground_truth : Any
197
198    Returns
199    -------
200    dict[str, float]
201        Union of all calculators' output dicts. Later calculators may
202        overwrite earlier ones if they emit the same key.
203
204    Raises
205    ------
206    MetricError
207        If no calculators are registered for ``task``.
208    """
209    calcs = _CALCULATORS.get(task)
210    if not calcs:
211        raise MetricError(
212            f"No metric calculator registered for task {task.value!r}.",
213            hint=(
214                "Either import the built-in calculators (e.g. "
215                "`import rpx_benchmark.metrics`) or register a custom "
216                "one via @register_metric."
217            ),
218        )
219    merged: Dict[str, float] = {}
220    for calc in calcs:
221        merged.update(calc.compute(prediction, ground_truth))
222    return merged
223
224
225# --------------------------------------------------------------------------- #
226# Facade: MetricSuite
227# --------------------------------------------------------------------------- #
228
229
230@dataclass
231class BenchmarkResult:
232    """Outcome of running a :class:`BenchmarkRunner` against a dataset.
233
234    Attributes
235    ----------
236    task : TaskType
237        Which task was evaluated.
238    per_sample : list of dict
239        One dict per sample. Each dict mixes metric keys (numeric) and
240        metadata keys (``id``, ``phase``, ``difficulty``, ``scene``)
241        that :class:`MetricSuite.aggregate` silently skips when
242        computing means.
243    aggregated : dict[str, float]
244        Mean over the numeric metric keys in :attr:`per_sample`.
245    num_samples : int
246    """
247
248    task: TaskType
249    per_sample: List[Dict[str, Any]]
250    aggregated: Dict[str, float]
251    num_samples: int
252
253
254class MetricSuite:
255    """Thin wrapper around the metric registry used by the runner.
256
257    Kept as a class rather than a function because historical API
258    expects ``MetricSuite.for_task(...).evaluate(pred, gt)``. New code
259    can call :func:`compute_metrics` directly.
260    """
261
262    def __init__(self, task: TaskType) -> None:
263        self.task = task
264
265    @classmethod
266    def for_task(cls, task: TaskType) -> "MetricSuite":
267        """Create a suite for the given task.
268
269        Raises
270        ------
271        MetricError
272            If no calculators are registered for ``task``.
273        """
274        if not get_calculators(task):
275            raise MetricError(
276                f"No metric calculator registered for task {task.value!r}.",
277                hint="Import rpx_benchmark.metrics to trigger built-ins.",
278            )
279        return cls(task=task)
280
281    def evaluate(self, prediction: Any, ground_truth: Any) -> Dict[str, float]:
282        """Run every registered calculator and return merged results.
283
284        Raises
285        ------
286        MetricError
287            Propagated from individual calculators when inputs are
288            shape-mismatched or wrong-typed.
289        """
290        return compute_metrics(self.task, prediction, ground_truth)
291
292    def aggregate(self, per_sample: List[Dict[str, Any]]) -> Dict[str, float]:
293        """Mean over numeric metric keys; non-numeric metadata is skipped.
294
295        Parameters
296        ----------
297        per_sample : list of dict
298            Per-sample rows. May contain metric floats and metadata
299            strings/enums in the same dict.
300
301        Returns
302        -------
303        dict[str, float]
304            One float per numeric key. Empty dict if ``per_sample`` is
305            empty or contains no numeric values.
306        """
307        if not per_sample:
308            return {}
309        numeric_keys = [
310            k
311            for k, v in per_sample[0].items()
312            if isinstance(v, (int, float)) and not isinstance(v, bool)
313        ]
314        out: Dict[str, float] = {}
315        for k in numeric_keys:
316            vals = [
317                m[k]
318                for m in per_sample
319                if isinstance(m.get(k), (int, float)) and not isinstance(m.get(k), bool)
320            ]
321            if vals:
322                out[k] = float(np.mean(vals))
323        return out
324
325    def build_result(self, per_sample: List[Dict[str, Any]]) -> "BenchmarkResult":
326        """Convenience: wrap ``per_sample`` and its aggregate in a
327        :class:`BenchmarkResult`.
328        """
329        return BenchmarkResult(
330            task=self.task,
331            per_sample=per_sample,
332            aggregated=self.aggregate(per_sample),
333            num_samples=len(per_sample),
334        )
335
336
337__all__ = [
338    "MetricCalculator",
339    "MetricSuite",
340    "BenchmarkResult",
341    "available_metrics",
342    "clear_registry",
343    "compute_metrics",
344    "get_calculators",
345    "register_metric",
346    "unregister_metric",
347]