o ) i_@@sddlmZddlmZddlZddlZddlmZddl mm Z ddl m Z ddlmZddlmZddlmZddlmZd gZGd d d eZGd d d ejZGdddejZGdddejZGdd d ejZdS)) annotations)SequenceN) Convolution)SPADE)Act) get_act_layer)StrEnumSPADENetc@seZdZdZdZdZdS)UpsamplingModesbicubicnearestbilinearN)__name__ __module__ __qualname__r r r rrc/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/spade_network.pyr sr csFeZdZdZddejddifdfdfdd ZddZddZZ S)SPADENetResBlocka, Creates a Residual Block with SPADE normalisation. Args: spatial_dims: number of spatial dimensions in_channels: number of input channels out_channels: number of output channels label_nc: number of semantic channels that will be taken into account in SPADE normalisation blocks spade_intermediate_channels: number of intermediate channels in the middle conv. layers in SPADE normalisation blocks norm: base normalisation type used on top of SPADE kernel_size: convolutional kernel size INSTANCEnegative_slope皙? spatial_dimsint in_channels out_channelslabel_ncspade_intermediate_channelsnorm str | tupleact kernel_sizec st||_||_t|j|j|_|j|jk|_t||j|jddd|_t||j|jddd|_ t ||_ t ||j||||d|_ t ||j||||d|_|jrmt||j|jdddd|_t ||j||||d|_dSdS)N)rrrr!r)rnorm_ncr"rhidden_channelsr)rrrr!rr")super__init__rrminZ int_channelslearned_shortcutrconv_0conv_1r activationrnorm_0norm_1conv_snorm_s) selfrrrrrrr!r" __class__rrr'2sd    zSPADENetResBlock.__init__cCsH|||}|||||}|||||}||}|SN)shortcutr*r,r-r+r.)r1xsegx_sdxoutrrrforwardps zSPADENetResBlock.forwardcCs$|jr||||}|S|}|Sr4)r)r/r0)r1r6r7r8rrrr5ws zSPADENetResBlock.shortcut)rrrrrrrrrrrr r!r r"r) rrr__doc__r LEAKYRELUr'r;r5 __classcell__rrr2rr$s >rcsLeZdZdZddejddiffdfdd ZddZddZddZ Z S) SPADEEncoderaT Encoding branch of a VAE compatible with a SPADE-like generator Args: spatial_dims: number of spatial dimensions in_channels: number of input channels z_dim: latent space dimension of the VAE containing the image sytle information channels: number of output after each downsampling block input_shape: spatial input shape of the tensor, necessary to do the reshaping after the linear layers of the autoencoder (HxWx[D]) kernel_size: convolutional kernel size norm: normalisation layer type act: activation type rrrrrrrz_dimchannels Sequence[int] input_shaper"rr r!c  s,t|_|_|_t||krtd|t|D]\} } | dt|| dt|kr=td| | t|fq|_fddjD_ g} j} t|D]\} }| t || |d|||d|} qUt | _t jtj jdjd_t jtj jdjd_dS) N?Length of parameter input shape must match spatial_dims; got %sEach dimension of your input must be divisible by 2 ** (autoencoder depth).The shape in position %d, %d is not divisible by %d. cg|] }|dtjqSrE)lenrA.0s_r1rr z)SPADEEncoder.__init__..)rrrstridesr"rr!) in_features out_features)r&r'rr@rArI ValueError enumeraterClatent_spatial_shapeappendrnn ModuleListblocksLinearnpprodfc_mufc_var)r1rrr@rArCr"rr!s_indrLrZZch_init_ch_valuer2rMrr'sL       zSPADEEncoder.__init__cCsB|jD]}||}q||dd}||}||}||fSNrrQ)rZviewsizer^r_r1r6blockmulogvarrrrr;s    zSPADEEncoder.forwardcCsF|jD]}||}q||dd}||}||}|||Src)rZrdrer^r_reparameterizerfrrrencodes     zSPADEEncoder.encodecCs&td|}t|}|||S)Ng?)torchexp randn_likemul)r1rhristdepsrrrrjs zSPADEEncoder.reparameterize)rrrrr@rrArBrCrBr"rrr r!r ) rrrr<rr=r'r;rkrjr>rrr2rr?s 1r?csXeZdZdZddddejddifejddifdejjfd$fdd Z d%d&d"d#Z Z S)' SPADEDecodera Decoder branch of a SPADE-like generator. It can be used independently, without an encoding branch, behaving like a GAN, or coupled to a SPADE encoder. Args: label_nc: number of semantic labels spatial_dims: number of spatial dimensions out_channels: number of output channels label_nc: number of semantic channels used for the SPADE normalisation blocks input_shape: spatial input shape of the tensor, necessary to do the reshaping after the linear layers channels: number of output after each downsampling block z_dim: latent space dimension of the VAE containing the image sytle information (None if encoder is not used) is_vae: whether the decoder is going to be coupled to an autoencoder or not (true: yes, false: no) spade_intermediate_channels: number of channels in the intermediate layers of the SPADE normalisation blocks norm: base normalisation type act: activation layer type last_act: activation layer type for the last layer of the network (can differ from previous) kernel_size: convolutional kernel size upsampling_mode: upsampling mode (nearest, bilinear etc.) NTrrrrrrrrrrCrBrA list[int]r@ int | Noneis_vaeboolrrr r!last_actstr | tuple | Noner"upsampling_modestrc st|_|_|_|_t||krtd|t|D]\}}|dt||dt|kr@td||t|fq!fdd|D_ jsZt |||d| d_ njre|duretdt |tj |d_|_g}|jtj jd| d _t|dd D]\}}|t||||d ||| | | d qtj |_t ||d || d d| d| d _dS)NrDrErFcrGrH)rI num_channelsrJrMrrrN rOz)SPADEDecoder.__init__..r)rrrr"zqIf the network is used in VAE-GAN mode, parameter z_dim (number of latent channels in the VAE) must be populated.) scale_factormoderQr%)rrrrrrr"r!)rrrpaddingr"rr!)r&r'rurrr{rIrTrUrVr conv_initrXr[r\r]fcr@rWrlUpsample upsamplingrrYrZ last_conv)r1rrrrCrAr@rurrr!rwr"ryr`rLrZZch_indrbr2rMrr'sh            zSPADEDecoder.__init__ztorch.Tensor | NonecCs|jstj|t|jd}||}n+|dur,|jdur,tj| d|jtj | d}| |}|j d|jdg|j}|jD] }|||}||}qA||}|S)a  Args: seg: input BxCxHxW[xD] semantic map on which the output is conditioned on z: latent vector output by the encoder if self.is_vae is True. When is_vae is False, z is a random noise vector. Returns: )reNr)dtypedevicerQ)ruF interpolatetuplerVrr@rlrandnrefloat32 get_devicerrdr{rZrr)r1r7rr6 res_blockrrrr;4s        zSPADEDecoder.forward)rrrrrrrCrBrArsr@rtrurvrrrr r!r rwrxr"rryrzr4)rr) rrrr<rr=r r valuer'r;r>rrr2rrrs  HrrcsneZdZdZddddejddifejddifdejjfd,fdd Z d-d.d%d&Z d/d'd(Z d-d0d*d+Z Z S)1r a SPADE Network, implemented based on the code by Park, T et al. in "Semantic Image Synthesis with Spatially-Adaptive Normalization" (https://github.com/NVlabs/SPADE) Args: spatial_dims: number of spatial dimensions in_channels: number of input channels out_channels: number of output channels label_nc: number of semantic channels used for the SPADE normalisation blocks input_shape: spatial input shape of the tensor, necessary to do the reshaping after the linear layers channels: number of output after each downsampling block z_dim: latent space dimension of the VAE containing the image sytle information (None if encoder is not used) is_vae: whether the decoder is going to be coupled to an autoencoder (true) or not (false) spade_intermediate_channels: number of channels in the intermediate layers of the SPADE normalisation blocks norm: base normalisation type act: activation layer type last_act: activation layer type for the last layer of the network (can differ from previous) kernel_size: convolutional kernel size upsampling_mode: upsampling mode (nearest, bilinear etc.) NTrrrrrrrrrrrCrBrArsr@rtrurvrrr r!rwrxr"ryrzcst||_||_||_||_||_||_|jr0|dur#tdn t |||||| | | d|_ |}| t |||||||| | | | | |d |_ dS)NzVThe latent space dimension mapped by parameter z_dim cannot be None is is_vae is True.)rrr@rArCr"rr!) rrrrCrAr@rurrr!rwr"ry)r&r'rurrrArrCrTr?encoderreverserrdecoder)r1rrrrrCrAr@rurrr!rwr"rydecoder_channelsr2rrr'hsJ    zSPADENet.__init__r7 torch.Tensorr6rcCsFd}|jr||\}}|j||}|||||fS|||fSr4)rurrjr)r1r7r6rz_muZz_logvarrrrr;s zSPADENet.forwardcCs|jr |j|SdSr4)rurrk)r1r6rrrrks zSPADENet.encodercCs |||Sr4)r)r1r7rrrrdecodes zSPADENet.decode)rrrrrrrrrCrBrArsr@rtrurvrrrr r!r rwrxr"rryrzr4)r7rr6r)r6r)r7rrr)rrrr<rr=r r rr'r;rkrr>rrr2rr Qs   ; ) __future__rtypingrnumpyr\rltorch.nnrXtorch.nn.functional functionalrmonai.networks.blocksrZ monai.networks.blocks.spade_normrmonai.networks.layersrmonai.networks.layers.utilsrmonai.utils.enumsr__all__r Modulerr?rrr rrrrs"       [W{