Repository: open-mmlab/mmeval (language: python)

Issue:
[Feature] Add End Point Error
## Motivation
As title.

## Modification
1. `mmeval/metrics/__init__.py`
2. `mmeval/metrics/end_point_error.py`
3. `tests/test_metrics/test_end_point_error.py`


Discussion / hints:

Fix patch:
diff --git a/mmeval/metrics/__init__.py b/mmeval/metrics/__init__.py
index a1950769..2be9692b 100644
--- a/mmeval/metrics/__init__.py
+++ b/mmeval/metrics/__init__.py
@@ -1,6 +1,7 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 
 from .accuracy import Accuracy
+from .end_point_error import EndPointError
 from .f_metric import F1Metric
 from .hmean_iou import HmeanIoU
 from .mean_iou import MeanIoU
@@ -8,5 +9,6 @@
 from .voc_map import VOCMeanAP
 
 __all__ = [
-    'Accuracy', 'MeanIoU', 'VOCMeanAP', 'OIDMeanAP', 'HmeanIoU', 'F1Metric'
+    'Accuracy', 'MeanIoU', 'VOCMeanAP', 'OIDMeanAP', 'EndPointError',
+    'F1Metric', 'HmeanIoU'
 ]
diff --git a/mmeval/metrics/end_point_error.py b/mmeval/metrics/end_point_error.py
new file mode 100644
index 00000000..7d440c83
--- /dev/null
+++ b/mmeval/metrics/end_point_error.py
@@ -0,0 +1,168 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+
+import numpy as np
+from typing import TYPE_CHECKING, List, Optional, Sequence, Tuple, overload
+
+from mmeval.core.base_metric import BaseMetric
+from mmeval.core.dispatcher import dispatch
+from mmeval.utils import try_import
+
+if TYPE_CHECKING:
+    import torch
+else:
+    torch = try_import('torch')
+
+
+class EndPointError(BaseMetric):
+    """EndPointError evaluation metric.
+
+    EndPointError is a widely used evaluation metric for optical flow
+    estimation.
+
+    This metric supports 2 kinds of inputs, i.e. `numpy.ndarray` and
+    `torch.Tensor`, and the implementation for the calculation depends on
+    the inputs type.
+
+    Args:
+        **kwargs: Keyword arguments passed to :class:`BaseMetric`.
+
+    Examples:
+
+        >>> from mmeval import EndPointError
+        >>> epe = EndPointError()
+
+    Use NumPy implementation:
+
+        >>> import numpy as np
+        >>> predictions = np.array(
+        ...     [[[10., 5.], [0.1, 3.]],
+        ...     [[3., 15.2], [2.4, 4.5]]])
+        >>> labels = np.array(
+        ...     [[[10.1, 4.8], [10, 3.]],
+        ...     [[6., 10.2], [2.0, 4.1]]])
+        >>> valid_masks = np.array([[1., 1.], [1., 0.]])
+        >>> epe(predictions, labels, valid_masks)
+        {'EPE': 5.318186230865093}
+
+    Use PyTorch implementation:
+
+        >>> import torch
+        >>> predictions = torch.Tensor(
+        ...     [[[10., 5.], [0.1, 3.]],
+        ...     [[3., 15.2], [2.4, 4.5]]])
+        >>> labels = torch.Tensor(
+        ...     [[[10.1, 4.8], [10, 3.]],
+        ...     [[6., 10.2], [2.0, 4.1]]])
+        >>> valid_masks = torch.Tensor([[1., 1.], [1., 0.]])
+        >>> epe(predictions, labels, valid_masks)
+        {'EPE': 5.3181863}
+
+    Accumulate batch:
+
+        >>> for i in range(10):
+        ...     predictions = torch.randn(10, 10, 2)
+        ...     labels = torch.randn(10, 10, 2)
+        ...     epe.add(predictions, labels)
+        >>> epe.compute()  # doctest: +SKIP
+    """
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+
+    def add(self, predictions: Sequence, labels: Sequence, valid_masks: Optional[Sequence] = None) -> None:  # type: ignore # yapf: disable # noqa: E501
+        """Process one batch of predictions and labels.
+
+        Args:
+            predictions (Sequence): Predicted sequence of flow map with
+                shape (H, W, 2).
+            labels (Sequence): The ground truth sequence of flow map with
+                shape (H, W, 2).
+            valid_masks (Sequence): The Sequence of valid mask for labels with
+                shape (H, W). If it is None, this function will automatically
+                generate a map filled with 1. Defaults to None.
+        """
+
+        for idx, (prediction, label) in enumerate(zip(predictions, labels)):
+            assert prediction.shape == label.shape, 'The shape of ' \
+                '`prediction` and `label` should be the same, but got: ' \
+                f'{prediction.shape} and {label.shape}'
+
+            assert prediction.shape[-1] == 2, 'The last dimension of ' \
+                f'`prediction` should be 2, but got: {prediction.shape[-1]}'
+
+            if valid_masks is not None:
+                valid_mask = valid_masks[idx]
+            else:
+                valid_mask = None
+            epe, num_valid_pixel = self.end_point_error_map(
+                prediction, label, valid_mask)
+            self._results.append((epe, num_valid_pixel))
+
+    @overload  # type: ignore
+    @dispatch
+    def end_point_error_map(
+            self,
+            prediction: np.ndarray,
+            label: np.ndarray,
+            valid_mask: Optional[np.ndarray] = None) -> Tuple[np.ndarray, int]:
+        """Calculate end point error map.
+
+        Args:
+            prediction (np.ndarray): Prediction with shape (H, W, 2).
+            label (np.ndarray): Ground truth with shape (H, W, 2).
+            valid_mask (np.ndarray, optional): Valid mask with shape (H, W).
+
+        Returns:
+            Tuple: The mean of end point error and the numbers of valid labels.
+        """
+        if valid_mask is None:
+            valid_mask = np.ones_like(prediction[..., 0])
+        epe_map = np.sqrt(np.sum((prediction - label)**2, axis=-1))
+        val = valid_mask.reshape(-1) >= 0.5
+        epe = epe_map.reshape(-1)[val]
+        return epe.mean(keepdims=True), int(val.sum())
+
+    @dispatch
+    def end_point_error_map(
+            self,
+            prediction: 'torch.Tensor',
+            label: 'torch.Tensor',
+            valid_mask: Optional['torch.Tensor'] = None
+    ) -> Tuple[np.ndarray, int]:
+        """Calculate end point error map.
+
+        Args:
+            prediction (torch.Tensor): Prediction with shape (H, W, 2).
+            label (torch.Tensor): Ground truth with shape (H, W, 2).
+            valid_mask (torch.Tensor, optional): Valid mask with shape (H, W).
+
+        Returns:
+            Tuple: The mean of end point error and the numbers of valid labels.
+        """
+        if valid_mask is None:
+            valid_mask = torch.ones_like(prediction[..., 0])
+        epe_map = torch.sqrt(torch.sum((prediction - label)**2, dim=-1))
+        val = valid_mask.reshape(-1) >= 0.5
+        epe = epe_map.reshape(-1)[val]
+        return epe.mean().cpu().numpy(), int(val.sum())
+
+    def compute_metric(self, results: List[Tuple[np.ndarray, int]]) -> dict:
+        """Compute the EndPointError metric.
+
+        This method would be invoked in `BaseMetric.compute` after distributed
+        synchronization.
+
+        Args:
+            results (List[np.ndarray]): This list has already been synced
+                across all ranks. This is a list of `np.ndarray`, which is
+                the end point error map between the prediction and the label.
+
+        Returns:
+            Dict: The computed metric, with following key:
+
+                - EPE, the mean end point error of all pairs.
+        """
+        epe_overall = sum(res[0] * res[1] for res in results)
+        valid_pixels = sum(res[1] for res in results)
+        metric_results = {'EPE': epe_overall / valid_pixels}
+        return metric_results


