U
    PhD                     @  s   d dl mZ d dlmZ d dlmZ d dlZd dlmZ d dl	m  m
Z d dlmZ d dlmZmZ d dlmZmZ d dlmZ d	d
ddgZG dd dejZG dd dejZG dd dejZG dd dejZG dd dejZe Z Z Z!dS )    )annotations)Sequence)OptionalN)pad)ConvolutionUpSample)ConvPool)ensure_tuple_repBasicUnetPlusPlusBasicunetPlusPlusbasicunetplusplusBasicUNetPlusPlusKernelModifiedc                   @  s   e Zd ZdZdd ZdS )MCDropout3du>   MC Dropout：无论 model.train()/eval() 都启用随机失活c                 C  s   t j|| jd| jdS )NT)ptraininginplace)F	dropout3dr   r   )selfx r   b/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/networks/nets/basic_unetplusplus_origin.pyforward&   s    zMCDropout3d.forwardN)__name__
__module____qualname____doc__r   r   r   r   r   r   #   s   r   c                      s:   e Zd ZdZdddddddd	d
d
d
d
 fddZ  ZS )TwoConvztwo convolutions.           r!   r!      r#   r!   r#   r#   r#   intstr | tupleboolfloat | tupleOptional[tuple])
spatial_dimsin_chnsout_chnsactnormbiasdropoutkernel_sizepaddingstridec                   sZ   t    t|||||||||	d	}t|||||||||	d	}| d| | d| dS )  
        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.

        )r-   r.   r0   r/   r1   r2   conv_0conv_1N)super__init__r   
add_module)r   r*   r+   r,   r-   r.   r/   r0   r1   r2   r3   r5   r6   	__class__r   r   r8   -   s2    
zTwoConv.__init__)r   r    r"   r$   r   r   r   r   r8   __classcell__r   r   r:   r   r   *   s   
    r   c                      s8   e Zd ZdZddddddddd	d	d
	 fddZ  ZS )Downz-maxpooling downsampling and two convolutions.r   r    r"   r%   r&   r'   r(   r)   )	r*   r+   r,   r-   r.   r/   r0   r1   r2   c
                   sR   t    td|f dd}
t|||||||||	d	}| d|
 | d| dS )r4   MAX   rA   r#   r1   r1   r2   max_poolingconvsN)r7   r8   r	   r   r9   )r   r*   r+   r,   r-   r.   r/   r0   r1   r2   rD   rE   r:   r   r   r8   d   s    
zDown.__init__)r   r    r"   r<   r   r   r:   r   r>   a   s
   
   r>   c                      sV   e Zd ZdZdd	d	d	d	d
d
ddddddddddd fddZdddddZ  ZS )UpCatzHupsampling, concatenation with the encoder feature map, two convolutionsr   deconvdefaultlinearTr    r"   r%   r&   r'   r(   strznn.Module | str | Nonezbool | Noner)   )r*   r+   cat_chnsr,   r-   r.   r/   r0   upsamplepre_convinterp_modealign_cornershalvesis_padr1   r2   c                   sv   t    |	dkr |
dkr |}n|r,|d n|}t|||d|	|
|||d	| _t||| |||||||d	| _|| _dS )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.

        nontrainableNrA   r@   )moderM   rN   rO   r1   rC   )r7   r8   r   rL   r   rE   rQ   )r   r*   r+   rK   r,   r-   r.   r/   r0   rL   rM   rN   rO   rP   rQ   r1   r2   up_chnsr:   r   r   r8      s6    *
zUpCat.__init__torch.TensorzOptional[torch.Tensor])r   x_ec                 C  s   |  |}|dk	rtj|tjr| jrt|jd }dg|d  }t|D ]4}|j| d  |j| d  krLd||d d < qLtj	j
||d}| tj||gdd}n
| |}|S )z

        Args:
            x: features to be upsampled.
            x_e: optional features from the encoder, if None, this branch is not in use.
        NrA   r   r#   	replicatedim)rL   torchjit
isinstanceTensorrQ   lenshaperangenn
functionalr   rE   cat)r   r   rV   x_0
dimensionsspir   r   r   r      s    
 
zUpCat.forward)	r   rG   rH   rI   TTTr    r"   )r   r   r   r   r8   r   r=   r   r   r:   r   rF      s            2GrF   c                      sn   e Zd Zddddddddd	fd
ddifddddfdddddddddddd fddZddddZ  ZS )BasicUNetPlusPlusr!   r#   rA   )    ri   @         ri   F	LeakyReLUg?T)negative_sloper   instanceaffiner   rG   r%   zSequence[int]r'   r&   r(   rJ   float)r*   in_channelsout_channelsfeaturesdeep_supervisionr-   r.   r/   r0   rL   	dropout_pc                   sB  t    || _|| _t|d}td| d t|||d ||||	dd| _t||d |d ||||	dd| _	t||d |d ||||	dd| _
t||d |d	 ||||	dd| _t||d	 |d
 ||||	dd| _t||d |d |d ||||	|
ddd| _t||d |d |d ||||	|
ddd| _t||d	 |d |d ||||	|
ddd| _t||d
 |d	 |d	 ||||	|
ddd| _t||d |d d |d ||||	|
ddd| _t||d |d d |d ||||	|
ddd| _t||d	 |d d |d ||||	|
ddd| _t||d |d d	 |d ||||	|
ddd| _t||d |d d	 |d ||||	|
ddd| _t||d |d d
 |d ||||	|
ddd| _| jr| jdkrt| jnt | _td|f |d |dd| _td|f |d |dd| _td|f |d |dd| _td|f |d |dd| _ dS )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 <https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html>`_,
                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`

           zBasicUNetPlusPlus features: .r   )r!   r!      rB   r#   rA   r!      F)rP   r1      convN)!r7   r8   ru   rv   r
   printr   conv_0_0r>   conv_1_0conv_2_0conv_3_0conv_4_0rF   	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   ra   Identitymc_dropout_outr   final_conv_0_1final_conv_0_2final_conv_0_3final_conv_0_4)r   r*   rr   rs   rt   ru   r-   r.   r/   r0   rL   rv   fear:   r   r   r8      s   A












  
  
  
  zBasicUNetPlusPlus.__init__rU   )r   c                 C  s$  |  |}| |}| ||}| |}| ||}| |tj||gdd}| |}| 	||}	| 
|	tj||gdd}
| |
tj|||gdd}| |}| ||}| |tj||	gdd}| |tj|||
gdd}| |tj||||gdd}| |}| |}|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])``.
        r#   rX   )r~   r   r   r   r   r   rZ   rc   r   r   r   r   r   r   r   r   r   r   r   )r   r   x_0_0x_1_0x_0_1x_2_0x_1_1x_0_2x_3_0x_2_1x_1_2x_0_3x_4_0x_3_1x_2_2x_1_3x_0_4
output_0_4outputr   r   r   r     s&    






zBasicUNetPlusPlus.forward)r   r   r   r8   r   r=   r   r   r:   r   rh      s   
(  rh   )"
__future__r   collections.abcr   typingr   rZ   torch.nnra   torch.nn.functionalrb   r   numpyr   monai.networks.blocksr   r   monai.networks.layers.factoriesr   r	   monai.utilsr
   __all__	Dropout3dr   
Sequentialr   r>   ModulerF   rh   r   r   r   r   r   r   r   <module>   s,   7,e  B