U Ph6@sddlmZddlmZddlZddlZddlmZddl mm Z ddl m Z mZmZddlmZddlmZmZddlmZdd gZGd ddejZGd d d eZdS) ) annotations)SequenceN)ResBlockget_conv_layerget_upsample_layer)Dropout) get_act_layerget_norm_layer) UpsampleMode SegResNet SegResNetVAEcseZdZdZdddddddd ifd d difd dd d dejf dddddddddddddd fdd ZddZddZddddZ d d!d"d#d$Z d d%d d&d'd(Z d d d"d)d*Z Z S)+r a SegResNet based on `3D MRI brain tumor segmentation using autoencoder regularization `_. The module does not include the variational autoencoder (VAE). The model supports 2D or 3D inputs. Args: spatial_dims: spatial dimension of the input data. Defaults to 3. init_filters: number of output channels for initial convolution layer. Defaults to 8. in_channels: number of input channels for the network. Defaults to 1. out_channels: number of output channels for the network. Defaults to 2. dropout_prob: probability of an element to be zero-ed. Defaults to ``None``. act: activation type and arguments. Defaults to ``RELU``. norm: feature normalization type and arguments. Defaults to ``GROUP``. norm_name: deprecating option for feature normalization type. num_groups: deprecating option for group norm. parameters. use_conv_final: if add a final convolution block to output. Defaults to ``True``. blocks_down: number of down sample blocks in each layer. Defaults to ``[1,2,2,4]``. blocks_up: number of up sample blocks in each layer. Defaults to ``[1,1,1]``. upsample_mode: [``"deconv"``, ``"nontrainable"``, ``"pixelshuffle"``] The mode of upsampling manipulations. Using the ``nontrainable`` modes cannot guarantee the model's reproducibility. Defaults to``nontrainable``. - ``deconv``, uses transposed convolution layers. - ``nontrainable``, uses non-trainable `linear` interpolation. - ``pixelshuffle``, uses :py:class:`monai.networks.blocks.SubpixelUpsample`. NRELUinplaceTGROUP num_groupsrrrrrrint float | None tuple | strstrbooltupleUpsampleMode | str) spatial_dims init_filters in_channels out_channels dropout_probactnorm norm_nameruse_conv_final blocks_down blocks_up upsample_modecst|dkrtd||_||_||_| |_| |_||_||_ t ||_ |rz| dkrntd|ddd| if}||_ t| |_| |_t||||_||_|\|_|_|||_|dk rttj|f||_dS)N)rr z"`spatial_dims` can only be 2 or 3.groupzDeprecating option 'norm_name=z', please use 'norm' instead.r)super__init__ ValueErrorr r!r"r)r*r$r%ract_modlowerr&r r+r(rconvInit_make_down_layers down_layers_make_up_layers up_layers up_samples_make_final_conv conv_finalrDROPOUTdropout)selfr r!r"r#r$r%r&r'rr(r)r*r+ __class__R/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/networks/nets/segresnet.pyr.;s0       zSegResNet.__init__cst}jjjjf\}}t|D]j\}}|d||dkr\tdddnt}tj |ffddt |D}| |q,|S)Nrr)stridecsg|]}tjdqS)r&r%rr%.0_Zlayer_in_channelsr&r<r r?r@ rsz/SegResNet._make_down_layers..) nn ModuleListr)r r!r& enumeraterIdentity Sequentialrangeappend)r<r4r)filtersiitempre_convZ down_layerr?rGr@r3gs  zSegResNet._make_down_layersc stt}}jjjjjf\}}}t|}t|D]r}|d||| tj fddt||D| tj t dddt d|dgqD||fS)Nrcs"g|]}tdjdqS)rrBrCrDr&Zsample_in_channelsr<r r?r@rHsz-SegResNet._make_up_layers..r kernel_sizer+) rIrJr+r*r r!r&lenrNrOrMrr)r<r6r7r+r*rPZn_uprQr?rTr@r5ws2  zSegResNet._make_up_layers)r#c Cs2tt|j|j|jd|jt|j|j|dddS)Nnamer channelsrT)rVbias)rIrMr r&r r!r0r)r<r#r?r?r@r8s zSegResNet._make_final_conv torch.Tensorz'tuple[torch.Tensor, list[torch.Tensor]])xreturncCsH||}|jdk r||}g}|jD]}||}||q(||fSN)r2r$r;r4rO)r<r^down_xdownr?r?r@encodes     zSegResNet.encodezlist[torch.Tensor])r^rar_cCsPtt|j|jD](\}\}}||||d}||}q|jrL||}|S)Nr)rKzipr7r6r(r9)r<r^rarQupuplr?r?r@decodes   zSegResNet.decodecCs&||\}}||||}|Sr`)rcreverserg)r<r^rar?r?r@forwards zSegResNet.forward)__name__ __module__ __qualname____doc__r NONTRAINABLEr.r3r5r8rcrgri __classcell__r?r?r=r@r s*  ,,  cseZdZdZdddddddd d d d ifd ddifd ddejfddddddddddddddddfdd ZddZddd d!d"Zd#d$Z Z S)%r a SegResNetVAE based on `3D MRI brain tumor segmentation using autoencoder regularization `_. The module contains the variational autoencoder (VAE). The model supports 2D or 3D inputs. Args: input_image_size: the size of images to input into the network. It is used to determine the in_features of the fc layer in VAE. vae_estimate_std: whether to estimate the standard deviations in VAE. Defaults to ``False``. vae_default_std: if not to estimate the std, use the default value. Defaults to 0.3. vae_nz: number of latent variables in VAE. Defaults to 256. Where, 128 to represent mean, and 128 to represent std. spatial_dims: spatial dimension of the input data. Defaults to 3. init_filters: number of output channels for initial convolution layer. Defaults to 8. in_channels: number of input channels for the network. Defaults to 1. out_channels: number of output channels for the network. Defaults to 2. dropout_prob: probability of an element to be zero-ed. Defaults to ``None``. act: activation type and arguments. Defaults to ``RELU``. norm: feature normalization type and arguments. Defaults to ``GROUP``. use_conv_final: if add a final convolution block to output. Defaults to ``True``. blocks_down: number of down sample blocks in each layer. Defaults to ``[1,2,2,4]``. blocks_up: number of up sample blocks in each layer. Defaults to ``[1,1,1]``. upsample_mode: [``"deconv"``, ``"nontrainable"``, ``"pixelshuffle"``] The mode of upsampling manipulations. Using the ``nontrainable`` modes cannot guarantee the model's reproducibility. Defaults to``nontrainable``. - ``deconv``, uses transposed convolution layers. - ``nontrainable``, uses non-trainable `linear` interpolation. - ``pixelshuffle``, uses :py:class:`monai.networks.blocks.SubpixelUpsample`. Fg333333?r rrrNrrTrrrrz Sequence[int]rfloatrrz str | tuplerrr)input_image_sizevae_estimate_stdvae_default_stdvae_nzr r!r"r#r$r%r&r(r)r*r+c stj||||| | | | | ||d ||_d|_dt|jdfdd|jD|_||_||_||_ | | ||_ dS)N) r r!r"r#r$r%r&r(r)r*r+rrcsg|]}|dqS)rr?)rEszoomr?r@rHsz)SegResNetVAE.__init__..) r-r.rrsmallest_filtersrXr) fc_insizersrtru_prepare_vae_modulesr8vae_conv_final)r<rrrsrtrur r!r"r#r$r%r&r(r)r*r+r=rxr@r.s,zSegResNetVAE.__init__c Csdt|jd}|j|}t|jt|j}t t |j |j |d|j t|j ||jdddt |j |j |jd|j |_t||j|_t||j|_t|j||_t t|j |j|ddt|j ||jdt |j |j |d|j |_dS)NrrrYT)rAr\rUrW)rXr)r!rrznpprodr{rIrMr r&r r0rvae_downLinearruvae_fc1vae_fc2vae_fc3rr+vae_fc_up_sample)r<ryZ v_filtersZtotal_elementsr?r?r@r|s& z!SegResNetVAE._prepare_vae_modulesr]) net_input vae_inputc Cs.||}|d|jj}||}t|}|d|jr||}t |}dt |d|dt d|dd}|||}n |j }t |d}|||}||}||}|d|jg|j}||}t|j|jD]\}} ||}| |}q||}t ||} || } | S)z Args: net_input: the original input of the network. vae_input: the input of VAE module, which is also the output of the network's encoder. Fg?rg:0yE>r)rviewr in_featurestorch randn_likerequires_grad_rsrFsoftplusmeanlogrtrr0rzr{rrdr7r6r}mse_loss) r<rrZx_vaeZz_meanZ z_mean_randZz_sigmaZ vae_reg_lossrerfZ vae_mse_lossvae_lossr?r?r@ _get_vae_losss0      0       zSegResNetVAE._get_vae_losscCsL|}||\}}||}|||}|jrD|||}||fS|dfSr`)rcrhrgtrainingr)r<r^rrarrr?r?r@riFs  zSegResNetVAE.forward) rjrkrlrmr rnr.r|rriror?r?r=r@r s&#  0,') __future__rcollections.abcrnumpyr~rtorch.nnrItorch.nn.functional functionalrZ%monai.networks.blocks.segresnet_blockrrrmonai.networks.layers.factoriesrmonai.networks.layers.utilsrr monai.utilsr __all__Moduler r r?r?r?r@ s