o , i1Q@sddlmZddlZddlmZddlmZddlZddl Z ddl m Z ddl m Z ddlmZmZmZmZddlmZmZddlmZmZd d gZddddZdddZGddde jZGddde jZGdd d e jZGdd d eZ dS)) annotationsN)Callable)Union)UpSample)ActConvNorm split_args) get_act_layerget_norm_layer) UpsampleMode has_option SegResNetDS SegResNetDS2 resolution tuple | listn_stages int | Nonecst|}t|}t|dkstdttt||tj fddt tD}|rH|tkrH|d|g|t}|S|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).0inlb/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/segresnet_ds.py 2s*z)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|] }|dkr dndqS)rrrkr/rrrAs z aniso_kernel..cSsg|]}|dqSr rr1rrrrBs)r*r!)r/ kernel_sizepaddingrr3r aniso_kernel:s r6cs0eZdZdZ  ddfdd ZddZZS) SegResBlockz Residual network block used SegResNet based on `3D MRI brain tumor segmentation using autoencoder regularization `_. r0relu spatial_dimsint in_channelsnorm tuple | strr4 tuple | intactreturnNonecstt|ttfrtdd|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)rNrr1rrr _z'SegResBlock.__init__..r)namer9channelsrF)r; out_channelsr4strider5biasN)super__init__ isinstancerlistr norm1r act1rCONVconv1norm2act2conv2)selfr9r;r<r4r?r5 __class__rrrJLs0      zSegResBlock.__init__c Cs8|}|||||||}||7}|SN)rSrRrQrPrNrM)rTxidentityrrrforwardys(zSegResBlock.forward)r0r8) r9r:r;r:r<r=r4r>r?r=r@rA)__name__ __module__ __qualname____doc__rJrZ __classcell__rrrUrr7Fs  -r7csHeZdZdZ        d d!fdd Zd"ddZd"ddZZS)# 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. r0 rr8batchrrrNr9r: init_filtersr;r?r=r< blocks_downr head_modulenn.Module | Noneanisotropic_scales tuple | Nonec  stdvr tdtttdfdr$dddtttddr9ddd||rCt|dnd\} } t t j f|| dd d |_ t |_tt|D]W} t } |rqt|| nd \} } fd d t|| D}t j|| d<| t|dkrt t j fdd | | d| d<nt | d<|j| d9qc||_||_||_||_|_|_|_dS)Nrrr0%`spatial_dims` can only be 1, 2 or 3.raffinerTinplace)r0rrF)r;rFr4r5rGrHr0rrc g|] }tdqS)r9r;r4r<r?r7r_r?filtersr4r<r9rrrz*SegResEncoder.__init__..blocksr)r;rFrHr4rGr5 downsample)rIrJr$r r r setdefaultrr6rrO conv_initnn ModuleListlayersr*r! ModuleDict SequentialIdentityappendrgr;rfrer<r?r9)rTr9rer;r?r<rfrgrir5rtrlevelrGrxrUrurrJs\        zSegResEncoder.__init__rX torch.Tensorr@list[torch.Tensor]cCsTg}||}|jD]}|d|}|||d|}q |jdur(||}|S)Nrxry)r{r~rrg)rTrXoutputsrrrr_forwards      zSegResEncoder._forwardcC ||SrWrrTrXrrrrZ zSegResEncoder.forward)r0rarr8rbrcNN)r9r:rer:r;r:r?r=r<r=rfrrgrhrirj)rXrr@r)r[r\r]r^rJrrZr_rrrUrr`s Hr`cs`eZdZdZ         d+d,fdd Zdd Zd!d"Zd-d'd(Zd-d)d*ZZ 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. r0rarrr8rbrcNdeconvr9r:rer;rFr?r=r<rfr blocks_uprjdsdepth preprocessnn.Module | Callable | None upsample_modeUpsampleMode | strrc  stdvr td|_||_||_||_|_|_||_ t | d|_ | |_ | |_ | durJt| ttfs=tdtdd| DsJtdtttdfd radd d tttdd rvdd d d} | rt| t|d } | |_t|||| d |_t|d}|durd|}||_|d|t|_ t!|D]j}d| rt"| t||dnd\}}t#}t$| d|ddd|d<fddt!||D}tj%||d<t||| krt&t&j'f|dd d|d<nt(|d<|j )|q|dkrAt#t(t(t&t&j'f|dd dd}|j )|dSdS)Nrkrlrzresolution must be a tuplecss|]}|dkVqdS)rNr)rrrrrrB&rCz'SegResNetDS.__init__..zresolution must be positiverrmTrn)r)r9rer;r?r<rfri)rrroF)moder9r;rFr4 scale_factorrH align_cornersupsamplec rprqrrrsrurrr[rwz(SegResNetDS.__init__..rx)r;rFr4rHhead)rrxr)*rIrJr$r9rer;rFr?r<rfr'rrrrKrLr TypeErrorr#r r rrzrr.r!rir`encoderrr|r} up_layersr*r6rrrrrOrr)rTr9rer;rFr?r<rfrrrrrrin_uprrtrGrrxrUrurrJs             zSegResNetDS.__init__cCsN|jdurdt|jdg|j}|Sttjt|jdddd}|S)zb Calculate the factors (divisors) that the input image shape must be divisible by Nrrr)axis)rir!rfr9rLrprodr")rTdrrr shape_factorvs "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#)rTrXarrris_valid_shapes"zSegResNetDS.is_valid_shaperXrr@-Union[None, torch.Tensor, list[torch.Tensor]]cCs|jdur ||}||std|jd|||}||d}t|dkr;t j d|j |j dg}g}d}|j D],}|d|}||d7}|d|}t|j ||jkrj||d||d}qB||jr|t|dkr|dS|S) NInput spatial dims  must be divisible by rrdevicedtyperrxr)rrr$rrrreversepopr!torchzerosrrrrrtraining)rTrXx_downrrrrrrrs,          zSegResNetDS._forwardcCrrWrrrrrrZrzSegResNetDS.forward r0rarrr8rbrcNrNrNr9r:rer:r;r:rFr:r?r=r<r=rfrrrjrr:rrrrrrj)rXrr@r) r[r\r]r^rJrrrrZr_rrrUrrs&r $csTeZdZdZ         d,d-fdd Z d.d/d'd(Zd0d*d+ZZS)1raL SegResNetDS2 adds an additional decorder branch to SegResNetDS and is the image encoder of VISTA3D `_. 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. r0rarrr8rbrcNrr9r:rer;rFr?r=r<rfrrrjrrrrrrc s@tj||||||||| | | | d tdd|jD|_dS)N) r9rer;rFr?r<rfrrrrrcSsg|]}t|qSr)copydeepcopy)rlayerrrrrsz)SegResNetDS2.__init__..)rIrJr|r}rup_layers_auto) rTr9rer;rFr?r<rfrrrrrrUrrrJszSegResNetDS2.__init__TrXr with_pointbool with_labelr@ctuple[Union[None, torch.Tensor, list[torch.Tensor]], Union[None, torch.Tensor, list[torch.Tensor]]]c Cs|jdur ||}||std|jd|||}||d}t|dkr;t j d|j |j dg}g}g}|}|r~|rI| }d}|jD]+} | d|}|||}| d|}t|j||jkru|| d ||d}qN||}|rd}|jD]+} | d|}|||}| d|}t|j||jkr|| d ||d}q|t|dkr|dn|t|dkr|dfS|fS) z Args: x: input tensor. with_point: if true, return the point branch output. with_label: if true, return the label branch output. Nrrrrrrrxr)rrr$rrrrrr!rrrrclonerrrr) rTrXrrrrZ outputs_autox_rrrrrrZsH               2zSegResNetDS2.forwardFcCsR|jD] }| o | |_q|jD]}| |_q|jD]}| |_q dS)z Args: auto_freeze: if true, freeze the image encoder and the auto-branch. point_freeze: if true, freeze the image encoder and the point-branch. N)r parameters requires_gradrr)rTZ auto_freezeZ point_freezeparamrrr set_auto_grads  zSegResNetDS2.set_auto_gradrr)TT)rXrrrrrr@r)FF)r[r\r]r^rJrZrr_rrrUrrs$! 9rW)rrrr)r/r)! __future__rrcollections.abcrtypingrnumpyrrtorch.nnr|monai.networks.blocks.upsamplermonai.networks.layers.factoriesrrrr monai.networks.layers.utilsr r monai.utilsr r __all__r.r6Moduler7r`rrrrrrs&        :kE