Test patch:
diff --git a/tests/test_metrics/test_end_point_error.py b/tests/test_metrics/test_end_point_error.py
new file mode 100644
index 00000000..6224682a
--- /dev/null
+++ b/tests/test_metrics/test_end_point_error.py
@@ -0,0 +1,74 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+
+# yapf: disable
+
+import numpy as np
+import pytest
+
+from mmeval.core.base_metric import BaseMetric
+from mmeval.metrics import EndPointError
+from mmeval.utils import try_import
+
+torch = try_import('torch')
+
+
+def test_metric_interface_numpy():
+    epe = EndPointError()
+    assert isinstance(epe, BaseMetric)
+
+    results = epe(
+        np.random.normal(size=(10, 10, 2)),
+        np.random.normal(size=(10, 10, 2)))
+    assert isinstance(results, dict)
+
+
+@pytest.mark.skipif(torch is None, reason='PyTorch is not available!')
+def test_metric_interface_torch():
+    epe = EndPointError()
+    assert isinstance(epe, BaseMetric)
+
+    results = epe(torch.randn(10, 10, 2), torch.randn(10, 10, 2))
+    assert isinstance(results, dict)
+
+
+@pytest.mark.parametrize(
+    argnames=['predictions', 'labels', 'valid_masks'],
+    argvalues=[
+        (
+            np.array([[[10., 5.], [0.1, 3.]], [[3., 15.2], [2.4, 4.5]]]),
+            np.array([[[10.1, 4.8], [10, 3.]], [[6., 10.2], [2.0, 4.1]]]),
+            np.array([[1., 1.], [1., 0.3]])
+        )
+    ]
+)
+def test_metric_accurate(predictions, labels, valid_masks):
+    epe_target = np.linalg.norm((predictions - labels), ord=2, axis=-1)
+    epe_target = (epe_target[0].sum() + epe_target[1][0]) / 3
+
+    epe_metric = EndPointError()
+    epe_result = epe_metric(predictions, labels, valid_masks)
+    np.testing.assert_allclose(epe_result['EPE'], epe_target)
+
+
+@pytest.mark.skipif(torch is None, reason='PyTorch is not available!')
+def test_metamorphic_numpy_pytorch():
+    """Metamorphic testing for NumPy and PyTorch implementation."""
+
+    epe = EndPointError()
+    predictions = np.random.normal(size=(10, 10, 2))
+    labels = np.random.normal(size=(10, 10, 2))
+    np_results = epe(predictions, labels)
+
+    predictions = torch.from_numpy(predictions)
+    labels = torch.from_numpy(labels)
+    torch_results = epe(predictions, labels)
+
+    assert np_results.keys() == torch_results.keys()
+
+    for key in np_results:
+        np.testing.assert_allclose(
+            np_results[key], torch_results[key], rtol=1e-06)
+
+
+if __name__ == '__main__':
+    pytest.main([__file__, '-vv', '--capture=no'])