U Ph$@sddlmZddlmZddlZddlmZddlmZddl m Z dgZ Gdddej Z Gd d d ej ZGd d d ej ZGd ddej ZGdddej ZdS)) annotations)SequenceN) Convolution)Norm AttentionUnetcs<eZdZd ddddddfdd Zd d d d d ZZS) ConvBlockintSequence[int] | int) spatial_dims in_channels out_channels kernel_sizestridesc sVtt|||||dddtj|d t||||ddddtj|d g}tj||_dS)NNDArelu) r rrrrpadding adn_orderingactnormdropoutr )super__init__rrBATCHnn Sequentialconv)selfr rrrrrlayers __class__V/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/networks/nets/attentionunet.pyrs6  zConvBlock.__init__ torch.TensorxreturncCs||}|SN)r)rr'x_cr#r#r$forwardAs zConvBlock.forward)rr r __name__ __module__ __qualname__rr+ __classcell__r#r#r!r$rs &rcs8eZdZd ddddfdd Zddd d d ZZS) UpConvrr r )r rrc s.tt|||||ddtj|dd |_dS)NrrT)rrrrrr is_transposed)rrrrrup)rr rrrrrr!r#r$rHs zUpConv.__init__r%r&cCs||}|Sr))r4)rr'Zx_ur#r#r$r+Ws zUpConv.forward)rr2r r,r#r#r!r$r1Fsr1cs<eZdZd dddddfdd Zdddddd ZZS) AttentionBlockr r )r f_intf_gf_lc sttt|||ddd|ddttj|f||_tt|||ddd|ddttj|f||_tt||dddd|ddttj|fdt |_ t |_ dS)Nr rT)r rrrrrr conv_only) rrrrrrrW_gW_xSigmoidpsiReLUr)rr r6r7r8rr!r#r$r^sT    zAttentionBlock.__init__r%)gr'r(cCs4||}||}|||}||}||Sr))r:r;rr=)rr?r'g1x1r=r#r#r$r+s    zAttentionBlock.forward)r r,r#r#r!r$r5\s/r5cs:eZdZd dddddfdd Zd d d d d ZZS)AttentionLayerrr2r r nn.Module)r rr submodulecsTtt||||dd|_t|||||d|_t|d|||d|_||_dS)Nr2)r r7r8r6)r rrrr)r rrr) rrr5 attentionr1upconvrmergerD)rr rrrDup_kernel_sizerrr!r#r$rs* zAttentionLayer.__init__r%r&cCs:|||}|j||d}|tj||fdd}|S)N)r?r'r )dim)rFrDrErGtorchcat)rr'Z fromlowerZattZatt_mr#r#r$r+szAttentionLayer.forward)rr2r r,r#r#r!r$rBs rBc sZeZdZdZddddddddddfd d Zdddd d d dZdddddZZS)ra Attention Unet based on Otkay et al. "Attention U-Net: Learning Where to Look for the Pancreas" https://arxiv.org/abs/1804.03999 Args: spatial_dims: number of spatial dimensions of the input image. in_channels: number of the input channel. out_channels: number of the output classes. channels (Sequence[int]): sequence of channels. Top block first. The length of `channels` should be no less than 2. strides (Sequence[int]): stride to use for convolutions. kernel_size: convolution kernel size. up_kernel_size: convolution kernel size for transposed convolution layers. dropout: dropout ratio. Defaults to no dropout. rr r Sequence[int]r float)r rrchannelsrrrHrc st_|_|_|_|_|__t ||djd} t |d|ddddd} |_ ddddfd d jj} t | | | _dS) Nr)r rrrrr T)r rrrrrr9rLrC)rNrr(cst|dkrr|dd|dd}t|d|dtt|d|d|djjd|j|ddS|d|d|dSdS)Nr2r rr rrrrrr rrrDrHrr) lenrBrrrrrrH_get_bottom_layer)rNrsubblock _create_blockrrr r#r$rUs,  z-AttentionUnet.__init__.._create_block)rr dimensionsrrrNrrrrrrHrrmodel) rr rrrNrrrHrheadZreduce_channelsZencdecr!rTr$rs8  zAttentionUnet.__init__rC)rrrr(c Cs2t|j||t|j||||j|jd|j||jdS)NrOrP)rBrVrrrrH)rrrrr#r#r$rRs zAttentionUnet._get_bottom_layerr%r&cCs||}|Sr))rW)rr'Zx_mr#r#r$r+ s zAttentionUnet.forward)rrr )r-r.r/__doc__rrRr+r0r#r#r!r$rs"D) __future__rcollections.abcrrJtorch.nnr"monai.networks.blocks.convolutionsrmonai.networks.layers.factoriesr__all__Modulerr1r5rBrr#r#r#r$ s     -:#