U jPh( @s^ddlmZddlZddlmmZddlmZddl m Z ddl m Z Gddde Z dS)) annotationsN)fftn)_Loss) LossReductioncsZeZdZdZddejfdddddd fd d Zd d d d ddZd d dddZZ 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. Northointztuple[int] | NonestrzLossReduction | strNone) spatial_dimsfft_signal_sizefft_norm reductionreturncs>tjt|jd||_||_ttd|d|_||_ dS)Nr) super__init__rvaluer r tuplerangefft_dimr )selfr r r r __class__O/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/losses/spectral_loss.pyr+s zJukeboxLoss.__init__z torch.Tensor)inputtargetrcCsf||}||}tj||dd}|jtjjkr<|}n&|jtjjkrT| }n|jtj jkrb|S)Nnoner) _get_fft_amplitudeFmse_lossrrMEANrmeanSUMsumNONE)rrrZinput_amplitudeZtarget_amplitudelossrrrforward9s    zJukeboxLoss.forward)imagesrcCs<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) __name__ __module__ __qualname____doc__rr$rr*r! __classcell__rrrrrsr) __future__rr/torch.nn.functionalnn functionalr"Z torch.fftrtorch.nn.modules.lossr monai.utilsrrrrrr s