o , i]@stddlmZddlmZddlmZmZddlmZddl Z ddl m Z ddl m m Zddl mZddlmZddlmZdd lmZdd lmZdd lmZmZdd lmZmZdd lmZddl m!Z!m"Z"gdZ#Gddde j$Z%GdddeZ&Gddde"Z'Gddde'Z(Gddde!Z)Gddde j$Z*GdddeZ+GdddeZ,Gd d!d!e j$Z-dS)") annotations) OrderedDict)CallableSequence)partialN)Tensor)ADN) SimpleASPP)BackboneWithFPN) Convolution) ExtraFPNBlockFeaturePyramidNetwork)ConvNorm)get_norm_layer)ResNetResNetBottleneck) AttentionModule 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&ASIGMOID)r$r( adn_orderingr&)super__init__nn Sequentialr attentive_maprefine)self spatial_dims in_channels out_channelsr%r& __class__[/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/daf3d.pyr-:s  zAttentionModule.__init__cCs8|t||fd}|t|||fd}||fS)Nr#)r0torchcatr1)r2slfmlfattoutr8r8r9forwardQszAttentionModule.forward__name__ __module__ __qualname____doc__r-r@ __classcell__r8r8r6r9r-s  rcs8eZdZdZ      ddfdd 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 LEAKYRELUFr3intr4conv_out_channelsr5 int | None kernel_sizes Sequence[int] dilations norm_typetuple | str | None acti_typebiasboolreturnNonec st|||||||| t} |jD]} tddd} | | _td|dd| _| | } | | q| |_|dur?t ||}tdt |||d||d|_ dS)Nr#N)orderingr%norm_dimr')r3r4r5r$r%r&) r,r-r. ModuleListconvsr convradn _init_weightappendlenconv_k1) r2r3r4rOr5rQrSrTrVrWZ new_convs_convZtmp_convr6r8r9r-ts*        zDaf3dASPP.__init__cCs.|D]}t|tjrtjj|jq|SN)modules isinstancer.Conv3dr:initkaiming_normal_weight)r2r`mr8r8r9rbs  zDaf3dASPP._init_weight)NrGrHrLrMF)r3rNr4rNrOrNr5rPrQrRrSrRrTrUrVrUrWrXrYrZ)rBrCrDrEr-rbrFr8r8r6r9rWs!'rcs6eZdZdZdZddddddifffd d 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. norm: which normalization layer to use. Defaults to group. rIr'r#Nrr rc sttj|f}tt||d}t|tjr*t||||jd|dd|||jd}t |||||||d|_ ||d|_ |||jd|_ |||dd|ddd|_ t|_dS) N)namer3r#F)r$striderW)channelsr'r)r$r(rpgroupsrW)rCONVrrrir.r/ expansionr,r-bn1bn2bn3conv2PReLUrelu) r2 in_planesplanesr3rp downsampler% conv_type norm_layerr6r8r9r-s   zDaf3dResNetBottleneck.__init__)rBrCrDrErtr-rFr8r8r6r9rs rcs2eZdZdZddddddifffdd 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#Nrr rc sBt||||||ttj|f}|||d|ddddd|_dS)Nr'rIrF)r$rpr(dilationrrrW)r,r-rrsrx)r2r{r|r3rpr}r%r~r6r8r9r-s  z%Daf3dResNetDilatedBottleneck.__init__rBrCrDrEr-rFr8r8r6r9rs rcs8eZdZdZ         ddfdd 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?Tlayers list[int]block_inplanesr3rNn_input_channels conv1_t_sizetuple[int] | intconv1_t_stride no_max_poolrX shortcut_typestr widen_factorfloat 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#rIrI)r'r'r'F)r$rpr(rWrrr#)rprIr')r,r-rr{rrsrZGROUPconv1rur.ryrz _make_layerrlayer1layer2rlayer3layer4)r2rrr3rrrrrrrrrr~rTr6r8r9r-sL   zDaf3dResNet.__init__) r'r'rr#FrrrTT)rrrrr3rNrrNrrrrrrXrrrrrrNrrXrrXrr8r8r6r9rsrcs(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. cszttgdgd|ddd}t|}tj|dd|_tj|dd|_|d|_ |d|_ |d |_ dS) N)r'rJrKr')rIF)rrrrrr'rKr) r,r-rlistchildrenr.r/layer0rrrr)r2rnetZ net_modulesr6r8r9r-Gs    zDaf3dBackbone.__init__cCs6||}||}||}||}||}|Srg)rrrrr)r2xrrrrrr8r8r9r@Ws     zDaf3dBackbone.forwardrAr8r8r6r9r>s rcs0eZdZdZ ddfd d ZdddZZS)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 Nr3rNin_channels_listrr5 extra_blocksExtraFPNBlock | Nonec sft||||t|_|D] }|dkrtdt|||ddddddd fd }|j|qdS) Nrz(in_channels=0 is currently not supportedr#NAr"rrrr)r$r+r&r%)r,r-r.r^ inner_blocks ValueErrorr rc)r2r3rr5rr4inner_block_moduler6r8r9r-qs    zDaf3dFPN.__init__rdict[str, Tensor]rYc st}t}||dd}g}||tt|dddD]#}||||}|jdd}tj ||dd} || }| d|q%|j durW| |||\}}|dgfdd|ddD}t tt ||} | S) NrI trilinearsizemodercs,g|]}tj|dddddqS)feat1rINrrF interpolater).0lrr8r9 s,z$Daf3dFPN.forward..r#)rkeysvaluesget_result_from_inner_blocksrcrangerdshaperrinsertrrzip) r2rnamesx_values last_innerresultsidx inner_lateral feat_shapeinner_top_downr?r8rr9r@s    $zDaf3dFPN.forwardrg)r3rNrrr5rNrr)rrrYrrAr8r8r6r9r`s rcs(eZdZdZ  ddfdd 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 Nbackbone nn.Module return_layersdict[str, str]rrr5rNr3rPrrrYrZcst|||||||dur6t|drt|jtr|j}nt|jtjr(d}nt|jtj r2d}nt dt |||||_ dS)Nr3rIr'zZCould not determine value of `spatial_dims` from backbone, please provide explicit value.) r,r-hasattrrir3rNrr.Conv2drjrrfpn)r2rrrr5r3rr6r8r9r-s zDaf3dBackboneWithFPN.__init__)NN)rrrrrrr5rNr3rPrrrYrZrr8r8r6r9rs rcs*eZdZdZdfdd ZddZZS)raz DAF3D network based on 'Deep Attentive Features for Prostate Segmentation in 3D Transrectal Ultrasound' . 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 Fcs t||_tt|dddddgdddd |_tjd|d d |_d d ddf}dd ddf}t t dddd d||dt dddddd ||dt dddddd ||d|_ t ddd||d|_ t dddd d||d|_tjd|d d |_tdddddd|ddd |_dS)NrZfeat2Zfeat3Zfeat4)rrrr)rrrirr')rrrr5r3r#)r$rrrrprelug?)num_parametersrkrr)r3r4r5r$r+r%r&)r3r4r5r$r+r(r%r&)r3r4r5r%r&r)r$r+r%r&)r'r'r'r'))r#r#r#)r#rKrK)r# r)r#rT) r3r4rOr5rQrSrTrVrW)r,r- visual_outputrrbackbone_with_fpnr.rjpredict1r/r fuser attentionr1predict2raspp)r2r4r5r group_normZ act_prelur6r8r9r-sv      zDAF3D.__init__cs8t}fdd|D}t|dfdd|D}tt|\}}fdd|D}fdd|D}t|d}  | }  | } j rg||| g} fdd| D} | Sj rt j| ddd d } fd d|||D} | g| } | St j| ddd d } | S) Ncg|]}|qSr8)rrr<r2r8r9r*z!DAF3D.forward..r#csg|]}|qSr8)rr)r=r2r8r9r.scrr8r)rafrr8r9r2rcrr8r)ramrr8r9r5rc(g|]}tj|ddddqSrINrrrrorr8r9r?s(rIrrcrrrrrr8r9rCs)rrrrr:r;tuplerr1rrtrainingrrrr)r2rZsingle_layer_featuresZ supervised1Zattentive_features_mapsZ att_featuresZatt_mapsZ supervised2Z supervised3Z attentive_mlfrZsupervised_finaloutputZsupervised_innerr8)r=r2rr9r@%s.      z DAF3D.forward)FrAr8r8r6r9rs?r). __future__r collectionsrcollections.abcrr functoolsrr:torch.nnr.torch.nn.functional functionalrrmonai.networks.blocksrZmonai.networks.blocks.asppr Z(monai.networks.blocks.backbone_fpn_utilsr "monai.networks.blocks.convolutionsr Z-monai.networks.blocks.feature_pyramid_networkr r monai.networks.layers.factoriesrrmonai.networks.layers.utilsrZmonai.networks.nets.resnetrr__all__Modulerrrrrrrrrr8r8r8r9s4          *K.U"D.