o iV@sddlZddlmZmZddlZddlmZddlmZgZ    dd ejd ejd ejd e d e de de dejfddZ GdddejjZGdddejjZGdddejjZGdddejjZdS)N)OptionalUnion)Tensor) functional ref_channelTHz>:0yE>psd_spsd_nreference_vectorsolutiondiagonal_loadingdiag_epsepsreturnc Cs`|dkrt||||||}|S|dkrt|}n tj|||||d}t||||||}|S)aCompute the MVDR beamforming weights with ``solution`` argument. Args: psd_s (torch.Tensor): The complex-valued power spectral density (PSD) matrix of target speech. Tensor with dimensions `(..., freq, channel, channel)`. psd_n (torch.Tensor): The complex-valued power spectral density (PSD) matrix of noise. Tensor with dimensions `(..., freq, channel, channel)`. reference_vector (torch.Tensor): one-hot reference channel matrix. solution (str, optional): Solution to compute the MVDR beamforming weights. Options: [``ref_channel``, ``stv_evd``, ``stv_power``]. (Default: ``ref_channel``) diagonal_loading (bool, optional): If ``True``, enables applying diagonal loading to ``psd_n``. (Default: ``True``) diag_eps (float, optional): The coefficient multiplied to the identity matrix for diagonal loading. It is only effective when ``diagonal_loading`` is set to ``True``. (Default: ``1e-7``) eps (float, optional): Value to add to the denominator in the beamforming weight formula. (Default: ``1e-8``) Returns: torch.Tensor: the mvdr beamforming weight matrix rstv_evd)r r)Fmvdr_weights_soudenrtf_evd rtf_powermvdr_weights_rtf) r r r r r rrbeamform_vectorstvrX/home/ubuntu/.local/lib/python3.10/site-packages/torchaudio/transforms/_multi_channel.py_get_mvdr_vectors rcsLeZdZdZddededeffdd Zdd ejd e ejfd dZ Z S)PSDaCompute cross-channel power spectral density (PSD) matrix. .. devices:: CPU CUDA .. properties:: Autograd TorchScript Args: multi_mask (bool, optional): If ``True``, only accepts multi-channel Time-Frequency masks. (Default: ``False``) normalize (bool, optional): If ``True``, normalize the mask along the time dimension. (Default: ``True``) eps (float, optional): Value to add to the denominator in mask normalization. (Default: ``1e-15``) FTV瞯< multi_mask normalizercs t||_||_||_dSN)super__init__rrr)selfrrr __class__rrr"Ds  z PSD.__init__NspecgrammaskcCs2|dur |jr |jdd}t|||j|j}|S)a Args: specgram (torch.Tensor): Multi-channel complex-valued spectrum. Tensor with dimensions `(..., channel, freq, time)`. mask (torch.Tensor or None, optional): Time-Frequency mask for normalization. Tensor with dimensions `(..., freq, time)` if multi_mask is ``False`` or with dimensions `(..., channel, freq, time)` if multi_mask is ``True``. (Default: ``None``) Returns: torch.Tensor: The complex-valued PSD matrix of the input spectrum. Tensor with dimensions `(..., freq, channel, channel)` Ndim)rmeanrpsdrr)r#r&r'r,rrrforwardJs  z PSD.forward)FTrr ) __name__ __module__ __qualname____doc__boolfloatr"torchrrr- __classcell__rrr$rr7s $rcseZdZdZ      d!dededed ed ed ef fd d Z    d"de j de j de j de j de j deded edede j fddZ de j de j de j fddZ de j de j de j fddZ d#de j de j dee j de j fdd ZZS)$MVDRa Minimum Variance Distortionless Response (MVDR) module that performs MVDR beamforming with Time-Frequency masks. .. devices:: CPU CUDA .. properties:: Autograd TorchScript Based on https://github.com/espnet/espnet/blob/master/espnet2/enh/layers/beamformer.py We provide three solutions of MVDR beamforming. One is based on *reference channel selection* :cite:`souden2009optimal` (``solution=ref_channel``). .. math:: \textbf{w}_{\text{MVDR}}(f) = \frac{{{\bf{\Phi}_{\textbf{NN}}^{-1}}(f){\bf{\Phi}_{\textbf{SS}}}}(f)} {\text{Trace}({{{\bf{\Phi}_{\textbf{NN}}^{-1}}(f) \bf{\Phi}_{\textbf{SS}}}(f))}}\bm{u} where :math:`\bf{\Phi}_{\textbf{SS}}` and :math:`\bf{\Phi}_{\textbf{NN}}` are the covariance matrices of speech and noise, respectively. :math:`\bf{u}` is an one-hot vector to determine the reference channel. The other two solutions are based on the steering vector (``solution=stv_evd`` or ``solution=stv_power``). .. math:: \textbf{w}_{\text{MVDR}}(f) = \frac{{{\bf{\Phi}_{\textbf{NN}}^{-1}}(f){\bm{v}}(f)}} {{\bm{v}^{\mathsf{H}}}(f){\bf{\Phi}_{\textbf{NN}}^{-1}}(f){\bm{v}}(f)} where :math:`\bm{v}` is the acoustic transfer function or the steering vector. :math:`.^{\mathsf{H}}` denotes the Hermitian Conjugate operation. We apply either *eigenvalue decomposition* :cite:`higuchi2016robust` or the *power method* :cite:`mises1929praktische` to get the steering vector from the PSD matrix of speech. After estimating the beamforming weight, the enhanced Short-time Fourier Transform (STFT) is obtained by .. math:: \hat{\bf{S}} = {\bf{w}^\mathsf{H}}{\bf{Y}}, {\bf{w}} \in \mathbb{C}^{M \times F} where :math:`\bf{Y}` and :math:`\hat{\bf{S}}` are the STFT of the multi-channel noisy speech and the single-channel enhanced speech, respectively. For online streaming audio, we provide a *recursive method* :cite:`higuchi2017online` to update the PSD matrices of speech and noise, respectively. Args: ref_channel (int, optional): Reference channel for beamforming. (Default: ``0``) solution (str, optional): Solution to compute the MVDR beamforming weights. Options: [``ref_channel``, ``stv_evd``, ``stv_power``]. (Default: ``ref_channel``) multi_mask (bool, optional): If ``True``, only accepts multi-channel Time-Frequency masks. (Default: ``False``) diagonal_loading (bool, optional): If ``True``, enables applying diagonal loading to the covariance matrix of the noise. (Default: ``True``) diag_eps (float, optional): The coefficient multiplied to the identity matrix for diagonal loading. It is only effective when ``diagonal_loading`` is set to ``True``. (Default: ``1e-7``) online (bool, optional): If ``True``, updates the MVDR beamforming weights based on the previous covarience matrices. (Default: ``False``) Note: To improve the numerical stability, the input spectrogram will be converted to double precision (``torch.complex128`` or ``torch.cdouble``) dtype for internal computation. The output spectrogram is converted to the dtype of the input spectrogram to be compatible with other modules. Note: If you use ``stv_evd`` solution, the gradient of the same input may not be identical if the eigenvalues of the PSD matrix are not distinct (i.e. some eigenvalues are close or identical). rrFTrr r diag_loadingronlinec st|dvrtd|||_||_||_||_||_||_ t ||_ t d}t d}t d} t d} |d||d||d| |d| dS)N)rr stv_powerzK`solution` must be one of ["ref_channel", "stv_evd", "stv_power"]. Given {}r r mask_sum_s mask_sum_n)r!r" ValueErrorformatrr rr7rr8rr,r4zerosregister_buffer) r#rr rr7rr8r r r;r<r$rrr"s(         z MVDR.__init__rr r mask_smask_nr r rrc Cs|jr|jdd}|jdd}|jjdkr3||_||_|jdd|_|jdd|_t||||||| S| ||}| ||}||_||_|j|jdd|_|j|jdd|_t||||||| S)aRecursively update the MVDR beamforming vector. Args: psd_s (torch.Tensor): The complex-valued power spectral density (PSD) matrix of target speech. Tensor with dimensions `(..., freq, channel, channel)`. psd_n (torch.Tensor): The complex-valued power spectral density (PSD) matrix of noise. Tensor with dimensions `(..., freq, channel, channel)`. mask_s (torch.Tensor): Time-Frequency mask of the target speech. Tensor with dimensions `(..., freq, time)` if multi_mask is ``False`` or with dimensions `(..., channel, freq, time)` if multi_mask is ``True``. mask_n (torch.Tensor or None, optional): Time-Frequency mask of the noise. Tensor with dimensions `(..., freq, time)` if multi_mask is ``False`` or with dimensions `(..., channel, freq, time)` if multi_mask is ``True``. reference_vector (torch.Tensor): One-hot reference channel matrix. solution (str, optional): Solution to compute the MVDR beamforming weights. Options: [``ref_channel``, ``stv_evd``, ``stv_power``]. (Default: ``ref_channel``) diagonal_loading (bool, optional): If ``True``, enables applying diagonal loading to ``psd_n``. (Default: ``True``) diag_eps (float, optional): The coefficient multiplied to the identity matrix for diagonal loading. It is only effective when ``diagonal_loading`` is set to ``True``. (Default: ``1e-7``) eps (float, optional): Value to add to the denominator in the beamforming weight formula. (Default: ``1e-8``) Returns: torch.Tensor: The MVDR beamforming weight matrix. r(r)r:) rr+r ndimr sumr;r<r_get_updated_psd_speech_get_updated_psd_noise) r#r r rArBr r r rrrrr_get_updated_mvdr_vectors &     zMVDR._get_updated_mvdr_vectorcCL|j|j|jdd}d|j|jdd}|j|d||d}|S)aUpdate psd of speech recursively. Args: psd_s (torch.Tensor): The complex-valued power spectral density (PSD) matrix of target speech. Tensor with dimensions `(..., freq, channel, channel)`. mask_s (torch.Tensor): Time-Frequency mask of the target speech. Tensor with dimensions `(..., freq, time)`. Returns: torch.Tensor: The updated PSD matrix of target speech. rCr)r:.NN)r;rEr )r#r rA numerator denominatorrrrrF zMVDR._get_updated_psd_speechcCrI)aUpdate psd of noise recursively. Args: psd_n (torch.Tensor): The complex-valued power spectral density (PSD) matrix of noise. Tensor with dimensions `(..., freq, channel, channel)`. mask_n (torch.Tensor or None, optional): Time-Frequency mask of the noise. Tensor with dimensions `(..., freq, time)`. Returns: torch.Tensor: The updated PSD matrix of noise. rCr)r:rJ)r<rEr )r#r rBrKrLrrrrGrMzMVDR._get_updated_psd_noiseNr&c Cs|j}|jdkrtd|j|std|j|jtjkr&|}|dur3t dd|}| ||}| ||}tj | dd|j tjd}|d |jfd|jrl|||||||j|j|j}n t||||j|j|j}t||} | |S) a`Perform MVDR beamforming. Args: specgram (torch.Tensor): Multi-channel complex-valued spectrum. Tensor with dimensions `(..., channel, freq, time)` mask_s (torch.Tensor): Time-Frequency mask of target speech. Tensor with dimensions `(..., freq, time)` if multi_mask is ``False`` or with dimensions `(..., channel, freq, time)` if multi_mask is ``True``. mask_n (torch.Tensor or None, optional): Time-Frequency mask of noise. Tensor with dimensions `(..., freq, time)` if multi_mask is ``False`` or with dimensions `(..., channel, freq, time)` if multi_mask is ``True``. (Default: None) Returns: torch.Tensor: Single-channel complex-valued enhanced spectrum with dimensions `(..., freq, time)`. z?Expected at least 3D tensor (..., channel, freq, time). Found: ziThe type of ``specgram`` tensor must be ``torch.cfloat`` or ``torch.cdouble``. Found: Nz=``mask_n`` is not provided, use ``1 - mask_s`` as ``mask_n``.r:)devicedtype.)rQrDr=shape is_complexr4cfloatcdoublewarningswarnr,r?sizerPrfill_r8rHr r7rrrapply_beamformingto) r#r&rArBrQr r uw_mvdrspecgram_enhancedrrrr-#s2        z MVDR.forward)rrFTrFrTrrr )r.r/r0r1intstrr2r3r"r4rrHrFrGrr-r5rrr$rr6astE*   9r6c@sJeZdZdZ   ddedededeeefd ed ed ed efd dZ dS)RTFMVDRaMinimum Variance Distortionless Response (*MVDR* :cite:`capon1969high`) module based on the relative transfer function (RTF) and power spectral density (PSD) matrix of noise. .. devices:: CPU CUDA .. properties:: Autograd TorchScript Given the multi-channel complex-valued spectrum :math:`\textbf{Y}`, the relative transfer function (RTF) matrix or the steering vector of target speech :math:`\bm{v}`, the PSD matrix of noise :math:`\bf{\Phi}_{\textbf{NN}}`, and a one-hot vector that represents the reference channel :math:`\bf{u}`, the module computes the single-channel complex-valued spectrum of the enhanced speech :math:`\hat{\textbf{S}}`. The formula is defined as: .. math:: \hat{\textbf{S}}(f) = \textbf{w}_{\text{bf}}(f)^{\mathsf{H}} \textbf{Y}(f) where :math:`\textbf{w}_{\text{bf}}(f)` is the MVDR beamforming weight for the :math:`f`-th frequency bin, :math:`(.)^{\mathsf{H}}` denotes the Hermitian Conjugate operation. The beamforming weight is computed by: .. math:: \textbf{w}_{\text{MVDR}}(f) = \frac{{{\bf{\Phi}_{\textbf{NN}}^{-1}}(f){\bm{v}}(f)}} {{\bm{v}^{\mathsf{H}}}(f){\bf{\Phi}_{\textbf{NN}}^{-1}}(f){\bm{v}}(f)} Trrr&rtfr reference_channelr rrrc C$t||||||}t||} | S)a Args: specgram (torch.Tensor): Multi-channel complex-valued spectrum. Tensor with dimensions `(..., channel, freq, time)` rtf (torch.Tensor): The complex-valued RTF vector of target speech. Tensor with dimensions `(..., freq, channel)`. psd_n (torch.Tensor): The complex-valued power spectral density (PSD) matrix of noise. Tensor with dimensions `(..., freq, channel, channel)`. reference_channel (int or torch.Tensor): Specifies the reference channel. If the dtype is ``int``, it represents the reference channel index. If the dtype is ``torch.Tensor``, its shape is `(..., channel)`, where the ``channel`` dimension is one-hot. diagonal_loading (bool, optional): If ``True``, enables applying diagonal loading to ``psd_n``. (Default: ``True``) diag_eps (float, optional): The coefficient multiplied to the identity matrix for diagonal loading. It is only effective when ``diagonal_loading`` is set to ``True``. (Default: ``1e-7``) eps (float, optional): Value to add to the denominator in the beamforming weight formula. (Default: ``1e-8``) Returns: torch.Tensor: Single-channel complex-valued enhanced spectrum with dimensions `(..., freq, time)`. )rrrZ) r#r&rcr rdr rrr]spectrum_enhancedrrrr-r zRTFMVDR.forwardNTrr) r.r/r0r1rrr`r2r3r-rrrrrbWs,   rbc@sLeZdZdZ   ddedededeeefd ed ed ed e jfd dZ dS) SoudenMVDRaMinimum Variance Distortionless Response (*MVDR* :cite:`capon1969high`) module based on the method proposed by *Souden et, al.* :cite:`souden2009optimal`. .. devices:: CPU CUDA .. properties:: Autograd TorchScript Given the multi-channel complex-valued spectrum :math:`\textbf{Y}`, the power spectral density (PSD) matrix of target speech :math:`\bf{\Phi}_{\textbf{SS}}`, the PSD matrix of noise :math:`\bf{\Phi}_{\textbf{NN}}`, and a one-hot vector that represents the reference channel :math:`\bf{u}`, the module computes the single-channel complex-valued spectrum of the enhanced speech :math:`\hat{\textbf{S}}`. The formula is defined as: .. math:: \hat{\textbf{S}}(f) = \textbf{w}_{\text{bf}}(f)^{\mathsf{H}} \textbf{Y}(f) where :math:`\textbf{w}_{\text{bf}}(f)` is the MVDR beamforming weight for the :math:`f`-th frequency bin. The beamforming weight is computed by: .. math:: \textbf{w}_{\text{MVDR}}(f) = \frac{{{\bf{\Phi}_{\textbf{NN}}^{-1}}(f){\bf{\Phi}_{\textbf{SS}}}}(f)} {\text{Trace}({{{\bf{\Phi}_{\textbf{NN}}^{-1}}(f) \bf{\Phi}_{\textbf{SS}}}(f))}}\bm{u} Trrr&r r rdr rrrc Cre)a Args: specgram (torch.Tensor): Multi-channel complex-valued spectrum. Tensor with dimensions `(..., channel, freq, time)`. psd_s (torch.Tensor): The complex-valued power spectral density (PSD) matrix of target speech. Tensor with dimensions `(..., freq, channel, channel)`. psd_n (torch.Tensor): The complex-valued power spectral density (PSD) matrix of noise. Tensor with dimensions `(..., freq, channel, channel)`. reference_channel (int or torch.Tensor): Specifies the reference channel. If the dtype is ``int``, it represents the reference channel index. If the dtype is ``torch.Tensor``, its shape is `(..., channel)`, where the ``channel`` dimension is one-hot. diagonal_loading (bool, optional): If ``True``, enables applying diagonal loading to ``psd_n``. (Default: ``True``) diag_eps (float, optional): The coefficient multiplied to the identity matrix for diagonal loading. It is only effective when ``diagonal_loading`` is set to ``True``. (Default: ``1e-7``) eps (float, optional): Value to add to the denominator in the beamforming weight formula. (Default: ``1e-8``) Returns: torch.Tensor: Single-channel complex-valued enhanced spectrum with dimensions `(..., freq, time)`. )rrrZ) r#r&r r rdr rrr]rfrrrr-rgzSoudenMVDR.forwardNrh) r.r/r0r1rrr`r2r3r4r-rrrrris,  rir_)rVtypingrrr4r torchaudiorr__all__rar2r3rnnModulerr6rbrirrrrs@   )*w@