o .wi@sddlmZddlmZmZmZddlmZmZddl m Z ddl m Z ddl mZmZmZddlmZmZdd d giZesAd gZGd dde Zd S))Sequence)AnyOptionalUnion)Tensortensor)'non_intrusive_speech_quality_assessment)Metric)_LIBROSA_AVAILABLE_MATPLOTLIB_AVAILABLE_REQUESTS_AVAILABLE)_AX_TYPE_PLOT_OUT_TYPE#NonIntrusiveSpeechQualityAssessmentlibrosarequests(NonIntrusiveSpeechQualityAssessment.plotcseZdZUdZeed<eed<dZeed<dZeed<dZ eed<d Z e ed <d Z e ed <d e deddffdd ZdeddfddZdefddZddeeeedfdeedefddZZS)ruJ`Non-Intrusive Speech Quality Assessment`_ (NISQA v2.0) [1], [2]. 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 - ``nisqa`` (:class:`~torch.Tensor`): float tensor reduced across the batch with shape ``(5,)`` corresponding to overall MOS, noisiness, discontinuity, coloration and loudness in that order .. hint:: Using this metric requires you to have ``librosa`` and ``requests`` installed. Install as ``pip install librosa requests``. .. caution:: The ``forward`` and ``compute`` methods in this class return values reduced across the batch. To obtain values for each sample, you may use the functional counterpart :func:`~torchmetrics.functional.audio.nisqa.non_intrusive_speech_quality_assessment`. Args: fs: sampling frequency of input Raises: ModuleNotFoundError: If ``librosa`` or ``requests`` are not installed Example: >>> import torch >>> from torchmetrics.audio import NonIntrusiveSpeechQualityAssessment >>> _ = torch.manual_seed(42) >>> preds = torch.randn(16000) >>> nisqa = NonIntrusiveSpeechQualityAssessment(16000) >>> nisqa(preds) tensor([1.0433, 1.9545, 2.6087, 1.3460, 1.7117]) References: - [1] G. Mittag and S. Möller, "Non-intrusive speech quality assessment for super-wideband speech communication networks", in Proc. ICASSP, 2019. - [2] G. Mittag, B. Naderi, A. Chehadi and S. Möller, "NISQA: A deep CNN-self-attention model for multidimensional speech quality prediction with crowdsourced datasets", in Proc. INTERSPEECH, 2021. sum_nisqatotalFfull_state_updateis_differentiableThigher_is_betterplot_lower_boundg@plot_upper_boundfskwargsreturnNc sxtjd i|tr tstdt|tr|dkr!td|||_|j dt gddd|j dt ddddS) NziNISQA metric requires that librosa and requests are installed. Install as `pip install librosa requests`.rz9Argument `fs` expected to be a positive integer, but got r)rrrrrsum)defaultdist_reduce_fxr) super__init__r r ModuleNotFoundError isinstanceint ValueErrorr add_stater)selfrr __class__r!U/home/ubuntu/sommelier/.venv/lib/python3.10/site-packages/torchmetrics/audio/nisqa.pyr#Xsz,NonIntrusiveSpeechQualityAssessment.__init__predscCsPt||j|jj}|dd}|j|jdd7_|j|jd7_dS)zUpdate state with predictions.r)dimN) rrtordevicereshaperrshape)r)r- nisqa_batchr!r!r,updatefs  z*NonIntrusiveSpeechQualityAssessment.updatecCs |j|jS)zCompute metric.)rr)r)r!r!r,computeqs z+NonIntrusiveSpeechQualityAssessment.computevalaxcCs |||S)aFPlot 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 NonIntrusiveSpeechQualityAssessment >>> metric = NonIntrusiveSpeechQualityAssessment(16000) >>> metric.update(torch.randn(16000)) >>> fig_, ax_ = metric.plot() .. plot:: :scale: 75 >>> # Example plotting multiple values >>> import torch >>> from torchmetrics.audio import NonIntrusiveSpeechQualityAssessment >>> metric = NonIntrusiveSpeechQualityAssessment(16000) >>> values = [] >>> for _ in range(10): ... values.append(metric(torch.randn(16000))) >>> fig_, ax_ = metric.plot(values) )_plot)r)r8r9r!r!r,plotus &r)NN)__name__ __module__ __qualname____doc__r__annotations__rboolrrrfloatrr&rr#r6r7rrrr rr; __classcell__r!r!r*r,r#s ,      2N)collections.abcrtypingrrrtorchrr#torchmetrics.functional.audio.nisqartorchmetrics.metricr torchmetrics.utilities.importsr r r torchmetrics.utilities.plotr r__doctest_requires____doctest_skip__rr!r!r!r,s