text: # MIT License
  # 
  # Copyright (c) 2018 akanazawa
  # Copyright (c) 2021 Google LLC
  # 
  # Permission is hereby granted, free of charge, to any person obtaining a copy
  # of this software and associated documentation files (the "Software"), to deal
  # in the Software without restriction, including without limitation the rights
  # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  # copies of the Software, and to permit persons to whom the Software is
  # furnished to do so, subject to the following conditions:
  # 
  # The above copyright notice and this permission notice shall be included in all
  # copies or substantial portions of the Software.
  # 
  # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  # SOFTWARE.
  
  '''
  CNN building blocks.
  Taken from https://github.com/shubhtuls/factored3d/
  '''
  from __future__ import division
  from __future__ import print_function
  import torch
  import torchvision
  import torch.nn as nn
  import math
  import kornia
  
  class Flatten(nn.Module):
      def forward(self, x):
          return x.view(x.size()[0], -1)
  
  class Unsqueeze(nn.Module):
      def __init__(self, dim):
          super(Unsqueeze, self).__init__()
          self.dim = dim
  
      def forward(self, x):
          return x.unsqueeze(self.dim)
  
  ## fc layers
  def fc(batch_norm, nc_inp, nc_out):
      if batch_norm:
          return nn.Sequential(
              nn.Linear(nc_inp, nc_out, bias=True),
              nn.BatchNorm1d(nc_out),
              nn.LeakyReLU(0.2,inplace=True)
          )
      else:
          return nn.Sequential(
              nn.Linear(nc_inp, nc_out),
              nn.LeakyReLU(0.1,inplace=True)
          )
  
  def fc_stack(nc_inp, nc_out, nlayers, use_bn=True):
      modules = []
      for l in range(nlayers):
          modules.append(fc(use_bn, nc_inp, nc_out))
          nc_inp = nc_out
      encoder = nn.Sequential(*modules)
      net_init(encoder)
      return encoder
  
  ## 2D convolution layers
  def conv2d(batch_norm, in_planes, out_planes, kernel_size=3, stride=1):
      if batch_norm:
          return nn.Sequential(
              nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, bias=True),
              nn.BatchNorm2d(out_planes),
              nn.LeakyReLU(0.2,inplace=True)
          )
      else:
          return nn.Sequential(
              nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, bias=True),
              nn.LeakyReLU(0.2,inplace=True)
          )
  
  
  def deconv2d(in_planes, out_planes):
      return nn.Sequential(
          nn.ConvTranspose2d(in_planes, out_planes, kernel_size=4, stride=2, padding=1, bias=True),
          nn.LeakyReLU(0.2,inplace=True)
      )
  
  
  def upconv2d(in_planes, out_planes, mode='bilinear'):
      if mode == 'nearest':
          print('Using NN upsample!!')
      upconv = nn.Sequential(
          nn.Upsample(scale_factor=2, mode=mode, align_corners=True),
          nn.ReflectionPad2d(1),
          nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=1, padding=0),
          nn.LeakyReLU(0.2,inplace=True)
      )
      return upconv
  
  
  def decoder2d(nlayers, nz_shape, nc_input, use_bn=True, nc_final=1, nc_min=8, nc_step=1, init_fc=True, use_deconv=False, upconv_mode='bilinear'):
      ''' Simple 3D encoder with nlayers.
      
      Args:
          nlayers: number of decoder layers
          nz_shape: number of bottleneck
          nc_input: number of channels to start upconvolution from
          use_bn: whether to use batch_norm
          nc_final: number of output channels
          nc_min: number of min channels
          nc_step: double number of channels every nc_step layers
          init_fc: initial features are not spatial, use an fc & unsqueezing to make them 3D
      '''
      modules = []
      if init_fc:
          modules.append(fc(use_bn, nz_shape, nc_input))
          for d in range(3):
              modules.append(Unsqueeze(2))
      nc_output = nc_input
      for nl in range(nlayers):
          if (nl % nc_step==0) and (nc_output//2 >= nc_min):
              nc_output = nc_output//2
          if use_deconv:
              print('Using deconv decoder!')
              modules.append(deconv2d(nc_input, nc_output))
              nc_input = nc_output
              modules.append(conv2d(use_bn, nc_input, nc_output))
          else:
              modules.append(upconv2d(nc_input, nc_output, mode=upconv_mode))
              nc_input = nc_output
              modules.append(conv2d(use_bn, nc_input, nc_output))
  
      modules.append(nn.Conv2d(nc_output, nc_final, kernel_size=3, stride=1, padding=1, bias=True))
      decoder = nn.Sequential(*modules)
      net_init(decoder)
      return decoder
  
  
  ## 3D convolution layers
  def conv3d(batch_norm, in_planes, out_planes, kernel_size=3, stride=1):
      if batch_norm:
          return nn.Sequential(
              nn.Conv3d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, bias=True),
              nn.BatchNorm3d(out_planes),
              nn.LeakyReLU(0.2,inplace=True)
          )
      else:
          return nn.Sequential(
              nn.Conv3d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, bias=True),
              nn.LeakyReLU(0.2,inplace=True)
          )
  
  
  def deconv3d(batch_norm, in_planes, out_planes):
      if batch_norm:
          return nn.Sequential(
              nn.ConvTranspose3d(in_planes, out_planes, kernel_size=4, stride=2, padding=1, bias=True),
              nn.BatchNorm3d(out_planes),
              nn.LeakyReLU(0.2,inplace=True)
          )
      else:        
          return nn.Sequential(
              nn.ConvTranspose3d(in_planes, out_planes, kernel_size=4, stride=2, padding=1, bias=True),
              nn.LeakyReLU(0.2,inplace=True)
          )
  
  
  ## 3D Network Modules
  def encoder3d(nlayers, use_bn=True, nc_input=1, nc_max=128, nc_l1=8, nc_step=1, nz_shape=20):
      ''' Simple 3D encoder with nlayers.
      
      Args:
          nlayers: number of encoder layers
          use_bn: whether to use batch_norm
          nc_input: number of input channels
          nc_max: number of max channels
          nc_l1: number of channels in layer 1
          nc_step: double number of channels every nc_step layers      
          nz_shape: size of bottleneck layer
      '''
      modules = []
      nc_output = nc_l1
      for nl in range(nlayers):
          if (nl>=1) and (nl%nc_step==0) and (nc_output <= nc_max*2):
              nc_output *= 2
  
          modules.append(conv3d(use_bn, nc_input, nc_output, stride=1))
          nc_input = nc_output
          modules.append(conv3d(use_bn, nc_input, nc_output, stride=1))
          modules.append(torch.nn.MaxPool3d(kernel_size=2, stride=2))
  
      modules.append(Flatten())
      modules.append(fc_stack(nc_output, nz_shape, 2, use_bn=True))
      encoder = nn.Sequential(*modules)
      net_init(encoder)
      return encoder, nc_output
  
  
  def decoder3d(nlayers, nz_shape, nc_input, use_bn=True, nc_final=1, nc_min=8, nc_step=1, init_fc=True):
      ''' Simple 3D encoder with nlayers.
      
      Args:
          nlayers: number of decoder layers
          nz_shape: number of bottleneck
          nc_input: number of channels to start upconvolution from
          use_bn: whether to use batch_norm
          nc_final: number of output channels
          nc_min: number of min channels
          nc_step: double number of channels every nc_step layers
          init_fc: initial features are not spatial, use an fc & unsqueezing to make them 3D
      '''
      modules = []
      if init_fc:
          modules.append(fc(use_bn, nz_shape, nc_input))
          for d in range(3):
              modules.append(Unsqueeze(2))
      nc_output = nc_input
      for nl in range(nlayers):
          if (nl%nc_step==0) and (nc_output//2 >= nc_min):
              nc_output = nc_output//2
  
          modules.append(deconv3d(use_bn, nc_input, nc_output))
          nc_input = nc_output
          modules.append(conv3d(use_bn, nc_input, nc_output))
  
      modules.append(nn.Conv3d(nc_output, nc_final, kernel_size=3, stride=1, padding=1, bias=True))
      decoder = nn.Sequential(*modules)
      net_init(decoder)
      return decoder
  
  
  def net_init(net):
      for m in net.modules():
          if isinstance(m, nn.Linear):
              #n = m.out_features
              #m.weight.data.normal_(0, 0.02 / n) #this modified initialization seems to work better, but it's very hacky
              #n = m.in_features
              #m.weight.data.normal_(0, math.sqrt(2. / n)) #xavier
              m.weight.data.normal_(0, 0.02)
              if m.bias is not None:
                  m.bias.data.zero_()
  
          if isinstance(m, nn.Conv2d): #or isinstance(m, nn.ConvTranspose2d):
              #n = m.kernel_size[0] * m.kernel_size[1] * m.in_channels
              #m.weight.data.normal_(0, math.sqrt(2. / n)) #this modified initialization seems to work better, but it's very hacky
              m.weight.data.normal_(0, 0.02)
              if m.bias is not None:
                  m.bias.data.zero_()
  
          if isinstance(m, nn.ConvTranspose2d):
              # Initialize Deconv with bilinear weights.
              base_weights = bilinear_init(m.weight.data.size(-1))
              base_weights = base_weights.unsqueeze(0).unsqueeze(0)
              m.weight.data = base_weights.repeat(m.weight.data.size(0), m.weight.data.size(1), 1, 1)
              if m.bias is not None:
                  m.bias.data.zero_()
  
          if isinstance(m, nn.Conv3d) or isinstance(m, nn.ConvTranspose3d):
              #n = m.kernel_size[0] * m.kernel_size[1] * m.kernel_size[2] * m.in_channels
              #m.weight.data.normal_(0, math.sqrt(2. / n))
              m.weight.data.normal_(0, 0.02)
              if m.bias is not None:
                  m.bias.data.zero_()
  
          elif isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.BatchNorm3d):
              m.weight.data.fill_(1)
              m.bias.data.zero_()
  
  
  def bilinear_init(kernel_size=4):
      # Following Caffe's BilinearUpsamplingFiller
      # https://github.com/BVLC/caffe/pull/2213/files
      import numpy as np
      width = kernel_size
      height = kernel_size
      f = int(np.ceil(width / 2.))
      cc = (2 * f - 1 - f % 2) / (2.*f)
      weights = torch.zeros((height, width))
      for y in range(height):
          for x in range(width):
              weights[y, x] = (1 - np.abs(x / f - cc)) * (1 - np.abs(y / f - cc))
  
      return weights
  
  #------------- Modules ------------#
  #----------------------------------#
  class ResNetConv(nn.Module):
      def __init__(self, n_blocks=4):
          super(ResNetConv, self).__init__()
          self.resnet = torchvision.models.resnet18(pretrained=True)
          self.resnet.fc=None
          self.n_blocks = n_blocks
  
      def forward(self, x):
          n_blocks = self.n_blocks
          x = self.resnet.conv1(x)
          x = self.resnet.bn1(x)
          x = self.resnet.relu(x)
          x = self.resnet.maxpool(x)
  
          if n_blocks >= 1:
              x = self.resnet.layer1(x)
          if n_blocks >= 2:
              x = self.resnet.layer2(x)
          if n_blocks >= 3:
              x = self.resnet.layer3(x)
          if n_blocks >= 4:
              x = self.resnet.layer4(x)
          return x
  
  
  class Encoder(nn.Module):
      """
      Current:
      Resnet with 4 blocks (x32 spatial dim reduction)
      Another conv with stride 2 (x64)
      This is sent to 2 fc layers with final output nz_feat.
      """
  
      def __init__(self, input_shape, n_blocks=4, nz_feat=100, batch_norm=True):
          super(Encoder, self).__init__()
          self.resnet_conv = ResNetConv(n_blocks=4)
          self.enc_conv1 = conv2d(batch_norm, 512, 256, stride=2, kernel_size=4)
          nc_input = 256 * (input_shape[0] // 64) * (input_shape[1] // 64)
          self.enc_fc = fc_stack(nc_input, nz_feat, 2)
  
          net_init(self.enc_conv1)
  
      def forward(self, img):
          resnet_feat = self.resnet_conv.forward(img)
  
          out_enc_conv1 = self.enc_conv1(resnet_feat)
          out_enc_conv1 = out_enc_conv1.view(img.size(0), -1)
          feat = self.enc_fc.forward(out_enc_conv1)
          return feat
  
  class QuatPredictor(nn.Module):
      def __init__(self, nz_feat, nz_rot=4, classify_rot=False,n_bones=None,n_hypo=None):
          super(QuatPredictor, self).__init__()
          self.pred_layer = nn.Linear(nz_feat, nz_rot*n_bones*n_hypo)
          self.classify_rot = classify_rot
          self.nmesh = n_bones
          self.nhypo = n_hypo
          self.nz_feat = nz_feat
  
      def forward(self, feat):
          quat = self.pred_layer.forward(feat).view(-1,4)
          quat = quat.view(-1,self.nhypo,self.nmesh,4)
          quat[:,:,1:,3]+=10
          quat = quat.view(-1,4)
          if self.classify_rot:
              quat = torch.nn.functional.log_softmax(quat)
          else:
              quat = torch.nn.functional.normalize(quat)
          return kornia.quaternion_to_rotation_matrix(quat).view(-1,9)
  
      def reinit(self, n_hypo, good_hypo):
          prev_wt = self.pred_layer.weight.clone().view(n_hypo, -1, self.nz_feat)
          prev_bi = self.pred_layer.bias.clone().view(n_hypo,-1)
          
          prev_wt[~good_hypo] = prev_wt[good_hypo]
          prev_bi[~good_hypo] = prev_bi[good_hypo]
  
          self.pred_layer.weight.data = prev_wt.view(-1,self.nz_feat)
          self.pred_layer.bias.data   = prev_bi.view(-1,)
  
  
  class DepthPredictor(nn.Module):
      def __init__(self, nz,n_bones=None,offset=10):
          super(DepthPredictor, self).__init__()
          self.pred_layer = nn.Linear(nz, 1*n_bones)
          self.nz_feat = nz
          self.offset=offset
  
      def forward(self, feat):
          depth = self.pred_layer.forward(feat) + self.offset # mean is 10+1e-12, min is 1e-12
          depth = torch.nn.functional.relu(depth) + 1e-12
          return depth
      
      def reinit(self, n_hypo, good_hypo):
          prev_wt = self.pred_layer.weight.clone().view(n_hypo, -1, self.nz_feat)
          prev_bi = self.pred_layer.bias.clone().view(n_hypo,-1)
          
          prev_wt[~good_hypo] = prev_wt[good_hypo]
          prev_bi[~good_hypo] = prev_bi[good_hypo]
  
          self.pred_layer.weight.data = prev_wt.view(-1,self.nz_feat)
          self.pred_layer.bias.data   = prev_bi.view(-1,)
  
  class TransPredictor(nn.Module):
      """
      Outputs [tx, ty] or [tx, ty, tz]
      """
  
      def __init__(self, nz, orth=True,n_bones=None):
          super(TransPredictor, self).__init__()
          if orth:
              self.pred_layer = nn.Linear(nz, 2*n_bones)
          else:
              self.pred_layer = nn.Linear(nz, 3)
  
      def forward(self, feat):
          trans = self.pred_layer.forward(feat).view(-1,2)
          #trans = self.pred_layer.forward(feat)
          # print('trans: ( Mean = {}, Var = {} )'.format(trans.mean().data[0], trans.var().data[0]))
          return trans
  
  
  class PPointPredictor(nn.Module):
      def __init__(self, nz, orth=True):
          super(PPointPredictor, self).__init__()
          self.pred_layer = nn.Linear(nz, 2)
  
      def forward(self, feat):
          trans = self.pred_layer.forward(feat)
          # print('trans: ( Mean = {}, Var = {} )'.format(trans.mean().data[0], trans.var().data[0]))
          return trans
  
  
  class CodePredictor(nn.Module):
      def __init__(self, nz_feat=100, num_verts=1000,n_bones = None, n_hypo=None):
          super(CodePredictor, self).__init__()
          self.offset = 20
          torch.manual_seed(0)
          self.quat_predictor = QuatPredictor(nz_feat, n_bones=n_bones, n_hypo=n_hypo)
          self.scale_predictor = DepthPredictor(nz_feat,n_bones=n_hypo,offset=self.offset)
          self.trans_predictor = TransPredictor(nz_feat,n_bones=n_bones)
          self.depth_predictor = DepthPredictor(nz_feat,n_bones=n_bones,offset=self.offset)
          self.ppoint_predictor = PPointPredictor(nz_feat)
  
          self.nmesh = n_bones
          self.nhypo = n_hypo
  
      def forward(self, feat):
          scale_pred = self.scale_predictor.forward(feat)
          quat_pred = self.quat_predictor.forward(feat)
          
          trans_pred = self.trans_predictor.forward(feat)/10.
  
          depth_pred = self.depth_predictor.forward(feat)
          depth_pred = depth_pred.view(-1,1,self.nmesh)
          depth_pred[:,:,1:] =  (depth_pred[:,:,1:]-self.offset)/10.
          depth_pred = depth_pred.view(feat.shape[0],-1)
          
          ppoint_pred = self.ppoint_predictor.forward(feat)/10.
          return scale_pred, trans_pred, quat_pred, depth_pred, ppoint_pred
  
  
  
  if __name__ == '__main__':
      decoder2d(5, None, 256, use_deconv=True, init_fc=False)
      bilinear_init()
  