o , iO@s8ddlmZddlmZmZmZmZddlZddlm Z ddl m m Z ddl mZmZddl mZddlmZmZddlmZddlmZmZdd lmZed \ZZd gZGd d d eZGddde j Z!GdddeZ"GdddeZ#Gddde#Z$Gddde#Z%Gddde#Z&Gdd d e j'Z(dS)) annotations)OptionalSequenceTupleUnionN)ConvDenseBlock Convolution)squeeze_and_excitation)ActNorm)SkipConnection)get_dropout_layerget_pool_layer)optional_importr Quicknatcs eZdZdZfddZZS)SkipConnectionWithIdxa7 Combine the forward pass input with the result from the given submodule:: --+--submodule--o-- |_____________| The available modes are ``"cat"``, ``"add"``, ``"mul"``. Defaults to "cat" and dimension 1. Inherits from SkipConnection but provides the indizes with each forward pass. cst||fSN)superforward)selfinputindices __class__^/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/quicknat.pyr-zSkipConnectionWithIdx.forward)__name__ __module__ __qualname____doc__r __classcell__rrrrr#s rcs(eZdZdZfddZddZZS)SequentialWithIdxz A sequential container. Modules will be added to it in the order they are passed in the constructor. Own implementation to work with the new indices in the forward pass. cstj|dSrr__init__)rargsrrrr$9rzSequentialWithIdx.__init__cCs |D] }|||\}}q||fSrr)rrrmodulerrrr<szSequentialWithIdx.forwardrrrr r$rr!rrrrr"1s r"cs2eZdZdZd fdd Zd d fdd ZZS) ClassifierBlocka Returns a classifier block without an activation function at the top. It consists of a 1 * 1 convolutional layer which maps the input to a num_class channel feature map. The output is a probability map for each of the classes. Args: spatial_dims: number of spatial dimensions. in_channels: number of input channels. out_channels: number of classes to map to. strides: convolution stride. Defaults to 1. kernel_size: convolution kernel size. Defaults to 3. adn_ordering: a string representing the ordering of activation, normalization, and dropout. Defaults to "NDA". act: activation type and arguments. Defaults to PReLU. NAc st|||||||dSrr#)r spatial_dims in_channels out_channelsstrides kernel_sizeact adn_orderingrrrr$TszClassifierBlock.__init__r torch.Tensorcst|^}}}|dur0tj|dd\}}|d|dd}t|dkr,t||}|dfStdt |}|dfS)Nrdimz;Quicknat is a 2D architecture, please check your dimension.) sizetorchmaxviewlenFconv2d ValueErrorrr)rrweightsr_channeldimsZout_convrrrrWs   zClassifierBlock.forward)Nr))NN)rr1r'rrrrr(Bsr(cs<eZdZdZ    ddfdd ZddZddZZS)ConvConcatDenseBlocka This dense block is defined as a sequence of 'Convolution' blocks. It overwrite the '_get_layer' methodto change the ordering of Every convolutional layer is preceded by a batch-normalization layer and a Rectifier Linear Unit (ReLU) layer. The first two convolutional layers are followed by a concatenation layer that concatenates the input feature map with outputs of the current and previous convolutional blocks. Kernel size of two convolutional layers kept small to limit number of paramters. Appropriate padding is provided so that the size of feature maps before and after convolution remains constant. The output channels for each convolution layer is set to 64, which acts as a bottle- neck for feature map selectivity. The input channel size is variable, depending on the number of dense connections. The third convolutional layer is also preceded by a batch normalization and ReLU, but has a 1 * 1 kernel size to compress the feature map size to 64. Args: in_channles: variable depending on depth of the network seLayer: Squeeze and Excite block to be included, defaults to None, valid options are {'NONE', 'CSE', 'SSE', 'CSSE'}, dropout_layer: Dropout block to be included, defaults to None. :return: forward passed tensor N@r+intse_layerOptional[nn.Module] dropout_layerOptional[nn.Dropout2d]r.Sequence[int] | int num_filterscs`d|_tj|d|||gdd|if|d|dur|nt|_|dur)||_dSt|_dS)Nrr5instance num_features)r+r*channelsnormr.)countrr$nnIdentityrFrH)rr+rFrHr.rKrrrr${s  zConvConcatDenseBlock.__init__c Cs\|jdkr|jnd}|jd7_t|j||d||jdd|ifd}t|d|dS) a After ever convolutional layer the output is concatenated with the input and the layer before. The concatenated output is used as input to the next convolutional layer. Args: in_channels: number of input channels. out_channels: number of output channels. strides: convolution stride. is_top: True if this is the top block. r5)r4r4r4rLrM)r*r+r,r-r.r/rOadnconv)rPr.rr*r/rQ Sequential get_submodule)rr+r,dilationZ kernelsizerTrrr _get_layers   zConvConcatDenseBlock._get_layercCsd}|}|}|D]3}t|tjtjtjfrq ||}|dkr+|}tj||fdd}|dkr9tj|||fdd}|d}q ||}| |}|dfS)Nrr4r2) children isinstancerQ MaxPool2d MaxUnpool2d Dropout2dr7catrFrH)rrr?iresultresult1lrrrrs     zConvConcatDenseBlock.forward)NNrCrD) r+rErFrGrHrIr.rJrKrE)rrrr r$rXrr!rrrrrBhsrBcs0eZdZdZd fdd Zd fdd ZZS) Encodera Returns a convolution dense block for the encoding (down) part of a layer of the network. This Encoder block downpools the data with max_pool. Its output is used as input to the next layer down. New feature: it returns the indices of the max_pool to the decoder (up) path at the same layer to upsample the input. Args: in_channels: number of input channels. max_pool: predefined max_pool layer to downsample the data. se_layer: Squeeze and Excite block to be included, defaults to None. dropout: Dropout block to be included, defaults to None. kernel_size : kernel size of the convolutional layers. Defaults to 5*5 num_filters : number of input channels to each convolution block. Defaults to 64 r+rEct|||||||_dSr)rr$max_pool)rr+rerFdropoutr.rKrrrr$ zEncoder.__init__Ncs(||\}}t|d\}}||fSr)rerrrrrZ out_blockr?rrrrszEncoder.forwardr+rErr'rrrrrcsrcc.eZdZdZdfdd ZfddZZS) Decodera Returns a convolution dense block for the decoding (up) part of a layer of the network. This will upsample data with an unpool block before the forward. It uses the indices from corresponding encoder on it's level. Its output is used as input to the next layer up. Args: in_channels: number of input channels. un_pool: predefined unpool block. se_layer: predefined SELayer. Defaults to None. dropout: predefined dropout block. Defaults to None. kernel_size: Kernel size of convolution layers. Defaults to 5*5. num_filters: number of input channels to each convolution layer. Defaults to 64. r+rEcrdr)rr$un_pool)rr+rlrFrfr.rKrrrr$rgzDecoder.__init__cs&t|d\}}|||}|dfSr)rrrlrhrrrrs zDecoder.forwardrir'rrrrrksrkcrj) Bottlenecka Returns the bottom or bottleneck layer at the bottom of a network linking encoder to decoder halves. It consists of a 5 * 5 convolutional layer and a batch normalization layer to separate the encoder and decoder part of the network, restricting information flow between the encoder and decoder. Args: in_channels: number of input channels. se_layer: predefined SELayer. Defaults to None. dropout: predefined dropout block. Defaults to None. un_pool: predefined unpool block. max_pool: predefined maxpool block. kernel_size: Kernel size of convolution layers. Defaults to 5*5. num_filters: number of input channels to each convolution layer. Defaults to 64. r+rEcs$t|||||||_||_dSr)rr$rerl)rr+rFrfrerlr.rKrrrr$ s zBottleneck.__init__cs4||\}}t|d\}}|||}|dfSr)rerrrlrhrrrrs zBottleneck.forwardrir'rrrrrmsrmc sbeZdZdZddddddddddejejd f d'fdd Zd(dd Z e d!d"Z d)d%d&Z Z S)*ra Model for "Quick segmentation of NeuroAnaTomy (QuickNAT) based on a deep fully convolutional neural network. Refer to: "QuickNAT: A Fully Convolutional Network for Quick and Accurate Segmentation of Neuroanatomy by Abhijit Guha Roya, Sailesh Conjetib, Nassir Navabb, Christian Wachingera" QuickNAT has an encoder/decoder like 2D F-CNN architecture with 4 encoders and 4 decoders separated by a bottleneck layer. The final layer is a classifier block with softmax. The architecture includes skip connections between all encoder and decoder blocks of the same spatial resolution, similar to the U-Net architecture. All Encoder and Decoder consist of three convolutional layers all with a Batch Normalization and ReLU. The first two convolutional layers are followed by a concatenation layer that concatenates the input feature map with outputs of the current and previous convolutional blocks. The kernel size of the first two convolutional layers is 5*5, the third convolutional layer has a kernel size of 1*1. Data in the encode path is downsampled using max pooling layers instead of upsamling like UNet and in the decode path upsampled using max un-pooling layers instead of transpose convolutions. The pooling is done at the beginning of the block and the unpool afterwards. The indices of the max pooling in the Encoder are forwarded through the layer to be available to the corresponding Decoder. The bottleneck block consists of a 5 * 5 convolutional layer and a batch normalization layer to separate the encoder and decoder part of the network, restricting information flow between the encoder and decoder. The output feature map from the last decoder block is passed to the classifier block, which is a convolutional layer with 1 * 1 kernel size that maps the input to an N channel feature map, where N is the number of segmentation classes. To further explain this consider the first example network given below. This network has 3 layers with strides of 2 for each of the middle layers (the last layer is the bottom connection which does not down/up sample). Input data to this network is immediately reduced in the spatial dimensions by a factor of 2 by the first convolution of the residual unit defining the first layer of the encode part. The last layer of the decode part will upsample its input (data from the previous layer concatenated with data from the skip connection) in the first convolution. this ensures the final output of the network has the same shape as the input. The original QuickNAT implementation included a `enable_test_dropout()` mechanism for uncertainty estimation during testing. As the dropout layers are the only stochastic components of this network calling the train() method instead of eval() in testing or inference has the same effect. Args: num_classes: number of classes to segmentate (output channels). num_channels: number of input channels. num_filters: number of output channels for each convolutional layer in a Dense Block. kernel_size: size of the kernel of each convolutional layer in a Dense Block. kernel_c: convolution kernel size of classifier block kernel. stride_convolution: convolution stride. Defaults to 1. pool: kernel size of the pooling layer, stride_pool: stride for the pooling layer. se_block: Squeeze and Excite block type to be included, defaults to None. Valid options : NONE, CSE, SSE, CSSE, droup_out: dropout ratio. Defaults to no dropout. act: activation type and arguments. Defaults to PReLU. norm: feature normalization type and arguments. Defaults to instance norm. adn_ordering: a string representing the ordering of activation (A), normalization (N), and dropout (D). Defaults to "NA". See also: :py:class:`monai.networks.blocks.ADN`. Examples:: from monai.networks.nets import QuickNAT # network with max pooling by a factor of 2 at each layer with no se_block. net = QuickNAT( num_classes=3, num_channels=1, num_filters=64, pool = 2, se_block = "None" ) !r4rDrCr5NonerNA num_classesrE num_channelsrKr.rJkernel_c stride_convpool stride_poolse_blockstrdrop_outfloatr/Union[Tuple, str]rOr0returnc s| |_| |_| |_t|| tdd| ifddtd||dddfddtj ||d  d f dd d|_ dS)Nrfpr5)name dropout_dimr8T)r.stridereturn_indices ceil_mode)r~r*)r.rlayerrEr| nn.Modulecs|dkr |d}n t }|dkr._create_modelr4)rrEr|r) r/rOr0rr$ get_selayerr rrQr\model)rrqrrrKr.rsrtrurvrwryr/rOr0rrrr$^s  "zQuicknat.__init__cCsP|dkr td|S|dks|dkr&tstd|dkr!t|St|SdS)a@ Returns the SEBlock defined in the initialization of the QuickNAT model. Args: n_filters: encoding half of the layer se_block_type: defaults to None. Valid options are None, CSE, SSE, CSSE Returns: Appropriate SEBlock. SSE and CSSE not implemented in Monai yet. ZCSEr5SSEZCSSEzCPlease install squeeze_and_excitation locally to use SpatialSELayerN)seChannelSELayerflag ImportErrorse1ZSpatialSELayerZChannelSpatialSELayer)r n_filtersZ se_block_typerrrrs   zQuicknat.get_selayercCst|jS)zE Check if model parameters are allocated on the GPU. )next parametersis_cuda)rrrrrszQuicknat.is_cudarr1cCs||d\}}|Sr)r)rrr?rrrrszQuicknat.forward)rqrErrrErKrEr.rJrsrErtrErurErvrErwrxryrzr/r{rOr{r0rxr|ro)ro)rr1r|r1)rrrr r PRELUr INSTANCEr$rpropertyrrr!rrrrrs(G ; )) __future__rtypingrrrrr7torch.nnrQtorch.nn.functional functionalr;monai.networks.blocksrrr rmonai.networks.layers.factoriesr r "monai.networks.layers.simplelayersr monai.networks.layers.utilsr r monai.utilsrrr__all__rrUr"r(rBrcrkrmModulerrrrrs*      &]