Add plot() functionality to audio metrics

CHANGELOG.md

@@ -20,6 +20,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 )
 
 
+- Added support for plotting of audio metrics through `.plot()` method ([#1434](https://github.com/Lightning-AI/metrics/pull/1434))
+
+
 - Added `classes` to output from `MAP` metric ([#1419](https://github.com/Lightning-AI/metrics/pull/1419))
 
 

examples/plotting.py

@@ -17,6 +17,147 @@
 import torch
 
 
+def pesq_example():
+    """Plot PESQ audio example."""
+    from torchmetrics.audio.pesq import PerceptualEvaluationSpeechQuality
+
+    p = lambda: torch.randn(8000)
+    t = lambda: torch.randn(8000)
+
+    # plot single value
+    metric = PerceptualEvaluationSpeechQuality(8000, "nb")
+    metric.update(p(), t())
+    fig, ax = metric.plot()
+
+    # plot multiple values
+    metric = PerceptualEvaluationSpeechQuality(16000, "wb")
+    vals = [metric(p(), t()) for _ in range(10)]
+    fig, ax = metric.plot(vals)
+
+    return fig, ax
+
+
+def pit_example():
+    """Plot PIT audio example."""
+    from torchmetrics.audio.pit import PermutationInvariantTraining
+    from torchmetrics.functional import scale_invariant_signal_noise_ratio
+
+    p = lambda: torch.randn(3, 2, 5)
+    t = lambda: torch.randn(3, 2, 5)
+
+    # plot single value
+    metric = PermutationInvariantTraining(scale_invariant_signal_noise_ratio, "max")
+    metric.update(p(), t())
+    fig, ax = metric.plot()
+
+    # plot multiple values
+    metric = PermutationInvariantTraining(scale_invariant_signal_noise_ratio, "max")
+    vals = [metric(p(), t()) for _ in range(10)]
+    fig, ax = metric.plot(vals)
+
+    return fig, ax
+
+
+def sdr_example():
+    """Plot SDR audio example."""
+    from torchmetrics.audio.sdr import SignalDistortionRatio
+
+    p = lambda: torch.randn(8000)
+    t = lambda: torch.randn(8000)
+
+    # plot single value
+    metric = SignalDistortionRatio()
+    metric.update(p(), t())
+    fig, ax = metric.plot()
+
+    # plot multiple values
+    metric = SignalDistortionRatio()
+    vals = [metric(p(), t()) for _ in range(10)]
+    fig, ax = metric.plot(vals)
+
+    return fig, ax
+
+
+def si_sdr_example():
+    """Plot SI-SDR audio example."""
+    from torchmetrics.audio.sdr import ScaleInvariantSignalDistortionRatio
+
+    p = lambda: torch.randn(5)
+    t = lambda: torch.randn(5)
+
+    # plot single value
+    metric = ScaleInvariantSignalDistortionRatio()
+    metric.update(p(), t())
+    fig, ax = metric.plot()
+
+    # plot multiple values
+    metric = ScaleInvariantSignalDistortionRatio()
+    vals = [metric(p(), t()) for _ in range(10)]
+    fig, ax = metric.plot(vals)
+
+    return fig, ax
+
+
+def snr_example():
+    """Plot SNR audio example."""
+    from torchmetrics.audio.snr import SignalNoiseRatio
+
+    p = lambda: torch.randn(4)
+    t = lambda: torch.randn(4)
+
+    # plot single value
+    metric = SignalNoiseRatio()
+    metric.update(p(), t())
+    fig, ax = metric.plot()
+
+    # plot multiple values
+    metric = SignalNoiseRatio()
+    vals = [metric(p(), t()) for _ in range(10)]
+    fig, ax = metric.plot(vals)
+
+    return fig, ax
+
+
+def si_snr_example():
+    """Plot SI-SNR example."""
+    from torchmetrics.audio.snr import ScaleInvariantSignalNoiseRatio
+
+    p = lambda: torch.randn(4)
+    t = lambda: torch.randn(4)
+
+    # plot single value
+    metric = ScaleInvariantSignalNoiseRatio()
+    metric.update(p(), t())
+    fig, ax = metric.plot()
+
+    # plot multiple values
+    metric = ScaleInvariantSignalNoiseRatio()
+    vals = [metric(p(), t()) for _ in range(10)]
+    fig, ax = metric.plot(vals)
+
+    return fig, ax
+
+
+def stoi_example():
+    """Plot STOI example."""
+    from torchmetrics.audio.stoi import ShortTimeObjectiveIntelligibility
+
+    p = lambda: torch.randn(8000)
+    t = lambda: torch.randn(8000)
+
+    # plot single value
+    metric = ShortTimeObjectiveIntelligibility(8000, False)
+    metric.update(p(), t())
+    fig, ax = metric.plot()
+
+    # plot multiple values
+    metric = ShortTimeObjectiveIntelligibility(8000, False)
+    vals = [metric(p(), t()) for _ in range(10)]
+    fig, ax = metric.plot(vals)
+
+    return fig, ax
+
+
 def accuracy_example():
     """Plot Accuracy example."""
     from torchmetrics.classification import MulticlassAccuracy
@@ -85,6 +226,13 @@ def confusion_matrix_example():
 
     metrics_func = {
         "accuracy": accuracy_example,
+        "pesq": pesq_example,
+        "pit": pit_example,
+        "sdr": sdr_example,
+        "si-sdr": si_sdr_example,
+        "snr": snr_example,
+        "si-snr": si_snr_example,
+        "stoi": stoi_example,
         "mean_squared_error": mean_squared_error_example,
         "confusion_matrix": confusion_matrix_example,
     }

requirements/docs.txt

@@ -14,3 +14,4 @@ sphinx-copybutton>=0.3
 # integrations
 -r integrate.txt
 -r visual.txt
+-r audio.txt

src/torchmetrics/audio/pesq.py

@@ -11,16 +11,20 @@
 # 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
+from typing import Any, Optional, Sequence, Union
 
 from torch import Tensor, tensor
 
 from torchmetrics.functional.audio.pesq import perceptual_evaluation_speech_quality
 from torchmetrics.metric import Metric
-from torchmetrics.utilities.imports import _PESQ_AVAILABLE
+from torchmetrics.utilities.imports import _MATPLOTLIB_AVAILABLE, _PESQ_AVAILABLE
+from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE, plot_single_or_multi_val
 
 __doctest_requires__ = {"PerceptualEvaluationSpeechQuality": ["pesq"]}
 
+if not _MATPLOTLIB_AVAILABLE:
+    __doctest_skip__ = ["PerceptualEvaluationSpeechQuality.plot"]
+
 
 class PerceptualEvaluationSpeechQuality(Metric):
     """Calculate `Perceptual Evaluation of Speech Quality`_ (PESQ). It's a recognized industry standard for audio
@@ -80,6 +84,7 @@ class PerceptualEvaluationSpeechQuality(Metric):
     full_state_update: bool = False
     is_differentiable: bool = False
     higher_is_better: bool = True
+    plot_options: dict = {"lower_bound": 1.0, "upper_bound": 4.5}
 
     def __init__(
         self,
@@ -119,3 +124,49 @@ def update(self, preds: Tensor, target: Tensor) -> None:
     def compute(self) -> Tensor:
         """Compute metric."""
         return self.sum_pesq / self.total
+
+    def plot(
+        self, val: Optional[Union[Tensor, Sequence[Tensor]]] = 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: An matplotlib axis object. If provided will add plot to that axis
+
+        Returns:
+            Figure and Axes object
+
+        Raises:
+            ModuleNotFoundError:
+                If `matplotlib` is not installed
+
+        Examples:
+        .. plot::
+            :scale: 75
+
+            >>> # Example plotting a single value
+            >>> import torch
+            >>> from torchmetrics.audio.pesq import PerceptualEvaluationSpeechQuality
+            >>> metric = PerceptualEvaluationSpeechQuality(8000, 'nb')
+            >>> metric.update(torch.rand(8000), torch.rand(8000))
+            >>> fig_, ax_ = metric.plot()
+
+        .. plot::
+            :scale: 75
+
+            >>> # Example plotting multiple values
+            >>> import torch
+            >>> from torchmetrics.audio.pesq import PerceptualEvaluationSpeechQuality
+            >>> metric = PerceptualEvaluationSpeechQuality(8000, 'nb')
+            >>> values = [ ]
+            >>> for _ in range(10):
+            ...     values.append(metric(torch.rand(8000), torch.rand(8000)))
+            >>> fig_, ax_ = metric.plot(values)
+        """
+        val = val or self.compute()
+        fig, ax = plot_single_or_multi_val(
+            val, ax=ax, higher_is_better=self.higher_is_better, **self.plot_options, name=self.__class__.__name__
+        )
+        return fig, ax

src/torchmetrics/audio/pit.py

@@ -11,13 +11,20 @@
 # 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, Callable, Dict
+from typing import Any, Callable, Dict, Optional, Sequence, Union
 
 from torch import Tensor, tensor
 from typing_extensions import Literal
 
 from torchmetrics.functional.audio.pit import permutation_invariant_training
 from torchmetrics.metric import Metric
+from torchmetrics.utilities.imports import _MATPLOTLIB_AVAILABLE
+from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE, plot_single_or_multi_val
+
+__doctest_requires__ = {"PermutationInvariantTraining": ["pit"]}
+
+if not _MATPLOTLIB_AVAILABLE:
+    __doctest_skip__ = ["PermutationInvariantTraining.plot"]
 
 
 class PermutationInvariantTraining(Metric):
@@ -60,6 +67,7 @@ class PermutationInvariantTraining(Metric):
     is_differentiable: bool = True
     sum_pit_metric: Tensor
     total: Tensor
+    plot_options: dict = {"lower_bound": -10.0, "upper_bound": 1.0}
 
     def __init__(
         self,
@@ -90,3 +98,55 @@ def update(self, preds: Tensor, target: Tensor) -> None:
     def compute(self) -> Tensor:
         """Compute metric."""
         return self.sum_pit_metric / self.total
+
+    def plot(
+        self, val: Optional[Union[Tensor, Sequence[Tensor]]] = 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: An matplotlib axis object. If provided will add plot to that axis
+
+        Returns:
+            Figure and Axes object
+
+        Raises:
+            ModuleNotFoundError:
+                If `matplotlib` is not installed
+
+        Examples:
+        .. plot::
+            :scale: 75
+
+            >>> # Example plotting a single value
+            >>> import torch
+            >>> from torchmetrics.audio.pit import PermutationInvariantTraining
+            >>> from torchmetrics.functional import scale_invariant_signal_noise_ratio
+            >>> preds = torch.randn(3, 2, 5) # [batch, spk, time]
+            >>> target = torch.randn(3, 2, 5) # [batch, spk, time]
+            >>> metric = PermutationInvariantTraining(scale_invariant_signal_noise_ratio, 'max')
+            >>> metric.update(preds, target)
+            >>> fig_, ax_ = metric.plot()
+
+        .. plot::
+            :scale: 75
+
+            >>> # Example plotting multiple values
+            >>> import torch
+            >>> from torchmetrics.audio.pit import PermutationInvariantTraining
+            >>> from torchmetrics.functional import scale_invariant_signal_noise_ratio
+            >>> preds = torch.randn(3, 2, 5) # [batch, spk, time]
+            >>> target = torch.randn(3, 2, 5) # [batch, spk, time]
+            >>> metric = PermutationInvariantTraining(scale_invariant_signal_noise_ratio, 'max')
+            >>> values = [ ]
+            >>> for _ in range(10):
+            ...     values.append(metric(preds, target))
+            >>> fig_, ax_ = metric.plot(values)
+        """
+        val = val or self.compute()
+        fig, ax = plot_single_or_multi_val(
+            val, ax=ax, higher_is_better=self.higher_is_better, **self.plot_options, name=self.__class__.__name__
+        )
+        return fig, ax