o , i6@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 SegResNetVAEc seZdZdZdddddddd ifd d difd dd d dejf d6fd"d# Zd$d%Zd&d'Zd7d(d)Z d8d.d/Z d9d2d3Z d:d4d5Z 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_groupsrrrrrr spatial_dimsint init_filters in_channels out_channels dropout_prob float | Noneact tuple | strnorm norm_namestruse_conv_finalbool blocks_downtuple blocks_up upsample_modeUpsampleMode | strcst|dvr td||_||_||_| |_| |_||_||_ t ||_ |r=| dkr7td|ddd| if}||_ t| |_| |_t||||_||_|\|_|_|||_|durrttj|f||_dSdS)N)rr z"`spatial_dims` can only be 2 or 3.groupzDeprecating option 'norm_name=z', please use 'norm' instead.r)super__init__ ValueErrorrrrr'r)rr ract_modlowerr"r r*r%rconvInit_make_down_layers down_layers_make_up_layers up_layers up_samples_make_final_conv conv_finalrDROPOUTdropout)selfrrrrrr r"r#rr%r'r)r* __class___/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/segresnet.pyr.;s2       zSegResNet.__init__cst}jjjjf\}}t|D]6\}}|d||dkr.tdddnt}tj |gfddt |DR}| |q|S)Nrr)stridecsg|] }tjdqS)r"r rr .0_Zlayer_in_channelsr"r<rr?r@ rsz/SegResNet._make_down_layers..) nn ModuleListr'rrr" enumeraterIdentity Sequentialrangeappend)r<r4r'filtersiitempre_convZ down_layerr?rGr@r3gs  zSegResNet._make_down_layersc stt}}jjjjjf\}}}t|}t|D]9}|d||| tj fddt||D| tj t dddt d|dgq"||fS)Nrcs"g|] }tdjdqS)rrBrCrDr"Zsample_in_channelsr<rr?r@rHsz-SegResNet._make_up_layers..r kernel_sizer*) rIrJr*r)rrr"lenrNrOrMrr)r<r6r7r*r)rPZn_uprQr?rTr@r5ws2  zSegResNet._make_up_layersc Cs2tt|j|j|jd|jt|j|j|dddS)NnamerchannelsrT)rVbias)rIrMr r"rrr0r)r<rr?r?r@r8s zSegResNet._make_final_convx torch.Tensorreturn'tuple[torch.Tensor, list[torch.Tensor]]cCsH||}|jdur||}g}|jD] }||}||q||fSN)r2rr;r4rO)r<r]down_xdownr?r?r@encodes     zSegResNet.encoderblist[torch.Tensor]cCsPtt|j|jD]\}\}}||||d}||}q |jr&||}|S)Nr)rKzipr7r6r%r9)r<r]rbrQupuplr?r?r@decodes   zSegResNet.decodecCs&||\}}||||}|Sra)rdreverseri)r<r]rbr?r?r@forwards zSegResNet.forward)rrrrrrrrrrr r!r"r!r#r$rrr%r&r'r(r)r(r*r+)rr)r]r^r_r`)r]r^rbrer_r^)r]r^r_r^)__name__ __module__ __qualname____doc__r NONTRAINABLEr.r3r5r8rdrirk __classcell__r?r?r=r@r s,  ,    csleZdZdZdddddddd d d d ifd ddifd ddejfd4fd)d* Zd+d,Zd5d0d1Zd2d3Z Z S)6r 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 rrrNrrTrrrrinput_image_size Sequence[int]vae_estimate_stdr&vae_default_stdfloatvae_nzrrrrrrrr str | tupler"r!r%r'r(r)r*r+c stj||||| | | | | ||d ||_d|_dt|jdfdd|jD|_||_||_||_ | | ||_ dS)N) rrrrrr r"r%r'r)r*rrcsg|]}|dqS)rr?)rEszoomr?r@rHsz)SegResNetVAE.__init__..) r-r.rssmallest_filtersrXr' fc_insizerurvrx_prepare_vae_modulesr8vae_conv_final)r<rsrurvrxrrrrrr r"r%r'r)r*r=r|r@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'rrr~npprodrrIrMr r"rr0rvae_downLinearrxvae_fc1vae_fc2vae_fc3rr*vae_fc_up_sample)r<r}Z v_filtersZtotal_elementsr?r?r@rs&  z!SegResNetVAE._prepare_vae_modules net_inputr^ vae_inputc Cs.||}|d|jj}||}t|}|d|jrH||}t |}dt |d|dt d|dd}|||}n|j }t |d}|||}||}||}|d|jg|j}||}t|j|jD] \}} ||}| |}qy||}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_rurFsoftplusmeanlogrvrr0r~rrrfr7r6rmse_loss) r<rrZx_vaeZz_meanZ z_mean_randz_sigmaZ vae_reg_lossrgrhZ vae_mse_lossvae_lossr?r?r@ _get_vae_losss0      0       zSegResNetVAE._get_vae_losscCsL|}||\}}||}|||}|jr"|||}||fS|dfSra)rdrjritrainingr)r<r]rrbrrr?r?r@rkFs  zSegResNetVAE.forward)rsrtrur&rvrwrxrrrrrrrrrrrr ryr"r!r%r&r'r(r)r(r*r+)rr^rr^) rlrmrnror rpr.rrrkrqr?r?r=r@r s(#  , ') __future__rcollections.abcrnumpyrrtorch.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