text: # 
  # This code is based on 
  # ESPNetV2 (https://github.com/sacmehta/ESPNetv2)
  #
  #
  
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import numpy as np
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import math
  
  from .DCNv2.dcn_v2 import DCN
  BN_MOMENTUM = 0.1
  
  class CBR(nn.Module):
      '''
      This class defines the convolution layer with batch normalization and PReLU activation
      '''
  
      def __init__(self, nIn, nOut, kSize, stride=1, groups=1):
          '''
  
          :param nIn: number of input channels
          :param nOut: number of output channels
          :param kSize: kernel size
          :param stride: stride rate for down-sampling. Default is 1
          '''
          super().__init__()
          padding = int((kSize - 1) / 2)
          self.conv = nn.Conv2d(nIn, nOut, kSize, stride=stride, padding=padding, bias=False, groups=groups)
          self.bn = nn.BatchNorm2d(nOut)
          self.act = nn.PReLU(nOut)
  
      def forward(self, input):
          '''
          :param input: input feature map
          :return: transformed feature map
          '''
          output = self.conv(input)
          # output = self.conv1(output)
          output = self.bn(output)
          output = self.act(output)
          return output
  
  
  class BR(nn.Module):
      '''
          This class groups the batch normalization and PReLU activation
      '''
  
      def __init__(self, nOut):
          '''
          :param nOut: output feature maps
          '''
          super().__init__()
          self.bn = nn.BatchNorm2d(nOut)
          self.act = nn.PReLU(nOut)
  
      def forward(self, input):
          '''
          :param input: input feature map
          :return: normalized and thresholded feature map
          '''
          output = self.bn(input)
          output = self.act(output)
          return output
  
  
  class CB(nn.Module):
      '''
         This class groups the convolution and batch normalization
      '''
  
      def __init__(self, nIn, nOut, kSize, stride=1, groups=1):
          '''
          :param nIn: number of input channels
          :param nOut: number of output channels
          :param kSize: kernel size
          :param stride: optinal stide for down-sampling
          '''
          super().__init__()
          padding = int((kSize - 1) / 2)
          self.conv = nn.Conv2d(nIn, nOut, kSize, stride=stride, padding=padding, bias=False,
                                groups=groups)
          self.bn = nn.BatchNorm2d(nOut)
  
      def forward(self, input):
          '''
  
          :param input: input feature map
          :return: transformed feature map
          '''
          output = self.conv(input)
          output = self.bn(output)
          return output
  
  
  class C(nn.Module):
      '''
      This class is for a convolutional layer.
      '''
  
      def __init__(self, nIn, nOut, kSize, stride=1, groups=1):
          '''
  
          :param nIn: number of input channels
          :param nOut: number of output channels
          :param kSize: kernel size
          :param stride: optional stride rate for down-sampling
          '''
          super().__init__()
          padding = int((kSize - 1) / 2)
          self.conv = nn.Conv2d(nIn, nOut, kSize, stride=stride, padding=padding, bias=False,
                                groups=groups)
  
      def forward(self, input):
          '''
          :param input: input feature map
          :return: transformed feature map
          '''
          output = self.conv(input)
          return output
  
  
  class CDilated(nn.Module):
      '''
      This class defines the dilated convolution.
      '''
  
      def __init__(self, nIn, nOut, kSize, stride=1, d=1, groups=1):
          '''
          :param nIn: number of input channels
          :param nOut: number of output channels
          :param kSize: kernel size
          :param stride: optional stride rate for down-sampling
          :param d: optional dilation rate
          '''
          super().__init__()
          padding = int((kSize - 1) / 2) * d
          self.conv = nn.Conv2d(nIn, nOut,kSize, stride=stride, padding=padding, bias=False,
                                dilation=d, groups=groups)
  
      def forward(self, input):
          '''
          :param input: input feature map
          :return: transformed feature map
          '''
          output = self.conv(input)
          return output
  
  class CDilatedB(nn.Module):
      '''
      This class defines the dilated convolution with batch normalization.
      '''
  
      def __init__(self, nIn, nOut, kSize, stride=1, d=1, groups=1):
          '''
          :param nIn: number of input channels
          :param nOut: number of output channels
          :param kSize: kernel size
          :param stride: optional stride rate for down-sampling
          :param d: optional dilation rate
          '''
          super().__init__()
          padding = int((kSize - 1) / 2) * d
          self.conv = nn.Conv2d(nIn, nOut,kSize, stride=stride, padding=padding, bias=False,
                                dilation=d, groups=groups)
          self.bn = nn.BatchNorm2d(nOut)
  
      def forward(self, input):
          '''
          :param input: input feature map
          :return: transformed feature map
          '''
          return self.bn(self.conv(input))
  
  
  class EESP(nn.Module):
      '''
      This class defines the EESP block, which is based on the following principle
          REDUCE ---> SPLIT ---> TRANSFORM --> MERGE
      '''
  
      def __init__(self, nIn, nOut, stride=1, k=4, r_lim=7, down_method='esp'): #down_method --> ['avg' or 'esp']
          '''
          :param nIn: number of input channels
          :param nOut: number of output channels
          :param stride: factor by which we should skip (useful for down-sampling). If 2, then down-samples the feature map by 2
          :param k: # of parallel branches
          :param r_lim: A maximum value of receptive field allowed for EESP block
          :param down_method: Downsample or not (equivalent to say stride is 2 or not)
          '''
          super().__init__()
          self.stride = stride
          n = int(nOut / k)
          n1 = nOut - (k - 1) * n
          assert down_method in ['avg', 'esp'], 'One of these is suppported (avg or esp)'
          assert n == n1, "n(={}) and n1(={}) should be equal for Depth-wise Convolution ".format(n, n1)
          self.proj_1x1 = CBR(nIn, n, 1, stride=1, groups=k)
  
          # (For convenience) Mapping between dilation rate and receptive field for a 3x3 kernel
          map_receptive_ksize = {3: 1, 5: 2, 7: 3, 9: 4, 11: 5, 13: 6, 15: 7, 17: 8}
          self.k_sizes = list()
          for i in range(k):
              ksize = int(3 + 2 * i)
              # After reaching the receptive field limit, fall back to the base kernel size of 3 with a dilation rate of 1
              ksize = ksize if ksize <= r_lim else 3
              self.k_sizes.append(ksize)
          # sort (in ascending order) these kernel sizes based on their receptive field
          # This enables us to ignore the kernels (3x3 in our case) with the same effective receptive field in hierarchical
          # feature fusion because kernels with 3x3 receptive fields does not have gridding artifact.
          self.k_sizes.sort()
          self.spp_dw = nn.ModuleList()
          for i in range(k):
              d_rate = map_receptive_ksize[self.k_sizes[i]]
              self.spp_dw.append(CDilated(n, n, kSize=3, stride=stride, groups=n, d=d_rate))
          # Performing a group convolution with K groups is the same as performing K point-wise convolutions
          self.conv_1x1_exp = CB(nOut, nOut, 1, 1, groups=k)
          self.br_after_cat = BR(nOut)
          self.module_act = nn.PReLU(nOut)
          self.downAvg = True if down_method == 'avg' else False
  
      def forward(self, input):
          '''
          :param input: input feature map
          :return: transformed feature map
          '''
  
          # Reduce --> project high-dimensional feature maps to low-dimensional space
          output1 = self.proj_1x1(input)
          output = [self.spp_dw[0](output1)]
          # compute the output for each branch and hierarchically fuse them
          # i.e. Split --> Transform --> HFF
          for k in range(1, len(self.spp_dw)):
              out_k = self.spp_dw[k](output1)
              # HFF
              out_k = out_k + output[k - 1]
              output.append(out_k)
          # Merge
          expanded = self.conv_1x1_exp( # learn linear combinations using group point-wise convolutions
              self.br_after_cat( # apply batch normalization followed by activation function (PRelu in this case)
                  torch.cat(output, 1) # concatenate the output of different branches
              )
          )
          del output
          # if down-sampling, then return the concatenated vector
          # because Downsampling function will combine it with avg. pooled feature map and then threshold it
          if self.stride == 2 and self.downAvg:
              return expanded
  
          # if dimensions of input and concatenated vector are the same, add them (RESIDUAL LINK)
          if expanded.size() == input.size():
              expanded = expanded + input
  
          # Threshold the feature map using activation function (PReLU in this case)
          return self.module_act(expanded)
  
  
  class DownSampler(nn.Module):
      '''
      Down-sampling fucntion that has three parallel branches: (1) avg pooling,
      (2) EESP block with stride of 2 and (3) efficient long-range connection with the input.
      The output feature maps of branches from (1) and (2) are concatenated and then additively fused with (3) to produce
      the final output.
      '''
  
      def __init__(self, nin, nout, k=4, r_lim=9, reinf=True):
          '''
              :param nin: number of input channels
              :param nout: number of output channels
              :param k: # of parallel branches
              :param r_lim: A maximum value of receptive field allowed for EESP block
              :param reinf: Use long range shortcut connection with the input or not.
          '''
          super().__init__()
          nout_new = nout - nin
          self.eesp = EESP(nin, nout_new, stride=2, k=k, r_lim=r_lim, down_method='avg')
          self.avg = nn.AvgPool2d(kernel_size=3, padding=1, stride=2)
          if reinf:
              self.inp_reinf = nn.Sequential(
                  CBR(config_inp_reinf, config_inp_reinf, 3, 1),
                  CB(config_inp_reinf, nout, 1, 1)
              )
          self.act =  nn.PReLU(nout)
  
      def forward(self, input, input2=None):
          '''
          :param input: input feature map
          :return: feature map down-sampled by a factor of 2
          '''
          avg_out = self.avg(input)
          eesp_out = self.eesp(input)
          output = torch.cat([avg_out, eesp_out], 1)
  
          if input2 is not None:
              #assuming the input is a square image
              # Shortcut connection with the input image
              w1 = avg_out.size(2)
              while True:
                  input2 = F.avg_pool2d(input2, kernel_size=3, padding=1, stride=2)
                  w2 = input2.size(2)
                  if w2 == w1:
                      break
              output = output + self.inp_reinf(input2)
  
          return self.act(output)
  
  def fill_fc_weights(layers):
      for m in layers.modules():
          if isinstance(m, nn.Conv2d):
              nn.init.normal_(m.weight, std=0.001)
              if m.bias is not None:
                  nn.init.constant_(m.bias, 0)
  
  def fill_up_weights(up):
      w = up.weight.data
      f = math.ceil(w.size(2) / 2)
      c = (2 * f - 1 - f % 2) / (2. * f)
      for i in range(w.size(2)):
          for j in range(w.size(3)):
              w[0, 0, i, j] = \
                  (1 - math.fabs(i / f - c)) * (1 - math.fabs(j / f - c))
      for c in range(1, w.size(0)):
          w[c, 0, :, :] = w[0, 0, :, :] 
  
  class EESPNet(nn.Module):
      '''
      This class defines the ESPNetv2 architecture
      '''
  
      # def __init__(self, classes=1000, s=1):
      def __init__(self, heads, pretrained, head_conv, s=1):
  
          super().__init__()
          reps = [0, 3, 7, 3]  # how many times EESP blocks should be repeated at each spatial level.
          channels = 3
  
          r_lim = [13, 11, 9, 7, 5]  # receptive field at each spatial level
          K = [4]*len(r_lim) # No. of parallel branches at different levels
  
          base = 64 #base configuration
          config_len = 5
          config = [base] * config_len
          base_s = 0
          for i in range(config_len):
              if i== 0:
                  base_s = int(base * s)
                  base_s = math.ceil(base_s / K[0]) * K[0]
                  config[i] = base if base_s > base else base_s
              else:
                  config[i] = base_s * pow(2, i)
          # if s <= 1.5:
          #     config.append(1024)
          # elif s <= 2.0:
          #     config.append(1280)
          else:
              ValueError('Configuration not supported')
  
  
          global config_inp_reinf
          config_inp_reinf = 3
  
          self.inplanes = 1024
          self.deconv_with_bias = False
  
          self.input_reinforcement = True # True for the shortcut connection with input
  
          assert len(K) == len(r_lim), 'Length of branching factor array and receptive field array should be the same.'
  
          self.level1 = CBR(channels, config[0], 3, 2)  # 112 L1
  
          self.level2_0 = DownSampler(config[0], config[1], k=K[0], r_lim=r_lim[0], reinf=self.input_reinforcement)  # out = 56
  
          self.level3_0 = DownSampler(config[1], config[2], k=K[1], r_lim=r_lim[1], reinf=self.input_reinforcement) # out = 28
          self.level3 = nn.ModuleList()
          for i in range(reps[1]):
              self.level3.append(EESP(config[2], config[2], stride=1, k=K[2], r_lim=r_lim[2]))
  
          self.level4_0 = DownSampler(config[2], config[3], k=K[2], r_lim=r_lim[2], reinf=self.input_reinforcement) #out = 14
          self.level4 = nn.ModuleList()
          for i in range(reps[2]):
              self.level4.append(EESP(config[3], config[3], stride=1, k=K[3], r_lim=r_lim[3]))
  
          self.level5_0 = DownSampler(config[3], config[4], k=K[3], r_lim=r_lim[3]) #7
          self.level5 = nn.ModuleList()
          for i in range(reps[3]):
              self.level5.append(EESP(config[4], config[4], stride=1, k=K[4], r_lim=r_lim[4]))
  
          # expand the feature maps using depth-wise convolution followed by group point-wise convolution
          # self.level5.append(CBR(config[4], config[4], 3, 1, groups=config[4]))
          # self.level5.append(CBR(config[4], config[5], 1, 1, groups=K[4]))
  
          # used for deconv layers
          self.deconv_layers = self._make_deconv_layer(
              3,
              [256, 128, 64],
              [4, 4, 4],
          )
  
          self.heads = heads
          # for head in self.heads:
          #     classes = self.heads[head]
          #     if head_conv > 0:
          #         fc = nn.Sequential(
          #           nn.Conv2d(config[4], config[4],
          #             kernel_size=3, padding=1, bias=True, groups=config[4]),
          #           nn.ReLU(inplace=True),
          #           nn.Conv2d(config[4], classes, 
          #             kernel_size=1, stride=1, 
          #             padding=0, bias=True))
          #         if 'hm' in head:
          #             fc[-1].bias.data.fill_(-2.19)
          #         else:
          #             fill_fc_weights(fc)
          #     else:
          #         fc = nn.Conv2d(64, classes, 
          #           kernel_size=1, stride=1, 
          #           padding=0, bias=True)
          #         if 'hm' in head:
          #             fc.bias.data.fill_(-2.19)
          #         else:
          #             fill_fc_weights(fc)
          #     self.__setattr__(head, fc)
  
          for head in self.heads:
              classes = self.heads[head]
              if head_conv > 0:
                  fc = nn.Sequential(
                    nn.Conv2d(64, head_conv,
                      kernel_size=3, padding=1, bias=True),
                    nn.ReLU(inplace=True),
                    nn.Conv2d(head_conv, classes, 
                      kernel_size=1, stride=1, 
                      padding=0, bias=True))
                  if 'hm' in head:
                      fc[-1].bias.data.fill_(-2.19)
                  else:
                      fill_fc_weights(fc)
              else:
                  fc = nn.Conv2d(64, classes, 
                    kernel_size=1, stride=1, 
                    padding=0, bias=True)
                  if 'hm' in head:
                      fc.bias.data.fill_(-2.19)
                  else:
                      fill_fc_weights(fc)
              self.__setattr__(head, fc)
  
          # init weights 
          # self.init_params()
  
      def init_params(self):
          '''
          Function to initialze the parameters
          '''
          for m in self.modules():
              if isinstance(m, nn.Conv2d):
                  nn.init.kaiming_normal_(m.weight, mode='fan_out')
                  if m.bias is not None:
                      nn.init.constant_(m.bias, 0)
              elif isinstance(m, nn.BatchNorm2d):
                  nn.init.constant_(m.weight, 1)
                  nn.init.constant_(m.bias, 0)
              elif isinstance(m, nn.Linear):
                  nn.init.normal_(m.weight, std=0.001)
                  if m.bias is not None:
                      nn.init.constant_(m.bias, 0)
  
          # init of deconv layers
          for name, m in self.deconv_layers.named_modules():
              if isinstance(m, nn.BatchNorm2d):
                  nn.init.constant_(m.weight, 1)
                  nn.init.constant_(m.bias, 0)
  
      def _get_deconv_cfg(self, deconv_kernel, index):
          if deconv_kernel == 4:
              padding = 1
              output_padding = 0
          elif deconv_kernel == 3:
              padding = 1
              output_padding = 1
          elif deconv_kernel == 2:
              padding = 0
              output_padding = 0
  
          return deconv_kernel, padding, output_padding
  
      def _make_deconv_layer(self, num_layers, num_filters, num_kernels):
          assert num_layers == len(num_filters), \
              'ERROR: num_deconv_layers is different len(num_deconv_filters)'
          assert num_layers == len(num_kernels), \
              'ERROR: num_deconv_layers is different len(num_deconv_filters)'
  
          layers = []
          for i in range(num_layers):
              kernel, padding, output_padding = \
                  self._get_deconv_cfg(num_kernels[i], i)
  
              planes = num_filters[i]
              fc = DCN(self.inplanes, planes, 
                      kernel_size=(3,3), stride=1,
                      padding=1, dilation=1, deformable_groups=1)
              # fc = nn.Conv2d(self.inplanes, planes,
              #         kernel_size=3, stride=1, 
              #         padding=1, dilation=1, bias=False)
              # fill_fc_weights(fc)
              up = nn.ConvTranspose2d(
                      in_channels=planes,
                      out_channels=planes,
                      kernel_size=kernel,
                      stride=2,
                      padding=padding,
                      output_padding=output_padding,
                      bias=self.deconv_with_bias)
              fill_up_weights(up)
  
              layers.append(fc)
              layers.append(nn.BatchNorm2d(planes, momentum=BN_MOMENTUM))
              layers.append(nn.ReLU(inplace=True))
              layers.append(up)
              layers.append(nn.BatchNorm2d(planes, momentum=BN_MOMENTUM))
              layers.append(nn.ReLU(inplace=True))
              self.inplanes = planes
  
          return nn.Sequential(*layers)
  
      def forward(self, input, p=0.2):
          '''
          :param input: Receives the input RGB image
          '''
          out_l1 = self.level1(input)  # 112
          if not self.input_reinforcement:
              del input
              input = None
  
          out_l2 = self.level2_0(out_l1, input)  # 56
  
          out_l3_0 = self.level3_0(out_l2, input)  # down-sample
          for i, layer in enumerate(self.level3):
              if i == 0:
                  out_l3 = layer(out_l3_0)
              else:
                  out_l3 = layer(out_l3)
  
          out_l4_0 = self.level4_0(out_l3, input)  # down-sample
          for i, layer in enumerate(self.level4):
              if i == 0:
                  out_l4 = layer(out_l4_0)
              else:
                  out_l4 = layer(out_l4)
  
          out_l5_0 = self.level5_0(out_l4)  # down-sample
          for i, layer in enumerate(self.level5):
              if i == 0:
                  out_l5 = layer(out_l5_0)
              else:
                  out_l5 = layer(out_l5)
  
          # output_g = F.adaptive_avg_pool2d(out_l5, output_size=1)
          # output_g = F.dropout(output_g, p=p, training=self.training)
          # output_1x1 = output_g.view(output_g.size(0), -1)
  
          # deconv layers here 
          x = self.deconv_layers(out_l5)
          # x = out_l5
  
          ret = {} 
          for head in self.heads:
              ret[head] = self.__getattr__(head)(x)
          return [ret]
  
  def get_espv2_net(num_layers, heads, head_conv=1024):
  
      # input = torch.Tensor(1, 3, 224, 224).cuda()
      # model = EESPNet(classes=1000, s=1.0)
  
      # head convolution size 
      model = EESPNet(heads,
                   pretrained=False,
                   head_conv=head_conv,
                   s=1.0)
      
      return model