U PhW=@sddlmZddlmZddlmZddlZddlZddl m Z ddl m Z ddl mZmZmZmZddlmZmZddlmZmZd gZdd d d d dZd dddZGddde jZGddde jZGdd d e jZdS)) annotations)Callable)UnionN)UpSample)ActConvNorm split_args) get_act_layerget_norm_layer) UpsampleMode has_option SegResNetDSz tuple | listz int | None) resolutionn_stagescst|}t|}t|dks&tdttt||tj fddt tD}|r|tkr|d|g|t}n |d|}|S)aV A helper function to compute a schedule of scale at different downsampling levels, given the input resolution. .. code-block:: python scales_for_resolution(resolution=[1,1,5], n_stages=5) Args: resolution: input image resolution (in mm) n_stages: optionally the number of stages of the network rzResolution must be positivecs*g|]"}ttd|dkddqS))tuplenpwhere).0inlU/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/networks/nets/segresnet_ds.py 1sz)scales_for_resolution..rN) lenrarrayall ValueErrorfloorlog2maxastypeint32range)rrndimresscalesrrrscales_for_resolutions  " r+scalecs2fddttD}dd|D}||fS)z A helper function to compute kernel_size, padding and stride for the given scale Args: scale: scale from a current scale level cs g|]}|dkrdndqS)rrrkr,rrr@sz aniso_kernel..cSsg|] }|dqSrrr/rrrrAs)r'r)r- kernel_sizepaddingrr,r aniso_kernel9sr3cs:eZdZdZd dddddddfd d Zd d ZZS) SegResBlockz Residual network block used SegResNet based on `3D MRI brain tumor segmentation using autoencoder regularization `_. r.reluint tuple | strz tuple | intNone) spatial_dims in_channelsnormr1actreturncstt|ttfr,tdd|D}n|d}t|||d|_t||_t t j |f|||d|dd|_ t|||d|_ t||_ t t j |f|||d|dd|_dS) aY Args: spatial_dims: number of spatial dimensions, could be 1, 2 or 3. in_channels: number of input channels. norm: feature normalization type and arguments. kernel_size: convolution kernel size. Defaults to 3. act: activation type and arguments. Defaults to ``RELU``. css|]}|dVqdS)rNrr/rrr ^sz'SegResBlock.__init__..r)namer9channelsrF)r: out_channelsr1strider2biasN)super__init__ isinstancerlistr norm1r act1rCONVconv1norm2act2conv2)selfr9r:r;r1r<r2 __class__rrrEKs0      zSegResBlock.__init__c Cs8|}|||||||}||7}|SN)rNrMrLrKrIrH)rOxidentityrrrforwardxs(zSegResBlock.forward)r.r5)__name__ __module__ __qualname____doc__rErU __classcell__rrrPrr4Es  -r4c sVeZdZdZdd d d d d d d d dfdd ZdddddZdddddZZS) SegResEncodera~ SegResEncoder based on the encoder structure in `3D MRI brain tumor segmentation using autoencoder regularization `_. Args: spatial_dims: spatial dimension of the input data. Defaults to 3. init_filters: number of output channels for initial convolution layer. Defaults to 32. in_channels: number of input channels for the network. Defaults to 1. out_channels: number of output channels for the network. Defaults to 2. act: activation type and arguments. Defaults to ``RELU``. norm: feature normalization type and arguments. Defaults to ``BATCH``. blocks_down: number of downsample blocks in each layer. Defaults to ``[1,2,2,4]``. head_module: optional callable module to apply to the final features. anisotropic_scales: optional list of scale for each scale level. r. rr5batchrrrNr6r7rznn.Module | None tuple | None)r9 init_filtersr:r<r; blocks_down head_moduleanisotropic_scalesc  stdkrtdtttdfdrHdddtttddrrddd||rt|dnd\} } t t j f|| dd d |_ t |_tt|D]} t } |rt|| nd \} } fd d t|| D}t j|| d<| t|dkrVt t j fdd | | d| d<n t | d<|j| d9q||_||_||_||_|_|_|_dS)Nrrr.%`spatial_dims` can only be 1, 2 or 3.raffinerTinplace)r.rrF)r:rAr1r2rBrCr.rrc sg|]}tdqS)r9r:r1r;r<r4r_r<filtersr1r;r9rrrsz*SegResEncoder.__init__..blocksr)r:rArCr1rBr2 downsample)rDrEr!r r r setdefaultrr3rrJ conv_initnn ModuleListlayersr'r ModuleDict SequentialIdentityappendrcr:rbrar;r<r9)rOr9rar:r<r;rbrcrdr2rmrlevelrBrprPrnrrEs\        zSegResEncoder.__init__ torch.Tensorzlist[torch.Tensor]rSr=cCsTg}||}|jD]&}|d|}|||d|}q|jdk rP||}|S)Nrprq)rsrvrzrc)rOrSoutputsr{rrr_forwards      zSegResEncoder._forwardcCs ||SrRrrOrSrrrrUszSegResEncoder.forward)r.r\rr5r]r^NN)rVrWrXrYrErrUrZrrrPrr[s"Hr[csneZdZdZdd d d d d d d dd dddd fdd ZddZddZdddddZdddddZZ S) ra SegResNetDS based on `3D MRI brain tumor segmentation using autoencoder regularization `_. It is similar to https://docs.monai.io/en/stable/networks.html#segresnet, with several improvements including deep supervision and non-isotropic kernel support. Args: spatial_dims: spatial dimension of the input data. Defaults to 3. init_filters: number of output channels for initial convolution layer. Defaults to 32. in_channels: number of input channels for the network. Defaults to 1. out_channels: number of output channels for the network. Defaults to 2. act: activation type and arguments. Defaults to ``RELU``. norm: feature normalization type and arguments. Defaults to ``BATCH``. blocks_down: number of downsample blocks in each layer. Defaults to ``[1,2,2,4]``. blocks_up: number of upsample blocks (optional). dsdepth: number of levels for deep supervision. This will be the length of the list of outputs at each scale level. At dsdepth==1,only a single output is returned. preprocess: optional callable function to apply before the model's forward pass resolution: optional input image resolution. When provided, the network will first use non-isotropic kernels to bring image spacing into an approximately isotropic space. Otherwise, by default, the kernel size and downsampling is always isotropic. r.r\rrr5r]r^Ndeconvr6r7rr`znn.Module | Callable | NonezUpsampleMode | str) r9rar:rAr<r;rb blocks_updsdepth preprocess upsample_moderc  stdkrtd|_||_||_||_|_|_||_ t | d|_ | |_ | |_ | dk rt| ttfs|tdntdd| Dstdtttdfd rĈdd d tttdd rdd d d} | rt| t|d } | |_t|||| d |_t|d}|dkrDd|}||_|d|t|_ t!|D]}d| rt"| t||dnd\}}t#}t$| d|ddd|d<fddt!||D}tj%||d<t||| kr&t&t&j'f|dd d|d<n t(|d<|j )|qh|dkrt#t(t(t&t&j'f|dd dd}|j )|dS)Nrerfrzresolution must be a tuplecss|]}|dkVqdS)rNr)rrrrrr>%sz'SegResNetDS.__init__..zresolution must be positiverrgTrh)r)r9rar:r<r;rbrd)rrriF)moder9r:rAr1 scale_factorrC align_cornersupsamplec sg|]}tdqSrjrkrlrnrrrZsz(SegResNetDS.__init__..rp)r:rAr1rChead)rrpr)*rDrEr!r9rar:rAr<r;rbr$rrrrFrGr TypeErrorr r r rrrrr+rrdr[encoderrrtru up_layersr'r3rwrrxrrJryrz)rOr9rar:rAr<r;rbrrrrrrdn_uprrmrBr{rprPrnrrEs                 zSegResNetDS.__init__cCsL|jdkr&dt|jdg|j}n"ttjt|jdddd}|S)zb Calculate the factors (divisors) that the input image shape must be divisible by Nrrr)axis)rdrrbr9rGrprodr)rOdrrr shape_factorus "zSegResNetDS.shape_factorcCs*ddt|jdd|D}t|S)zx Calculate if the input shape is divisible by the minimum factors for the current network configuration cSsg|]\}}||dkqS)rr)rrjrrrrsz.SegResNetDS.is_valid_shape..rN)zipshaperr )rOrSarrris_valid_shapes"zSegResNetDS.is_valid_shaper|z-Union[None, torch.Tensor, list[torch.Tensor]]r}cCs|jdk r||}||s8td|jd|||}||d}t|dkrvt j d|j |j dg}g}d}|j D]X}|d|}||d7}|d|}t|j ||jkr||d||d}q||jrt|dkr|dS|S) NzInput spatial dims z must be divisible by rr)devicedtyperrpr)rrr!rrrreversepoprtorchzerosrrrrrztraining)rOrSZx_downr~rr{rrrrs,          zSegResNetDS._forwardcCs ||SrRrrrrrrUszSegResNetDS.forward) r.r\rrr5r]r^NrNrN) rVrWrXrYrErrrrUrZrrrPrrs$*r $)N) __future__rcollections.abcrtypingrnumpyrrtorch.nnrtmonai.networks.blocks.upsamplermonai.networks.layers.factoriesrrrr monai.networks.layers.utilsr r monai.utilsr r __all__r+r3Moduler4r[rrrrr s      :k