U Phc8@sddlmZddlZddlmZddlmZddlmZddl Z ddl m Z ddl m Z m Z ddlmZdd lmZdd lmZdd lmZdd lmZdd lmZmZddddgZGdddZeZeeeeGddde 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&eZdZdZddZddddZdS) 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)selfrV/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/networks/nets/flexible_unet.py__init__*sz FlexUNetEncoderRegister.__init__ztype[Any] | strnamec Cst|trHtd|d\}}|s*t|}|}t|tsHtd|dt|tsbt |d| }| }| }| }t|t|krt|krt|ksntt|D],\}} |||||||d} | |j| <qdS)z Register a given class to the encoder dict. Please notice that input class must be a subclass of BaseEncoder. zmonai.networks.netsrz 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_parameterslenAssertionError enumerater) rrtmp_name has_built_inZname_string_listZfeature_number_listZfeature_channel_listparameter_listcnt name_stringZcur_dictrrrregister_class-s.    4z&FlexUNetEncoderRegister.register_classN)__name__ __module__ __qualname____doc__rr1rrrrr!sc sNeZdZdZddddddddddd dd fd d ZddddddZZS) 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. intz Sequence[int] str | tuple float | tupleboolr z str | Nonez bool | None) spatial_dimsencoder_channelsdecoder_channelsactnormdropoutbiasupsamplepre_conv interp_mode align_cornersis_padc stt|dkrtdt|t|dkr:td|dgt|dd} t|ddddddg}dgt|d}|dg}t| |||D]6\}}}}|t|||||||||| | | || d qt ||_ 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)r;in_chnscat_chnsout_chnsr>r?r@rArBrCrDrEhalvesrF) superrr)r!listappendzipr r ModuleListblocks)rr;r<r=r>r?r@rArBrCrDrErF in_channelsZ skip_channelsrMrSZin_chnZskip_chnZout_chnZhalve __class__rrrqs<    zUNetDecoder.__init__zlist[torch.Tensor])features skip_connectcCsl|ddddd}|ddddd}|d}t|jD](\}}||krX||}nd}|||}q>|S)NrIrHr)r+rS)rrXrYskipsxiblockskiprrrforwards  zUNetDecoder.forward)rWr2r3r4r5rr_ __classcell__rrrUrr6Us(/r6cs2eZdZdZd dddddddfd d 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?r7ztuple | str | Nonefloat)r;rT out_channels kernel_sizer> scale_factorc spttj|f||||dd}t}|dkrBt||ddtjd}|dk rTt|} nt} t ||| dS)NrG)rTrfrgpaddingrd nontrainable)r;rhmoderCrD) r CONVrIdentityrrLINEARr rNr) rr;rTrfrgr>rh conv_layerZup_layerZ act_layerrUrrrs&  zSegmentationHead.__init__)rcNrd)r2r3r4r5rrarrrUrrbs rbcszeZdZdZdddddddfd d d ifd dd ddd f ddddddddddddddddfdd ZddddZZS)rzN A flexible implementation of UNet-like encoder-decoder architecture. F)@ rGbatchgMbP?g?)epsmomentumreluinplaceTgrjdefaultnearestr7r r:tupler8r9None)rTrfbackbone pretrainedr=r;r?r>r@ decoder_biasrBrCrDrFreturncs&t|tjkr0td|dtjd|dkr@tdtj|}||_||_|d}d|krvd|krvd |ks~td |d }|d krtd |d|}|d|_| |||dt |gt |d}|d}|f||_ 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 .)rGrcz spatial_dims can only be 2 or 3.rr;rTrzWThe backbone init method must have spatial_dims, in_channels and pretrained parameters.rzBFlexible unet can only accept no more than 5 encoder feature maps.NrH)r;rTrrr) r;r<r=r>r?r@rArBrDrCrErFrIrc)r;rTrfrgr>)rNrrrr!keysr~r;rYupdater|rOencoderr6decoderrbsegmentation_head)rrTrfr~rr=r;r?r>r@rrBrCrDrFrZencoder_parametersZencoder_feature_numr<Z encoder_typerUrrrs`5      zFlexibleUNet.__init__z torch.Tensor)inputscCs*|}||}|||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])``. )rrrYr)rrr[enc_outZ decoder_outZx_segrrrr_Js   zFlexibleUNet.forwardr`rrrUrrs   0j)% __future__rr#collections.abcrpydocrtypingrtorchrZmonai.networks.blocksrrmonai.networks.layers.factoriesr monai.networks.layers.utilsr monai.networks.netsr Zmonai.networks.nets.basic_unetr Zmonai.networks.nets.resnetr monai.utilsrr__all__rrr1Moduler6 Sequentialrbrrrrrr s.           /  Z,