text: # Copyright 2019 Google LLC
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     https://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  # DenseNet-BC 121/169/201 + composable (2016)
  # Paper: https://arxiv.org/pdf/1608.06993.pdf
  
  import tensorflow as tf
  from tensorflow.keras import Input, Model
  from tensorflow.keras.layers import ZeroPadding2D, Conv2D, MaxPooling2D, BatchNormalization, ReLU
  from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, AveragePooling2D, Concatenate
  
  class DenseNet(object):
      """ Construct a Densely Connected Convolution Neural Network """
      # Meta-parameter: number of residual blocks in each dense group
      groups = { 121 : [6, 12, 24, 16],	# DenseNet 121
                 169 : [6, 12, 32, 32],	# DenseNet 169
                 201 : [6, 12, 48, 32] }	# DenseNet 201
      # Meta-parameter: amount to reduce feature maps by (compression factor) during transition blocks
      reduction = 0.5
      # Meta-parameter: number of filters in a convolution block within a residual block (growth rate)
      n_filters = 32
      init_weights = 'he_normal'
      _model = None
  
      def __init__(self, n_layers, n_filters=32, reduction=0.5, input_shape=(224, 224, 3), n_classes=1000):
          """ Construct a Densely Connected Convolution Neural Network
              n_layers   : number of layers
              n_filters  : number of filters (growth rate)
              reduction  : anount to reduce feature maps by (compression factor)
              input_shape: input shape
              n_classes  : number of output classes
          """
          # The input vector
          inputs = Input(shape=input_shape)
  
          # The Stem Convolution Group
          x = self.stem(inputs, n_filters)
  
          # The Learner
          x = self.learner(x, self.groups[n_layers], n_filters, reduction)
  
          # Classifier for 1000 classes
          outputs = self.classifier(x, n_classes)
  
          # Instantiate the model
          self._model = Model(inputs, outputs)
  
      @property
      def model(self):
          return self._model
  
      @model.setter
      def model(self, _model):
          self._model = _model
  
      def stem(self, inputs, n_filters):
          """ Construct the Stem Convolution Group
              inputs   : input tensor
              n_filters: number of filters for the dense blocks (k)
          """
          # Pads input from 224x224 to 230x230
          x = ZeroPadding2D(padding=((3, 3), (3, 3)))(inputs)
      
          # First large convolution for abstract features for input 224 x 224 and output 112 x 112
          # Stem convolution uses 2 * k (growth rate) number of filters
          x = Conv2D(2 * n_filters, (7, 7), strides=(2, 2), use_bias=False, kernel_initializer=self.init_weights)(x)
          x = BatchNormalization()(x)
          x = ReLU()(x)
      
          # Add padding so when downsampling we fit shape 56 x 56
          x = ZeroPadding2D(padding=((1, 1), (1, 1)))(x)
          x = MaxPooling2D((3, 3), strides=2)(x)
          return x
      
      def learner(self, x, groups, n_filters, reduction):
          """ Construct the Learner
              x         : input to the learner
              blocks    : set of number of blocks per group
              n_filters : number of filters (growth rate)
              reduction : amount to reduce (compress) feature maps by
          """
          # pop off the list the last dense block
          last = groups.pop()
  
          # Create the dense groups and interceding transition blocks
          for n_blocks in groups:
              x = DenseNet.group(x, n_blocks, n_filters, reduction)
  
          # Add the last dense group w/o a following transition block
          x = DenseNet.group(x, last, n_filters)
          return x
  
      @staticmethod
      def group(x, n_blocks, n_filters, reduction=None, init_weights=None):
          """ Construct a Dense Block
              x         : input to the block
              n_blocks  : number of residual blocks in dense block
              n_filters : number of filters in convolution layer 
              reduction : amount to reduce (compress) feature maps by
          """
          # Construct a group of residual blocks
          for _ in range(n_blocks):
              x = DenseNet.residual_block(x, n_filters, init_weights=init_weights)
  
          # Construct interceding transition block
          if reduction is not None:
              x = DenseNet.trans_block(x, reduction)
          return x
  
      @staticmethod
      def residual_block(x, n_filters, init_weights=None):
          """ Construct a Residual Block
              x        : input to the block
              n_filters: number of filters in convolution layer in residual block
          """
          if init_weights is None:
              init_weights = DenseNet.init_weights
              
          # Remember input tensor into residual block
          shortcut = x 
      
          # BN-RE-Conv pre-activation form of convolutions
  
          # Dimensionality expansion, expand filters by 4 (DenseNet-B)
          x = BatchNormalization()(x)
          x = ReLU()(x)
          x = Conv2D(4 * n_filters, (1, 1), strides=(1, 1), use_bias=False, kernel_initializer=init_weights)(x)
      
          # Bottleneck convolution
          # 3x3 convolution with padding=same to preserve same shape of feature maps
          x = BatchNormalization()(x)
          x = ReLU()(x)
          x = Conv2D(n_filters, (3, 3), strides=(1, 1), padding='same', use_bias=False, kernel_initializer=init_weights)(x)
  
          # Concatenate the input (identity) with the output of the residual block
          # Concatenation (vs. merging) provides Feature Reuse between layers
          x = Concatenate()([shortcut, x])
          return x
  
      @staticmethod
      def trans_block(x, reduction, init_weights=None):
          """ Construct a Transition Block
              x        : input layer
              reduction: percentage of reduction of feature maps
          """
          if init_weights is None:
              init_weights = DenseNet.init_weights
  
          # Reduce (compress) the number of feature maps (DenseNet-C)
          # shape[n] returns a class object. We use int() to cast it into the dimension size
          n_filters = int( int(x.shape[3]) * reduction)
      
          # BN-LI-Conv pre-activation form of convolutions
  
          # Use 1x1 linear projection convolution
          x = BatchNormalization()(x)
          x = Conv2D(n_filters, (1, 1), strides=(1, 1), use_bias=False, kernel_initializer=init_weights)(x)
  
          # Use mean value (average) instead of max value sampling when pooling reduce by 75%
          x = AveragePooling2D((2, 2), strides=(2, 2))(x)
          return x
      
      def classifier(self, x, n_classes):
          """ Construct the Classifier Group
              x         : input to the classifier
              n_classes : number of output classes
          """
          # Global Average Pooling will flatten the 7x7 feature maps into 1D feature maps
          x = GlobalAveragePooling2D()(x)
          # Fully connected output layer (classification)
          x = Dense(n_classes, activation='softmax', kernel_initializer=self.init_weights)(x)
          return x
      
  # Example
  # densenet = DenseNet(121)
  