o  i( @s^ddlmZddlZddlmmZddlmZddl m Z ddl m Z Gddde Z dS)) annotationsN)fftn)_Loss) LossReductioncs@eZdZdZddejfdfdd ZdddZdddZZ S) JukeboxLossa Calculate spectral component based on the magnitude of Fast Fourier Transform (FFT). Based on: Dhariwal, et al. 'Jukebox: A generative model for music.' https://arxiv.org/abs/2005.00341 Args: spatial_dims: number of spatial dimensions. fft_signal_size: signal size in the transformed dimensions. See torch.fft.fftn() for more information. fft_norm: {``"forward"``, ``"backward"``, ``"ortho"``} Specifies the normalization mode in the fft. See torch.fft.fftn() for more information. reduction: {``"none"``, ``"mean"``, ``"sum"``} Specifies the reduction to apply to the output. Defaults to ``"mean"``. - ``"none"``: no reduction will be applied. - ``"mean"``: the sum of the output will be divided by the number of elements in the output. - ``"sum"``: the output will be summed. Northo spatial_dimsintfft_signal_sizetuple[int] | Nonefft_normstr reductionLossReduction | strreturnNonecs>tjt|jd||_||_ttd|d|_||_ dS)Nr) super__init__rvaluerr tuplerangefft_dimr )selfrr r r __class__\/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/losses/spectral_loss.pyr+s  zJukeboxLoss.__init__input torch.TensortargetcCsl||}||}tj||dd}|jtjjkr|}|S|jtjjkr,| }|S|jtj jkr4 |S)Nnoner) _get_fft_amplitudeFmse_lossrrMEANrmeanSUMsumNONE)rr r"Zinput_amplitudeZtarget_amplitudelossrrrforward9s  zJukeboxLoss.forwardimagescCs<t||j|j|jd}tt|dt|d}|S)z Calculate the amplitude of the fourier transformations representation of the images Args: images: Images that are to undergo fftn Returns: fourier transformation amplitude )sdimnormr)rr rr torchsqrtrealimag)rr.Zimg_fft amplituderrrr$Js "zJukeboxLoss._get_fft_amplitude) rr r r r r rrrr)r r!r"r!rr!)r.r!rr!) __name__ __module__ __qualname____doc__rr'rr-r$ __classcell__rrrrrs r) __future__rr2torch.nn.functionalnn functionalr%Z torch.fftrtorch.nn.modules.lossr monai.utilsrrrrrrs