[Feature] Add MattingMSE metric (#71)

* add mattingmse

* modify config.yml

* Modify some details

* Add test_matting_mse.py and simplify matting_mse.py.

* Synchronize config.yml in the main branch

* improve the document

* Fine-tune docstring

* Apply suggestions from code review

Modify format statement

Co-authored-by: yancong <32220263+ice-tong@users.noreply.github.com>

* update

* Update mmeval/metrics/matting_mse.py

* Update matting_mse.py

update

* Change metrics to lowercase

Co-authored-by: yancong <32220263+ice-tong@users.noreply.github.com>
Co-authored-by: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com>

Author: 3 people

docs/en/api/metrics.rst

@@ -44,3 +44,4 @@ Metrics
    BLEU
    SAD
    GradientError
+   MattingMSE

docs/zh_cn/api/metrics.rst

@@ -44,3 +44,4 @@ Metrics
    BLEU
    SAD
    GradientError
+   MattingMSE

mmeval/metrics/__init__.py

@@ -9,6 +9,7 @@
 from .gradient_error import GradientError
 from .hmean_iou import HmeanIoU
 from .mae import MAE
+from .matting_mse import MattingMSE
 from .mean_iou import MeanIoU
 from .mse import MSE
 from .multi_label import AveragePrecision, MultiLabelMetric
@@ -27,5 +28,6 @@
     'F1Metric', 'HmeanIoU', 'SingleLabelMetric', 'COCODetectionMetric',
     'PCKAccuracy', 'MpiiPCKAccuracy', 'JhmdbPCKAccuracy', 'ProposalRecall',
     'PSNR', 'MAE', 'MSE', 'SSIM', 'SNR', 'MultiLabelMetric',
-    'AveragePrecision', 'AVAMeanAP', 'BLEU', 'SAD', 'GradientError'
+    'AveragePrecision', 'AVAMeanAP', 'BLEU', 'SAD', 'GradientError',
+    'MattingMSE'
 ]

mmeval/metrics/matting_mse.py

@@ -0,0 +1,87 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import numpy as np
+from typing import Dict, List, Sequence
+
+from mmeval.core import BaseMetric
+
+
+class MattingMSE(BaseMetric):
+    """Mean Squared Error metric for image matting.
+
+    This metric computes the per-pixel squared error average across all
+    pixels.
+    i.e. mean((a-b)^2)
+
+    Args:
+        **kwargs: Keyword parameters passed to :class:`BaseMetric`.
+
+    Note:
+        The current implementation assumes the image / alpha / trimap
+        a numpy array with pixel values ranging from 0 to 255.
+
+        The pred_alpha should be masked by trimap before passing
+        into this metric.
+
+        The trimap is the most commonly used prior knowledge. As the
+        name implies, trimap is a ternary graph and each pixel
+        takes one of {0, 128, 255}, representing the foreground, the
+        unknown and the background respectively.
+
+    Examples:
+
+        >>> from mmeval import MattingMSE
+        >>> import numpy as np
+        >>>
+        >>> matting_mse = MattingMSE()
+        >>> pred_alpha = np.zeros((32, 32), dtype=np.uint8)
+        >>> gt_alpha = np.ones((32, 32), dtype=np.uint8) * 255
+        >>> trimap = np.zeros((32, 32), dtype=np.uint8)
+        >>> trimap[:16, :16] = 128
+        >>> trimap[16:, 16:] = 255
+        >>> matting_mse(pred_alpha, gt_alpha, trimap)  # doctest: +ELLIPSIS
+        {'matting_mse': ...}
+    """
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+
+    def add(self, pred_alphas: Sequence[np.ndarray], gt_alphas: Sequence[np.ndarray], trimaps: Sequence[np.ndarray]) -> None:  # type: ignore # yapf: disable # noqa: E501
+        """Add MattingMSE score of batch to ``self._results``
+
+        Args:
+            pred_alphas (Sequence[np.ndarray]): Predict the probability
+                that pixels belong to the foreground.
+            gt_alphas (Sequence[np.ndarray]): Probability that the actual
+                pixel belongs to the foreground.
+            trimaps (Sequence[np.ndarray]): Broadly speaking, the trimap
+                consists of foreground and unknown region.
+        """
+
+        for pred_alpha, gt_alpha, trimap in zip(pred_alphas, gt_alphas,
+                                                trimaps):
+            assert pred_alpha.shape == gt_alpha.shape, 'The shape of ' \
+                '`pred_alpha` and `gt_alpha` should be the same, but got: ' \
+                f'{pred_alpha.shape} and {gt_alpha.shape}'
+
+            weight_sum = (trimap == 128).sum()
+            if weight_sum != 0:
+                mse_result = ((pred_alpha - gt_alpha)**2).sum() / weight_sum
+            else:
+                mse_result = 0
+
+            self._results.append(mse_result)
+
+    def compute_metric(self, results: List) -> Dict[str, float]:
+        """Compute the MattingMSE metric.
+
+        Args:
+            results (List): A list that consisting the MattingMSE score.
+                This list has already been synced across all ranks.
+
+        Returns:
+            Dict[str, float]: The computed MattingMSE metric.
+            The keys are the names of the metrics,
+            and the values are corresponding results.
+        """
+
+        return {'matting_mse': float(np.array(results).mean())}

tests/test_metrics/test_matting_mse.py

@@ -0,0 +1,17 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import numpy as np
+
+from mmeval.metrics import MattingMSE
+
+
+def test_matting_mse():
+    pred_alpha = np.zeros((32, 32), dtype=np.uint8)
+    gt_alpha = np.ones((32, 32), dtype=np.uint8) * 255
+    trimap = np.zeros((32, 32), dtype=np.uint8)
+    trimap[:16, :16] = 128
+    trimap[16:, 16:] = 255
+
+    matting_mse = MattingMSE()
+    metric_results = matting_mse(pred_alpha, gt_alpha, trimap)
+    assert isinstance(metric_results, dict)
+    np.testing.assert_almost_equal(metric_results['matting_mse'], 1.0)