o , ic8@sddlmZddlZddlmZddlmZddlmZddl Z ddl m Z ddl m Z m Z ddlmZdd lmZdd lmZdd lmZdd lmZdd lmZmZgdZGdddZeZeeeeGddde jZ Gddde j!Z"Gddde jZ#e#Z$dS)) annotationsN)Sequence)locate)Any)nn) BaseEncoderUpSample)Conv) get_act_layer)EfficientNetEncoder)UpCat) ResNetEncoder)InterpolateModeoptional_import) FlexibleUNetFlexUNetFLEXUNET_BACKBONEFlexUNetEncoderRegisterc@s"eZdZdZddZd ddZdS) raz A register to regist backbones for the flexible unet. All backbones can be found in register_dict. Please notice each output of backbone must be 2x downsample in spatial dimension of last output. For example, if given a 512x256 2D image and a backbone with 4 outputs. Then spatial size of each encoder output should be 256x128, 128x64, 64x32 and 32x16. cCs i|_dS)N) register_dict)selfrc/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/flexible_unet.py__init__*s z FlexUNetEncoderRegister.__init__nametype[Any] | strc Cst|tr$td|d\}}|st|}|}t|ts$td|dt|ts1t |d| }| }| }| }t|t|krYt|krYt|ks\JJt|D]\}} |||||||d} | |j| <q`dS)z Register a given class to the encoder dict. Please notice that input class must be a subclass of BaseEncoder. zmonai.networks.nets)rz Cannot find z class.zl would better be derived from monai.networks.blocks.BaseEncoder or implement all interfaces specified by it.)typefeature_numberfeature_channel parameterN) isinstancestrrrr ValueError issubclassrwarningswarnget_encoder_names num_outputsnum_channels_per_outputget_encoder_parameterslen enumerater) rrtmp_name has_built_inZname_string_listZfeature_number_listZfeature_channel_listparameter_listcnt name_stringZcur_dictrrrregister_class-s0    6 z&FlexUNetEncoderRegister.register_classN)rr)__name__ __module__ __qualname____doc__rr0rrrrr!srcs.eZdZdZdfdd Zdd ddZZS)! UNetDecoderaE UNet Decoder. This class refers to `segmentation_models.pytorch `_. Args: spatial_dims: number of spatial dimensions. encoder_channels: number of output channels for all feature maps in encoder. `len(encoder_channels)` should be no less than 2. decoder_channels: number of output channels for all feature maps in decoder. `len(decoder_channels)` should equal to `len(encoder_channels) - 1`. act: activation type and arguments. norm: feature normalization type and arguments. dropout: dropout ratio. bias: whether to have a bias term in convolution blocks in this decoder. upsample: upsampling mode, available options are ``"deconv"``, ``"pixelshuffle"``, ``"nontrainable"``. pre_conv: a conv block applied before upsampling. Only used in the "nontrainable" or "pixelshuffle" mode. interp_mode: {``"nearest"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``} Only used in the "nontrainable" mode. align_corners: set the align_corners parameter for upsample. Defaults to True. Only used in the "nontrainable" mode. is_pad: whether to pad upsampling features to fit the encoder spatial dims. spatial_dimsintencoder_channels Sequence[int]decoder_channelsact str | tuplenormdropout float | tuplebiasboolupsampler pre_conv str | None interp_mode align_corners bool | Noneis_padc stt|dkrtdt|t|dkrtd|dgt|dd} t|ddddddg}dgt|d}|dg}t| |||D]\}}}}|t|||||||||| | | || d qQt ||_ dS) Nz:the length of `encoder_channels` should be no less than 2.zD`len(decoder_channels)` should equal to `len(encoder_channels) - 1`.rTF)r6in_chnscat_chnsout_chnsr;r=r>r@rBrCrErFhalvesrH) superrr)r!listappendzipr r ModuleListblocks)rr6r8r:r;r=r>r@rBrCrErFrH in_channelsZ skip_channelsrOrUZin_chnZskip_chnZout_chnZhalve __class__rrrqs<    zUNetDecoder.__init__featureslist[torch.Tensor] skip_connectcCsl|ddddd}|ddddd}|d}t|jD]\}}||kr,||}nd}|||}q|S)NrKrJr)r*rU)rrZr\skipsxiblockskiprrrforwards  zUNetDecoder.forward)r6r7r8r9r:r9r;r<r=r<r>r?r@rArBr rCrDrEr rFrGrHrA)rY)rZr[r\r7r1r2r3r4rrb __classcell__rrrWrr5Us/r5cs*eZdZdZ   ddfdd ZZS)SegmentationHeada  Segmentation head. This class refers to `segmentation_models.pytorch `_. Args: spatial_dims: number of spatial dimensions. in_channels: number of input channels for the block. out_channels: number of output channels for the block. kernel_size: kernel size for the conv layer. act: activation type and arguments. scale_factor: multiplier for spatial size. Has to match input size if it is a tuple. N?r6r7rV out_channels kernel_sizer;tuple | str | None scale_factorfloatc spttj|f||||dd}t}|dkr!t||ddtjd}|dur*t|} nt} t ||| dS)NrI)rVrhripaddingrg nontrainable)r6rkmoderCrE) r CONVrIdentityrrLINEARr rPr) rr6rVrhrir;rk conv_layerZup_layerZ act_layerrWrrrs  zSegmentationHead.__init__)rfNrg) r6r7rVr7rhr7rir7r;rjrkrl)r1r2r3r4rrdrrrWrres rec sVeZdZdZdddddddfd d d ifd dd ddd f d,fd&d' Zd-d*d+ZZS).rzN A flexible implementation of UNet-like encoder-decoder architecture. F)@ rIbatchgMbP?g?)epsmomentumreluinplaceTgrndefaultnearestrVr7rhbackboner pretrainedrAr:tupler6r=r<r;r>r? decoder_biasrBrCrErHreturnNonecs*t|tjvrtd|dtjd|dvr tdtj|}||_||_|d}d|vr;d|vr;d |vs?td |d }|d krKtd |d|}|d|_| |||dt |gt |d}|d}|di||_ t |||||| | | | | d|d |_t||d|ddd|_dS)a. A flexible implement of UNet, in which the backbone/encoder can be replaced with any efficient or residual network. Currently the input must have a 2 or 3 spatial dimension and the spatial size of each dimension must be a multiple of 32 if is_pad parameter is False. Please notice each output of backbone must be 2x downsample in spatial dimension of last output. For example, if given a 512x256 2D image and a backbone with 4 outputs. Spatial size of each encoder output should be 256x128, 128x64, 64x32 and 32x16. Args: in_channels: number of input channels. out_channels: number of output channels. backbone: name of backbones to initialize, only support efficientnet and resnet right now, can be from [efficientnet-b0, ..., efficientnet-b8, efficientnet-l2, resnet10, ..., resnet200]. pretrained: whether to initialize pretrained weights. ImageNet weights are available for efficient networks if spatial_dims=2 and batch norm is used. MedicalNet weights are available for residual networks if spatial_dims=3 and in_channels=1. Default to False. decoder_channels: number of output channels for all feature maps in decoder. `len(decoder_channels)` should equal to `len(encoder_channels) - 1`,default to (256, 128, 64, 32, 16). spatial_dims: number of spatial dimensions, default to 2. norm: normalization type and arguments, default to ("batch", {"eps": 1e-3, "momentum": 0.1}). act: activation type and arguments, default to ("relu", {"inplace": True}). dropout: dropout ratio, default to 0.0. decoder_bias: whether to have a bias term in decoder's convolution blocks. upsample: upsampling mode, available options are``"deconv"``, ``"pixelshuffle"``, ``"nontrainable"``. pre_conv:a conv block applied before upsampling. Only used in the "nontrainable" or "pixelshuffle" mode, default to `default`. interp_mode: {``"nearest"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``} Only used in the "nontrainable" mode. is_pad: whether to pad upsampling features to fit features from encoder. Default to True. If this parameter is set to "True", the spatial dim of network input can be arbitrary size, which is not supported by TensorRT. Otherwise, it must be a multiple of 32. zinvalid model_name z found, must be one of .)rIrfz spatial_dims can only be 2 or 3.rr6rVrzWThe backbone init method must have spatial_dims, in_channels and pretrained parameters.rzBFlexible unet can only accept no more than 5 encoder feature maps.NrJ)r6rVrrr) r6r8r:r;r=r>r@rBrErCrFrHrKrf)r6rVrhrir;r)rPrrrr!keysrr6r\updaterrQencoderr5decoderresegmentation_head)rrVrhrrr:r6r=r;r>rrBrCrErHrZencoder_parametersZencoder_feature_numr8Z encoder_typerWrrrsZ 5     zFlexibleUNet.__init__inputs torch.TensorcCs*|}||}|||j}||}|S)as Do a typical encoder-decoder-header inference. Args: inputs: input should have spatially N dimensions ``(Batch, in_channels, dim_0[, dim_1, ..., dim_N])``, N is defined by `dimensions`. Returns: A torch Tensor of "raw" predictions in shape ``(Batch, out_channels, dim_0[, dim_1, ..., dim_N])``. )rrr\r)rrr^Zenc_outZ decoder_outZx_segrrrrbJs   zFlexibleUNet.forward)rVr7rhr7rr rrAr:rr6r7r=r<r;r<r>r?rrArBr rCr rEr rHrArr)rrrcrrrWrrs   jr)% __future__rr#collections.abcrpydocrtypingrtorchrmonai.networks.blocksrrmonai.networks.layers.factoriesr monai.networks.layers.utilsr monai.networks.netsr Zmonai.networks.nets.basic_unetr Zmonai.networks.nets.resnetr monai.utilsrr__all__rrr0Moduler5 Sequentialrerrrrrrs0          /  Z,