U Ph\ @sjddlmZddlmZddlmZmZddlZddlm Z ddl m m Z ddlm Z ddlmZddlmZddlmZdd lmZdd lmZmZdd lmZmZdd lmZmZd ddddddddg ZGdd d e j Z!GdddeZ"GdddeZ#Gddde#Z$GdddeZ%Gddde j Z&GdddeZ'GdddeZ(Gddde j Z)dS)) annotations) OrderedDict)CallableSequenceN)Tensor)ADN) SimpleASPP)BackboneWithFPN) Convolution) ExtraFPNBlockFeaturePyramidNetwork)ConvNorm)ResNetResNetBottleneckAttentionModule Daf3dASPPDaf3dResNetBottleneckDaf3dResNetDilatedBottleneck Daf3dResNet Daf3dBackboneDaf3dFPNDaf3dBackboneWithFPNDAF3Dcs8eZdZdZddddfdffdd Zd d ZZS) ra Attention Module as described in 'Deep Attentive Features for Prostate Segmentation in 3D Transrectal Ultrasound' . Returns refined single layer feature (SLF) and attentive map Args: spatial_dims: dimension of inputs. in_channels: number of input channels (channels of slf and mlf). out_channels: number of output channels (channels of attentive map and refined slf). norm: normalization type. act: activation type. group @ num_groups num_channelsPRELUc sttt|||d||dt|||dd||dt|||ddddd|_tt|||d||dt|||dd||dt|||dd||d|_dS)N) kernel_sizenormact)r"paddingr#r$AZSIGMOID)r"r& adn_orderingr$)super__init__nn Sequentialr attentive_maprefine)self spatial_dims in_channels out_channelsr#r$ __class__N/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/networks/nets/daf3d.pyr*8s& zAttentionModule.__init__cCs8|t||fd}|t|||fd}||fS)Nr!)r-torchcatr.)r/slfmlfattoutr5r5r6forwardOszAttentionModule.forward__name__ __module__ __qualname____doc__r*r= __classcell__r5r5r3r6r+s  c sBeZdZdZddddd d d d d d d d fdd ZddZZS)ra Atrous Spatial Pyramid Pooling module as used in 'Deep Attentive Features for Prostate Segmentation in 3D Transrectal Ultrasound' . Core functionality as in SimpleASPP, but after each layerwise convolution a group normalization is added. Further weight initialization for convolutions is provided in _init_weight(). Additional possibility to specify the number of final output channels. Args: spatial_dims: number of spatial dimensions, could be 1, 2, or 3. in_channels: number of input channels. conv_out_channels: number of output channels of each atrous conv. out_channels: number of output channels of final convolution. If None, uses len(kernel_sizes) * conv_out_channels kernel_sizes: a sequence of four convolutional kernel sizes. Defaults to (1, 3, 3, 3) for four (dilated) convolutions. dilations: a sequence of four convolutional dilation parameters. Defaults to (1, 2, 4, 6) for four (dilated) convolutions. norm_type: final kernel-size-one convolution normalization type. Defaults to batch norm. acti_type: final kernel-size-one convolution activation type. Defaults to leaky ReLU. bias: whether to have a bias term in convolution blocks. Defaults to False. According to `Performance Tuning Guide `_, if a conv layer is directly followed by a batch norm layer, bias should be False. Raises: ValueError: When ``kernel_sizes`` length differs from ``dilations``. Nr!r%r%r%r!BATCH LEAKYRELUFint int | Nonez Sequence[int]ztuple | str | NoneboolNone) r0r1conv_out_channelsr2 kernel_sizes dilations norm_type acti_typebiasreturnc st|||||||| t} |jD]:} tddd} | | _td|dd| _| | } | | q(| |_|dkr~t ||}tdt |||d||d|_ dS)Nr!N)orderingr#norm_dimr%)r0r1r2r"r#r$) r)r*r+ ModuleListconvsr convradn _init_weightappendlenconv_k1) r/r0r1rOr2rPrQrRrSrTZ new_convs_convZtmp_convr3r5r6r*rs8       zDaf3dASPP.__init__cCs.|D] }t|tjrtjj|jq|SN)modules isinstancer+Conv3dr7initkaiming_normal_weight)r/r[mr5r5r6r]s  zDaf3dASPP._init_weight)NrDrErIrJF)r?r@rArBr*r]rCr5r5r3r6rUs!&'cs&eZdZdZdZdfdd ZZS) ra~ ResNetBottleneck block as used in 'Deep Attentive Features for Prostate Segmentation in 3D Transrectal Ultrasound' . Instead of Batch Norm Group Norm is used, instead of ReLU PReLU activation is used. Initial expansion is 2 instead of 4 and second convolution uses groups. Args: in_planes: number of input channels. planes: number of output channels (taking expansion into account). spatial_dims: number of spatial dimensions of the input image. stride: stride to use for second conv layer. downsample: which downsample layer to use. rFr%r!Nc sttj|f}ttj|f}t|tjrVt||||jd|dd|d||jd}t ||||||d|d|_ |d|d|_ |d||jd|_ |||dd|ddd|_ t|_dS)Nr!F)r"striderTrrr%)r"r&rjgroupsrT)rGROUPr CONVrdr+r, expansionr)r*bn1bn2bn3conv2PReLUrelu)r/ in_planesplanesr0rj downsamplerR conv_typer3r5r6r*s zDaf3dResNetBottleneck.__init__)r%r!N)r?r@rArBrnr*rCr5r5r3r6rscs"eZdZdZdfdd ZZS)ra- ResNetDilatedBottleneck as used in 'Deep Attentive Features for Prostate Segmentation in 3D Transrectal Ultrasound' . Same as Daf3dResNetBottleneck but dilation of 2 is used in second convolution. Args: in_planes: number of input channels. planes: number of output channels (taking expansion into account). spatial_dims: number of spatial dimensions of the input image. stride: stride to use for second conv layer. downsample: which downsample layer to use. r%r!Nc s@t|||||ttj|f}|||d|ddddd|_dS)Nr%rFrF)r"rjr&dilationrkrT)r)r*r rmrr)r/rurvr0rjrwrxr3r5r6r*sz%Daf3dResNetDilatedBottleneck.__init__)r%r!Nr?r@rArBr*rCr5r5r3r6rs cs>eZdZdZdd d d d d d d ddd d d d fdd ZZS)ra ResNet as used in 'Deep Attentive Features for Prostate Segmentation in 3D Transrectal Ultrasound' . Uses two Daf3dResNetBottleneck blocks followed by two Daf3dResNetDilatedBottleneck blocks. Args: layers: how many layers to use. block_inplanes: determine the size of planes at each step. Also tunable with widen_factor. spatial_dims: number of spatial dimensions of the input image. n_input_channels: number of input channels for first convolutional layer. conv1_t_size: size of first convolution layer, determines kernel and padding. conv1_t_stride: stride of first convolution layer. no_max_pool: bool argument to determine if to use maxpool layer. shortcut_type: which downsample block to use. Options are 'A', 'B', default to 'B'. - 'A': using `self._downsample_basic_block`. - 'B': kernel_size 1 conv + norm. widen_factor: widen output for each layer. num_classes: number of output (classifications). feed_forward: whether to add the FC layer for the output, default to `True`. bias_downsample: whether to use bias term in the downsampling block when `shortcut_type` is 'B', default to `True`. r%r!FB?T list[int]rKztuple[int] | intrMstrfloat) layersblock_inplanesr0n_input_channels conv1_t_sizeconv1_t_stride no_max_pool shortcut_type widen_factor num_classes feed_forwardbias_downsamplec stt||||||||| | | | d|_ttj|f} ttj|f}| ||jddddd|_|dd|_ t |_ | t|d|d|||_|j t|d |d ||dd |_|j t|d |d ||d d |_|j t|d |d ||d d |_dS) Nrr{)r!rFrF)r%r%r%F)r"rjr&rTrrr!)rjrFr%)r)r*rrur rmrrlconv1ror+rsrt _make_layerrlayer1layer2rlayer3layer4)r/rrr0rrrrrrrrrrxrRr3r5r6r*sj  zDaf3dResNet.__init__) r%r%r{r!Fr|r}r~TTrzr5r5r3r6rscs(eZdZdZfddZddZZS)ra  Backbone for 3D Feature Pyramid Network in DAF3D module based on 'Deep Attentive Features for Prostate Segmentation in 3D Transrectal Ultrasound' . Args: n_input_channels: number of input channels for the first convolution. csttddddgddddg|dd d }t|}tj|dd|_tj|dd |_|d |_ |d|_ |d |_ dS) Nr%rGrHrFF)rrrrrr{) r)r*rlistchildrenr+r,layer0rrrr)r/rnetZ net_modulesr3r5r6r*?s      zDaf3dBackbone.__init__cCs6||}||}||}||}||}|Srb)rrrrr)r/xrrrrrr5r5r6r=Os      zDaf3dBackbone.forwardr>r5r5r3r6r6s cs>eZdZdZd dddddfdd Zd d d d d ZZS)ra0 Feature Pyramid Network as used in 'Deep Attentive Features for Prostate Segmentation in 3D Transrectal Ultrasound' . Omits 3x3x3 convolution of layer_blocks and interpolates resulting feature maps to be the same size as feature map with highest resolution. Args: spatial_dims: 2D or 3D images in_channels_list: number of channels for each feature map that is passed to the module out_channels: number of channels of the FPN representation extra_blocks: if provided, extra operations will be performed. It is expected to take the fpn features, the original features and the names of the original features as input, and returns a new list of feature maps and their corresponding names NrKrExtraFPNBlock | None)r0in_channels_listr2 extra_blocksc sft||||t|_|D]@}|dkr4tdt|||ddddddd fd }|j|q dS) Nrz(in_channels=0 is currently not supportedr!NAr rrrr)r"r(r$r#)r)r*r+rY inner_blocks ValueErrorr r^)r/r0rr2rr1inner_block_moduler3r5r6r*is   zDaf3dFPN.__init__zdict[str, Tensor])rrUc st}t}||dd}g}||tt|dddD]F}||||}|jdd}tj ||dd} || }| d|qJ|j dk r| |||\}}|dgfdd|ddD}t tt ||} | S) NrF trilinearsizemodercs,g|]$}tj|dddddqS)feat1rFNrrF interpolater).0lrr5r6 sz$Daf3dFPN.forward..r!)rkeysvaluesget_result_from_inner_blocksr^ranger_shaperrinsertrrzip) r/rnamesx_values last_innerresultsidx inner_lateral feat_shapeinner_top_downr<r5rr6r=s     $zDaf3dFPN.forward)Nr>r5r5r3r6rXsc s4eZdZdZd ddddddd d fd d ZZS)ra Same as BackboneWithFPN but uses custom Daf3DFPN as feature pyramid network Args: backbone: backbone network return_layers: a dict containing the names of the modules for which the activations will be returned as the key of the dict, and the value of the dict is the name of the returned activation (which the user can specify). in_channels_list: number of channels for each feature map that is returned, in the order they are present in the OrderedDict out_channels: number of channels in the FPN. spatial_dims: 2D or 3D images extra_blocks: if provided, extra operations will be performed. It is expected to take the fpn features, the original features and the names of the original features as input, and returns a new list of feature maps and their corresponding names Nz nn.Modulezdict[str, str]rrKrLrrN)backbone return_layersrr2r0rrUcst|||||||dkrlt|dr. The network consists of a 3D Feature Pyramid Network which is applied on the feature maps of a 3D ResNet, followed by a custom Attention Module and an ASPP module. During training the supervised signal consists of the outputs of the FPN (four Single Layer Features, SLFs), the outputs of the attention module (four Attentive Features) and the final prediction. They are individually compared to the ground truth, the final loss consists of a weighted sum of all individual losses (see DAF3D tutorial for details). There is an additional possiblity to return all supervised signals as well as the Attentive Maps in validation mode to visualize inner functionality of the network. Args: in_channels: number of input channels. out_channels: number of output channels. visual_output: whether to return all SLFs, Attentive Maps, Refined SLFs in validation mode can be used to visualize inner functionality of the network Fcst||_tt|ddddddddd gd d d |_tjd |d d|_ddddf}dd ddf}t t d ddd d||dt d ddd dd ||dt d ddd dd ||d|_ t d dd||d|_ t d ddd d||d|_tjd|d d|_td ddddd|ddd |_dS) NrZfeat2Zfeat3Zfeat4)rrrrrrrirr%)rrrr2r0r!)r"rrrrprelug?)num_parametersrfr)r0r1r2r"r(r#r$)r0r1r2r"r(r&r#r$)r0r1r2r#r$)r"r(r#r$)r%r%r%r%))r!r!r!)r!rHrH)r! r)r!rT) r0r1rOr2rPrQrRrSrT)r)r* visual_outputrrbackbone_with_fpnr+repredict1r,r fuser attentionr.predict2raspp)r/r1r2r group_normZ act_prelur3r5r6r*s~     zDAF3D.__init__cs6t}fdd|D}t|dfdd|D}tt|\}}fdd|D}fdd|D}t|d}  | }  | } j r||| g} fdd| D} nfj rt j| ddd d } fd d|||D} | g| } nt j| ddd d } | S) Ncsg|]}|qSr5)rrr9r/r5r6r"sz!DAF3D.forward..r!csg|]}|qSr5)rr)r:r/r5r6r&scsg|]}|qSr5r)rafrr5r6r*scsg|]}|qSr5r)ramrr5r6r-scs(g|] }tj|ddddqSrFNrrrrorr5r6r7srFrrcs(g|] }tj|ddddqSrrrrr5r6r;s)rrrrr7r8tuplerr.rrtrainingrrrr)r/rZsingle_layer_featuresZ supervised1Zattentive_features_mapsZ att_featuresZatt_mapsZ supervised2Z supervised3Z attentive_mlfrZsupervised_finaloutputZsupervised_innerr5)r:r/rr6r=s*     z DAF3D.forward)Fr>r5r5r3r6rs?)* __future__r collectionsrcollections.abcrrr7torch.nnr+torch.nn.functional functionalrrZmonai.networks.blocksrZmonai.networks.blocks.aspprZ(monai.networks.blocks.backbone_fpn_utilsr "monai.networks.blocks.convolutionsr Z-monai.networks.blocks.feature_pyramid_networkr r monai.networks.layers.factoriesr rZmonai.networks.nets.resnetrr__all__Modulerrrrrrrrrr5r5r5r6 s@         *K*U"D.