Add new Audio metric DNSMOS

CHANGELOG.md

@@ -12,6 +12,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Added a new audio metric `DNSMOS` ([#2525](https://github.com/PyTorchLightning/metrics/pull/2525))
+
+
 - Added `MetricInputTransformer` wrapper ([#2392](https://github.com/Lightning-AI/torchmetrics/pull/2392))
 
 

docs/source/audio/deep_noise_suppression_mean_opinion_score.rst

@@ -0,0 +1,21 @@
+.. customcarditem::
+   :header: Deep Noise Suppression Mean Opinion Score (DNSMOS)
+   :image: https://pl-flash-data.s3.amazonaws.com/assets/thumbnails/audio_classification.svg
+   :tags: Audio
+
+.. include:: ../links.rst
+
+##################################################
+Deep Noise Suppression Mean Opinion Score (DNSMOS)
+##################################################
+
+Module Interface
+________________
+
+.. autoclass:: torchmetrics.audio.dnsmos.DeepNoiseSuppressionMeanOpinionScore
+    :exclude-members: update, compute
+
+Functional Interface
+____________________
+
+.. autofunction:: torchmetrics.functional.audio.dnsmos.deep_noise_suppression_mean_opinion_score

docs/source/links.rst

@@ -110,6 +110,7 @@
 .. _Theils Uncertainty coefficient: https://en.wikipedia.org/wiki/Uncertainty_coefficient
 .. _Perceptual Evaluation of Speech Quality: https://en.wikipedia.org/wiki/Perceptual_Evaluation_of_Speech_Quality
 .. _pesq package: https://github.com/ludlows/python-pesq
+.. _Deep Noise Suppression performance evaluation based on Mean Opinion Score: https://arxiv.org/abs/2010.15258
 .. _Cees Taal's website: http://www.ceestaal.nl/code/
 .. _pystoi package: https://github.com/mpariente/pystoi
 .. _stoi ref1: https://ieeexplore.ieee.org/abstract/document/5495701

requirements/audio.txt

@@ -6,3 +6,7 @@ pesq >=0.0.4, <0.0.5
 pystoi >=0.3.0, <0.5.0
 torchaudio >=0.10.0, <2.5.0
 gammatone >=1.0.0, <1.1.0
+librosa >=0.9.0, <0.11.0
+onnxruntime-gpu >=1.12.0, <1.19; sys_platform != 'darwin'
+onnxruntime >=1.12.0, <1.19; sys_platform == 'darwin' # installing onnxruntime-gpu failed on macos
+requests >=2.19.0, <2.32.0

src/torchmetrics/audio/__init__.py

@@ -24,6 +24,8 @@
 )
 from torchmetrics.utilities.imports import (
     _GAMMATONE_AVAILABLE,
+    _LIBROSA_AVAILABLE,
+    _ONNXRUNTIME_AVAILABLE,
     _PESQ_AVAILABLE,
     _PYSTOI_AVAILABLE,
     _TORCHAUDIO_AVAILABLE,
@@ -54,3 +56,8 @@
     from torchmetrics.audio.srmr import SpeechReverberationModulationEnergyRatio
 
     __all__ += ["SpeechReverberationModulationEnergyRatio"]
+
+if _LIBROSA_AVAILABLE and _ONNXRUNTIME_AVAILABLE:
+    from torchmetrics.audio.dnsmos import DeepNoiseSuppressionMeanOpinionScore
+
+    __all__ += ["DeepNoiseSuppressionMeanOpinionScore"]

src/torchmetrics/audio/dnsmos.py

@@ -0,0 +1,173 @@
+# Copyright The Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Any, Optional, Sequence, Union
+
+import torch
+from torch import Tensor, tensor
+
+from torchmetrics.functional.audio.dnsmos import deep_noise_suppression_mean_opinion_score
+from torchmetrics.metric import Metric
+from torchmetrics.utilities.imports import (
+    _LIBROSA_AVAILABLE,
+    _MATPLOTLIB_AVAILABLE,
+    _ONNXRUNTIME_AVAILABLE,
+    _REQUESTS_AVAILABLE,
+)
+from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE
+
+__doctest_requires__ = {"DeepNoiseSuppressionMeanOpinionScore": ["requests", "librosa", "onnxruntime"]}
+
+if not _MATPLOTLIB_AVAILABLE:
+    __doctest_skip__ = ["DeepNoiseSuppressionMeanOpinionScore.plot"]
+
+
+class DeepNoiseSuppressionMeanOpinionScore(Metric):
+    """Calculate `Deep Noise Suppression performance evaluation based on Mean Opinion Score`_ (DNSMOS).
+
+    Human subjective evaluation is the ”gold standard” to evaluate speech quality optimized for human perception.
+    Perceptual objective metrics serve as a proxy for subjective scores. The conventional and widely used metrics
+    require a reference clean speech signal, which is unavailable in real recordings. The no-reference approaches
+    correlate poorly with human ratings and are not widely adopted in the research community. One of the biggest
+    use cases of these perceptual objective metrics is to evaluate noise suppression algorithms. DNSMOS generalizes
+    well in challenging test conditions with a high correlation to human ratings in stack ranking noise suppression
+    methods. More details can be found in `DNSMOS paper <https://arxiv.org/abs/2010.15258>`_ and
+    `DNSMOS P.835 paper <https://arxiv.org/abs/2110.01763>`_.
+
+
+    As input to ``forward`` and ``update`` the metric accepts the following input
+
+    - ``preds`` (:class:`~torch.Tensor`): float tensor with shape ``(...,time)``
+
+    As output of ``forward`` and ``compute`` the metric returns the following output
+
+    - ``dnsmos`` (:class:`~torch.Tensor`): float tensor of DNSMOS values reduced across the batch
+        with shape ``(...,4)`` indicating [p808_mos, mos_sig, mos_bak, mos_ovr] in the last dim.
+
+    .. note:: using this metric requires you to have ``librosa``, ``onnxruntime`` and ``requests`` installed.
+        Install as ``pip install torchmetrics['audio']`` or alternatively `pip install librosa onnxruntime-gpu requests`
+        (if you do not have GPU enabled machine install `onnxruntime` instead of `onnxruntime-gpu`)
+
+    .. note:: the ``forward`` and ``compute`` methods in this class return a reduced DNSMOS value
+        for a batch. To obtain the DNSMOS value for each sample, you may use the functional counterpart in
+        :func:`~torchmetrics.functional.audio.dnsmos.deep_noise_suppression_mean_opinion_score`.
+
+    Args:
+        fs: sampling frequency
+        personalized: whether interfering speaker is penalized
+        device: the device used for calculating DNSMOS, can be cpu or cuda:n, where n is the index of gpu.
+            If None is given, then the device of input is used.
+        num_threads: number of threads to use for onnxruntime CPU inference.
+
+    Raises:
+        ModuleNotFoundError:
+            If ``librosa``, ``onnxruntime`` or ``requests`` packages are not installed
+
+    Example:
+        >>> from torch import randn
+        >>> from torchmetrics.audio import DeepNoiseSuppressionMeanOpinionScore
+        >>> g = torch.manual_seed(1)
+        >>> preds = randn(8000)
+        >>> dnsmos = DeepNoiseSuppressionMeanOpinionScore(8000, False)
+        >>> dnsmos(preds)
+        tensor([2.2285, 2.1132, 1.3972, 1.3652], dtype=torch.float64)
+
+    """
+
+    sum_dnsmos: Tensor
+    total: Tensor
+    full_state_update: bool = False
+    is_differentiable: bool = False
+    higher_is_better: bool = True
+    plot_lower_bound: float = 0
+    plot_upper_bound: float = 5
+
+    def __init__(
+        self,
+        fs: int,
+        personalized: bool,
+        device: Optional[str] = None,
+        num_threads: Optional[int] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        if not _LIBROSA_AVAILABLE or not _ONNXRUNTIME_AVAILABLE or not _REQUESTS_AVAILABLE:
+            raise ModuleNotFoundError(
+                "DNSMOS metric requires that librosa, onnxruntime and requests are installed."
+                " Install as `pip install librosa onnxruntime-gpu requests`."
+            )
+
+        self.fs = fs
+        self.personalized = personalized
+        self.cal_device = device
+        self.num_threads = num_threads
+
+        self.add_state("sum_dnsmos", default=tensor([0, 0, 0, 0], dtype=torch.float64), dist_reduce_fx="sum")
+        self.add_state("total", default=tensor(0), dist_reduce_fx="sum")
+
+    def update(self, preds: Tensor) -> None:
+        """Update state with predictions."""
+        metric_batch = deep_noise_suppression_mean_opinion_score(
+            preds,
+            self.fs,
+            self.personalized,
+            self.cal_device,
+            self.num_threads,
+        ).to(self.sum_dnsmos.device)
+
+        self.sum_dnsmos += metric_batch.reshape(-1, 4).sum(dim=0)
+        self.total += metric_batch.reshape(-1, 4).shape[0]
+
+    def compute(self) -> Tensor:
+        """Compute metric."""
+        return self.sum_dnsmos / self.total
+
+    def plot(self, val: Union[Tensor, Sequence[Tensor], None] = None, ax: Optional[_AX_TYPE] = None) -> _PLOT_OUT_TYPE:
+        """Plot a single or multiple values from the metric.
+
+        Args:
+            val: Either a single result from calling ``metric.forward`` or ``metric.compute`` or a list of these
+                results. If no value is provided, will automatically call ``metric.compute`` and plot that result.
+            ax: A matplotlib axis object. If provided will add plot to that axis
+
+        Returns:
+            Figure and Axes object
+
+        Raises:
+            ModuleNotFoundError:
+                If ``matplotlib`` is not installed
+
+        .. plot::
+            :scale: 75
+
+            >>> # Example plotting a single value
+            >>> import torch
+            >>> from torchmetrics.audio import DeepNoiseSuppressionMeanOpinionScore
+            >>> metric = DeepNoiseSuppressionMeanOpinionScore(8000, False)
+            >>> metric.update(torch.rand(8000))
+            >>> fig_, ax_ = metric.plot()
+
+        .. plot::
+            :scale: 75
+
+            >>> # Example plotting multiple values
+            >>> import torch
+            >>> from torchmetrics.audio import DeepNoiseSuppressionMeanOpinionScore
+            >>> metric = DeepNoiseSuppressionMeanOpinionScore(8000, False)
+            >>> values = [ ]
+            >>> for _ in range(10):
+            ...     values.append(metric(torch.rand(8000)))
+            >>> fig_, ax_ = metric.plot(values)
+
+        """
+        return self._plot(val, ax)

src/torchmetrics/functional/audio/__init__.py

@@ -24,6 +24,8 @@
 )
 from torchmetrics.utilities.imports import (
     _GAMMATONE_AVAILABLE,
+    _LIBROSA_AVAILABLE,
+    _ONNXRUNTIME_AVAILABLE,
     _PESQ_AVAILABLE,
     _PYSTOI_AVAILABLE,
     _TORCHAUDIO_AVAILABLE,
@@ -55,3 +57,8 @@
     from torchmetrics.functional.audio.srmr import speech_reverberation_modulation_energy_ratio
 
     __all__ += ["speech_reverberation_modulation_energy_ratio"]
+
+if _LIBROSA_AVAILABLE and _ONNXRUNTIME_AVAILABLE:
+    from torchmetrics.functional.audio.dnsmos import deep_noise_suppression_mean_opinion_score
+
+    __all__ += ["deep_noise_suppression_mean_opinion_score"]