o - iT@s ddlmZddlmZddlmZddlZddlmZddl mm Z ddl m Z ddlmZddlmZddlmZmZdd lmZgd ZGd d d ejZGd ddejZGdddejZGdddejZGdddejZGdddejZ GdddejZ!e!Z"Z#Z$dS)) annotations)Sequence)OptionalN) checkpoint) Convolution)UpSample)ConvPool)ensure_tuple_rep)BasicUnetPlusPlusBasicunetPlusPlusbasicunetplusplusBasicUNetPlusPlusTranscs(eZdZdZfddZddZZS)Attention_blockz Attention Block c stt|ttj||dddddt||_ttj||dddddt||_ttj|ddddddtdt |_ tj dd|_ dS)NrT) kernel_sizestridepaddingbias)inplace) superr__init__nn SequentialConv3d BatchNorm3dW_gW_xSigmoidpsiReLUrelu)selfF_gF_lF_int __class__w/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/basic_unetplusplus_modified_trans.pyr)szAttention_block.__init__cCs\||}||}|j|jkrtj||jddddd}|||}||}||S)N trilinearF)sizemode align_corners)rrshapeF interpolater!r)r"gxg1x1rr(r(r)forward=s    zAttention_block.forward__name__ __module__ __qualname____doc__rr6 __classcell__r(r(r&r)r$s rc@seZdZdZdddZdS) MCDropout3du>MC Dropout：无论 model.train()/eval() 都启用随机失活input torch.TensorreturncCstj||jd|jdS)NT)ptrainingr)r0 dropout3drAr)r"r>r(r(r)r6NszMCDropout3d.forwardN)r>r?r@r?)r8r9r:r;r6r(r(r(r)r=Ksr=c,eZdZdZ    ddfdd ZZS)TwoConvztwo convolutions.rHrHrrrHr spatial_dimsintin_chnsout_chnsact str | tuplenormrbooldropout float | tupler Sequence[int]rOptional[tuple]rSequence[int] | intc s\tt||||||||| | d } t||||||||| d } |d| |d| dS)a Args: spatial_dims: number of spatial dimensions. in_chns: number of input channels. out_chns: number of output channels. act: activation type and arguments. norm: feature normalization type and arguments. bias: whether to have a bias term in convolution blocks. dropout: dropout ratio. Defaults to no dropout. )rNrPrRrrrstrides)rNrPrRrrrconv_0conv_1N)rrr add_module) r"rJrLrMrNrPrrRrrrrXrYr&r(r)rUs4   zTwoConv.__init__)rFrGrIr)rJrKrLrKrMrKrNrOrPrOrrQrRrSrrTrrUrrVr8r9r:r;rr<r(r(r&r)rER rEcrD)Downz-maxpooling downsampling and two convolutions.rFrGrImaxpoolrJrKrLrMrNrOrPrrQrRrSrrTrrUdownsample_modestrc  s~t| dkrtd|fdd} |d| d} n| dkr"d} ntd| t||||||||| | d } |d | d S) a Args: spatial_dims: number of spatial dimensions. in_chns: number of input channels. out_chns: number of output channels. act: activation type and arguments. norm: feature normalization type and arguments. bias: whether to have a bias term in convolution blocks. dropout: dropout ratio. Defaults to no dropout. downsample_mode: 'maxpool' or 'strideconv'. Defaults to 'maxpool'. r^MAXr*r*rr max_poolingr strideconvzUnsupported downsample mode: )rrrconvsN)rrr rZ ValueErrorrE)r"rJrLrMrNrPrrRrrr_rd conv_striderfr&r(r)rs*   z Down.__init__)rFrGrIr^)rJrKrLrKrMrKrNrOrPrOrrQrRrSrrTrrUr_r`r[r(r(r&r)r]r\r]cs&eZdZdfdd ZddZZS) TransformerEncoderr皙?csLt||_tj|||dd}tj||d|_tt dd||_ dS)NT)d_modelnheadrR batch_first) num_layersri@) rr in_channelsrTransformerEncoderLayerritransformer_encoder Parametertorchrandnpositional_encoding)r"rp num_headsrorR encoder_layerr&r(r)rs   zTransformerEncoder.__init__cCs|j\}}}}}|||dddd}|jd}|jddd} tj| |dd} | ddd} || }||} | ddd|||||} | S)Nrr*rlinear)r,r-)r/viewpermutervr0r1rr)r"r3bcdhwZx_flatseq_lenZpos_encoding_baseZpos_encoding_resized pos_encodingencodedoutputr(r(r)r6s  zTransformerEncoder.forward)rjrrkr8r9r:rr6r<r(r(r&r)ris ricsBeZdZdZ          d,d-fd$d% Zd.d*d+ZZS)/UpCatzHupsampling, concatenation with the encoder feature map, two convolutionsrFdeconvdefaultrzTrGrIFrJrKrLcat_chnsrMrNrOrPrrQrRrSupsampler`pre_convnn.Module | str | None interp_moder. bool | Nonehalvesis_padrrTrrU attentionc st| dkr| dur|}n| r|dn|}t|||d| | | | |d |_t||||||||||d |_||_||_|jrKt|||dd|_ dSdS)a Args: spatial_dims: number of spatial dimensions. in_chns: number of input channels to be upsampled. cat_chns: number of channels from the encoder. out_chns: number of output channels. act: activation type and arguments. norm: feature normalization type and arguments. bias: whether to have a bias term in convolution blocks. dropout: dropout ratio. Defaults to no dropout. 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. halves: whether to halve the number of channels during upsampling. This parameter does not work on ``nontrainable`` mode if ``pre_conv`` is `None`. is_pad: whether to pad upsampling features to fit features from encoder. Defaults to True. attention: whether to use attention gate. Defaults to False. nontrainableNr*rb)r-rrr.r)rr)r#r$r%) rrrrrErfr use_attentionrattention_gate)r"rJrLrrMrNrPrrRrrrr.rrrrrup_chnsr&r(r)rsB ,    zUpCat.__init__r3r?x_eOptional[torch.Tensor]cCs||}|dur\|}|jr|j||d}|jrNt|jd}dg|d}t|D]}|j| d|j| dkrDd||dd<q*tjj ||d}| tj ||gdd}|S| |}|S)z Args: x: features to be upsampled. x_e: optional features from the encoder, if None, this branch is not in use. N)r2r3r*rr replicatedim) rrrrlenr/rangertr functionalpadrfcat)r"r3rx_0 x_e_non_none dimensionsspir(r(r)r68s&    z UpCat.forward) rFrrrzTTTrGrIF)"rJrKrLrKrrKrMrKrNrOrPrOrrQrRrSrr`rrrr`r.rrrQrrQrrTrrUrrQ)r3r?rrr7r(r(r&r)rs Nrc sReZdZddddddddd fd d difdd d d f d&fd d! Zd'd$d%ZZS)(rrHrr*) r@rF LeakyReLUrkT)negative_sloperinstanceaffinerFrrJrKrp out_channelsfeaturesrTdeep_supervisionrQrNrOrPrrRrSrr` dropout_pfloatc snt||_| |_t|d} td| dt||| d|||| dd|_t|| d| d|||| dd|_ t|| d| d|||| dd|_ t|| d| d |||| dd d |_ t|| d | d |||| dd d |_ t | d d d d|_t|| d| d| d|||| | dddd |_t|| d| d| d|||| | dddd |_t|| d | d| d|||| | dddd |_t|| d | d | d |||| | dddd |_t|| d| dd| d|||| | dddd |_t|| d| dd| d|||| | dddd |_t|| d | dd| d|||| | dddd |_t|| d| dd | d|||| | dddd |_t|| d| dd | d|||| | dddd |_t|| d| dd | d|||| | dddd |_|jrx|jdkrxt|jnt|_td|f| d|dd|_td|f| d|dd|_ td|f| d|dd|_!td|f| d|dd|_"dS)a A UNet++ implementation with 1D/2D/3D supports. Based on: Zhou et al. "UNet++: A Nested U-Net Architecture for Medical Image Segmentation". 4th Deep Learning in Medical Image Analysis (DLMIA) Workshop, DOI: https://doi.org/10.48550/arXiv.1807.10165 Args: spatial_dims: number of spatial dimensions. Defaults to 3 for spatial 3D inputs. in_channels: number of input channels. Defaults to 1. out_channels: number of output channels. Defaults to 2. features: six integers as numbers of features. Defaults to ``(32, 32, 64, 128, 256, 32)``, - the first five values correspond to the five-level encoder feature sizes. - the last value corresponds to the feature size after the last upsampling. deep_supervision: whether to prune the network at inference time. Defaults to False. If true, returns a list, whose elements correspond to outputs at different nodes. act: activation type and arguments. Defaults to LeakyReLU. norm: feature normalization type and arguments. Defaults to instance norm. bias: whether to have a bias term in convolution blocks. Defaults to True. According to `Performance Tuning Guide `_, if a conv layer is directly followed by a batch norm layer, bias should be False. dropout: dropout ratio. Defaults to no dropout. upsample: upsampling mode, available options are ``"deconv"``, ``"pixelshuffle"``, ``"nontrainable"``. Examples:: # for spatial 2D >>> net = BasicUNetPlusPlus(spatial_dims=2, features=(64, 128, 256, 512, 1024, 128)) # for spatial 2D, with deep supervision enabled >>> net = BasicUNetPlusPlus(spatial_dims=2, features=(64, 128, 256, 512, 1024, 128), deep_supervision=True) # for spatial 2D, with group norm >>> net = BasicUNetPlusPlus(spatial_dims=2, features=(64, 128, 256, 512, 1024, 128), norm=("group", {"num_groups": 4})) # for spatial 3D >>> net = BasicUNetPlusPlus(spatial_dims=3, features=(32, 32, 64, 128, 256, 32)) See Also - :py:class:`monai.networks.nets.BasicUNet` - :py:class:`monai.networks.nets.DynUNet` - :py:class:`monai.networks.nets.UNet` zBasicUNetPlusPlus features: .r)rHrHrcrr*rHre)rr_rj)rprwroFT)rrrconvN)#rrrrr printrEconv_0_0r]conv_1_0conv_2_0conv_3_0conv_4_0ritransformer_bottleneckr upcat_0_1 upcat_1_1 upcat_2_1 upcat_3_1 upcat_0_2 upcat_1_2 upcat_2_2 upcat_0_3 upcat_1_3 upcat_0_4r=rIdentitymc_dropout_outrfinal_conv_0_1final_conv_0_2final_conv_0_3final_conv_0_4) r"rJrprrrrNrPrrRrrfear&r(r)rWs A                       zBasicUNetPlusPlusTrans.__init__r3r?cCs4||}||}|||}||}|||}||tj||gdd}t|j |}| ||} | | tj||gdd} | | tj|||gdd} t|j |} t|j| } || |}||tj|| gdd}||tj||| gdd}||tj|||| gdd}||}||}|g}|S)a Args: x: input should have spatially N dimensions ``(Batch, in_channels, dim_0[, dim_1, ..., dim_N-1])``, N is defined by `dimensions`. It is recommended to have ``dim_n % 16 == 0`` to ensure all maxpooling inputs have even edge lengths. Returns: A torch Tensor of "raw" predictions in shape ``(Batch, out_channels, dim_0[, dim_1, ..., dim_N-1])``. rr)rrrrrrrtrrrrrrrrrrrrrr)r"r3x_0_0x_1_0x_0_1x_2_0x_1_1x_0_2x_3_0x_2_1x_1_2x_0_3Zx_4_0_before_transx_4_0x_3_1x_2_2x_1_3x_0_4 output_0_4rr(r(r)r6zs(           zBasicUNetPlusPlusTrans.forward)rJrKrprKrrKrrTrrQrNrOrPrOrrQrRrSrr`rr)r3r?rr(r(r&r)rVs   %r)% __future__rcollections.abcrtypingrrttorch.nnrtorch.nn.functionalrr0torch.utils.checkpointr"monai.networks.blocks.convolutionsrmonai.networks.blocks.upsamplermonai.networks.layers.factoriesrr monai.utils.miscr __all__Moduler Dropout3dr=rrEr]rirrr r r r(r(r(r)s,       '86'oU