o , i*@sddlmZddlmZddlmZddlZddlmZddl m Z m Z ddl m Z mZddlmZgdZGd d d ejZGd d d ejZGd ddejZGdddejZeZZZdS)) annotations)Sequence)OptionalN) ConvolutionUpSample)ConvPool)ensure_tuple_rep) BasicUnet Basicunet basicunet BasicUNetc&eZdZdZ ddfdd ZZS)TwoConvztwo convolutions. spatial_dimsintin_chnsout_chnsact str | tuplenormbiasbooldropout float | tuplec sVtt|||||||dd}t|||||||dd} |d||d| dS) 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. )rrrrpaddingconv_0conv_1N)super__init__r add_module) selfrrrrrrrrr  __class__`/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/basic_unet.pyr"s  zTwoConv.__init__rrrrrrrrrrrrrrr__name__ __module__ __qualname____doc__r" __classcell__r'r'r%r(r rcr)Downz-maxpooling downsampling and two convolutions.rrrrrrrrrrrrc sLttd|fdd}t|||||||} |d||d| dS)rMAX kernel_size max_poolingconvsN)r!r"rrr#) r$rrrrrrrr7r8r%r'r(r"@s  z Down.__init__r)r*r+r'r'r%r(r2=r1r2cs<eZdZdZ       d$d%fdd Zd&d"d#ZZS)'UpCatzHupsampling, concatenation with the encoder feature map, two convolutionsrdeconvdefaultlinearTrrrcat_chnsrrrrrrrrupsamplestrpre_convnn.Module | str | None interp_mode align_corners bool | Nonehalvesis_padc snt| dkr| dur|}n| r|dn|}t|||d| | | | d|_t|||||||||_||_dS)a6 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. nontrainableNr4)moder@rBrC)r!r"rr>rr8rF)r$rrr=rrrrrr>r@rBrCrErFZup_chnsr%r'r(r"_s ( zUpCat.__init__x torch.Tensorx_eOptional[torch.Tensor]cCs||}|durXtj|tjrX|jrJt|jd}dg|d}t|D]}|j| d|j| dkr@d||dd<q&tj j ||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. Nr4rr replicate)dim)r>torchjit isinstanceTensorrFlenshaperangenn functionalpadr8cat)r$rIrKx_0 dimensionsspir'r'r(forwards    z UpCat.forward)rr:r;r<TTT)rrrrr=rrrrrrrrrrrr>r?r@rArBr?rCrDrErrFr)rIrJrKrL)r,r-r.r/r"r^r0r'r'r%r(r9\s :r9c sNeZdZddddddddfd d difdd d f d"fdd Zd#d d!ZZS)$r rr4) r`@r` LeakyReLUg?T)negative_slopeinplaceinstanceaffinerr:rr in_channels out_channelsfeatures Sequence[int]rrrrrrrr>r?c srtt|d} td| dt|||d|||||_t|| d| d|||||_t|| d| d|||||_t|| d| d|||||_ t|| d| d|||||_ t || d| d| d||||| |_ t || d| d| d||||| |_ t || d| d| d||||| |_t || d| d| d ||||| d d |_td |f| d |dd |_dS)u A UNet implementation with 1D/2D/3D supports. Based on: Falk et al. "U-Net – Deep Learning for Cell Counting, Detection, and Morphometry". Nature Methods 16, 67–70 (2019), DOI: http://dx.doi.org/10.1038/s41592-018-0261-2 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. 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 = BasicUNet(spatial_dims=2, features=(64, 128, 256, 512, 1024, 128)) # for spatial 2D, with group norm >>> net = BasicUNet(spatial_dims=2, features=(64, 128, 256, 512, 1024, 128), norm=("group", {"num_groups": 4})) # for spatial 3D >>> net = BasicUNet(spatial_dims=3, features=(32, 32, 64, 128, 256, 32)) See Also - :py:class:`monai.networks.nets.DynUNet` - :py:class:`monai.networks.nets.UNet` zBasicUNet features: .rrr4r_F)rEconvr5N)r!r"r printrrr2down_1down_2down_3down_4r9upcat_4upcat_3upcat_2upcat_1r final_conv) r$rrirjrkrrrrr>Zfear%r'r(r"s 9 &&&* zBasicUNet.__init__rIrJc Csp||}||}||}||}||}|||}|||}|||} || |} | | } | S)a Args: x: input should have spatially N dimensions ``(Batch, in_channels, dim_0[, dim_1, ..., dim_N-1])``, N is defined by `spatial_dims`. 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])``. ) rrsrtrurvrwrxryrzr{) r$rIx0x1x2x3x4u4u3u2u1logitsr'r'r(r^s         zBasicUNet.forward)rrrirrjrrkrlrrrrrrrrr>r?)rIrJ)r,r-r.r"r^r0r'r'r%r(r s  Jr ) __future__rcollections.abcrtypingrrOtorch.nnrVmonai.networks.blocksrrmonai.networks.layers.factoriesrr monai.utilsr __all__ Sequentialrr2Moduler9r r r r r'r'r'r(s    "Vh