text: # Copyright (c) Chris Choy (chrischoy@ai.stanford.edu) and Wei Dong (weidong@andrew.cmu.edu)
  #
  # Please cite the following papers if you use any part of the code.
  # - Christopher Choy, Wei Dong, Vladlen Koltun, Deep Global Registration, CVPR 2020
  # - Christopher Choy, Jaesik Park, Vladlen Koltun, Fully Convolutional Geometric Features, ICCV 2019
  # - Christopher Choy, JunYoung Gwak, Silvio Savarese, 4D Spatio-Temporal ConvNets: Minkowski Convolutional Neural Networks, CVPR 2019
  import torch
  import torch.nn as nn
  import MinkowskiEngine as ME
  from model.common import get_norm, get_nonlinearity
  
  from model.residual_block import get_block, conv, conv_tr, conv_norm_non
  
  
  class PyramidModule(ME.MinkowskiNetwork):
    NONLINEARITY = 'ELU'
    NORM_TYPE = 'BN'
    REGION_TYPE = ME.RegionType.HYPER_CUBE
  
    def __init__(self,
                 inc,
                 outc,
                 inner_inc,
                 inner_outc,
                 inner_module=None,
                 depth=1,
                 bn_momentum=0.05,
                 dimension=-1):
      ME.MinkowskiNetwork.__init__(self, dimension)
      self.depth = depth
  
      self.conv = nn.Sequential(
          conv_norm_non(
              inc,
              inner_inc,
              3,
              2,
              dimension,
              region_type=self.REGION_TYPE,
              norm_type=self.NORM_TYPE,
              nonlinearity=self.NONLINEARITY), *[
                  get_block(
                      self.NORM_TYPE,
                      inner_inc,
                      inner_inc,
                      bn_momentum=bn_momentum,
                      region_type=self.REGION_TYPE,
                      dimension=dimension) for d in range(depth)
              ])
      self.inner_module = inner_module
      self.convtr = nn.Sequential(
          conv_tr(
              in_channels=inner_outc,
              out_channels=inner_outc,
              kernel_size=3,
              stride=2,
              dilation=1,
              bias=False,
              region_type=self.REGION_TYPE,
              dimension=dimension),
          get_norm(
              self.NORM_TYPE, inner_outc, bn_momentum=bn_momentum, dimension=dimension),
          get_nonlinearity(self.NONLINEARITY))
  
      self.cat_conv = conv_norm_non(
          inner_outc + inc,
          outc,
          1,
          1,
          dimension,
          norm_type=self.NORM_TYPE,
          nonlinearity=self.NONLINEARITY)
  
    def forward(self, x):
      y = self.conv(x)
      if self.inner_module:
        y = self.inner_module(y)
      y = self.convtr(y)
      y = ME.cat(x, y)
      return self.cat_conv(y)
  
  
  class PyramidModuleINBN(PyramidModule):
    NORM_TYPE = 'INBN'
  
  
  class PyramidNet(ME.MinkowskiNetwork):
    NORM_TYPE = 'BN'
    NONLINEARITY = 'ELU'
    PYRAMID_MODULE = PyramidModule
    CHANNELS = [32, 64, 128, 128]
    TR_CHANNELS = [64, 128, 128, 128]
    DEPTHS = [1, 1, 1, 1]
    # None        b1, b2, b3, btr3, btr2
    #               1  2  3 -3 -2 -1
    REGION_TYPE = ME.RegionType.HYPER_CUBE
  
    # To use the model, must call initialize_coords before forward pass.
    # Once data is processed, call clear to reset the model before calling initialize_coords
    def __init__(self,
                 in_channels=3,
                 out_channels=32,
                 bn_momentum=0.1,
                 conv1_kernel_size=3,
                 normalize_feature=False,
                 D=3):
      ME.MinkowskiNetwork.__init__(self, D)
      self.conv1_kernel_size = conv1_kernel_size
      self.normalize_feature = normalize_feature
  
      self.initialize_network(in_channels, out_channels, bn_momentum, D)
  
    def initialize_network(self, in_channels, out_channels, bn_momentum, dimension):
      NORM_TYPE = self.NORM_TYPE
      NONLINEARITY = self.NONLINEARITY
      CHANNELS = self.CHANNELS
      TR_CHANNELS = self.TR_CHANNELS
      DEPTHS = self.DEPTHS
      REGION_TYPE = self.REGION_TYPE
  
      self.conv = conv_norm_non(
          in_channels,
          CHANNELS[0],
          kernel_size=self.conv1_kernel_size,
          stride=1,
          dimension=dimension,
          bn_momentum=bn_momentum,
          region_type=REGION_TYPE,
          norm_type=NORM_TYPE,
          nonlinearity=NONLINEARITY)
  
      pyramid = None
      for d in range(len(DEPTHS) - 1, 0, -1):
        pyramid = self.PYRAMID_MODULE(
            CHANNELS[d - 1],
            TR_CHANNELS[d - 1],
            CHANNELS[d],
            TR_CHANNELS[d],
            pyramid,
            DEPTHS[d],
            dimension=dimension)
      self.pyramid = pyramid
      self.final = nn.Sequential(
          conv_norm_non(
              TR_CHANNELS[0],
              TR_CHANNELS[0],
              kernel_size=3,
              stride=1,
              dimension=dimension),
          conv(TR_CHANNELS[0], out_channels, 1, 1, dimension=dimension))
  
    def forward(self, x):
      out = self.conv(x)
      out = self.pyramid(out)
      out = self.final(out)
  
      if self.normalize_feature:
        return ME.SparseTensor(
            out.F / (torch.norm(out.F, p=2, dim=1, keepdim=True) + 1e-8),
            coordinate_map_key=out.coordinate_map_key,
            coordinate_manager=out.coordinate_manager)
      else:
        return out
  
  
  class PyramidNet6(PyramidNet):
    CHANNELS = [32, 64, 128, 192, 256, 256]
    TR_CHANNELS = [64, 128, 192, 192, 256, 256]
    DEPTHS = [1, 1, 1, 1, 1, 1]
  
  
  class PyramidNet6NoBlock(PyramidNet6):
    DEPTHS = [0, 0, 0, 0, 0, 0]
  
  
  class PyramidNet6INBN(PyramidNet6):
    NORM_TYPE = 'INBN'
    PYRAMID_MODULE = PyramidModuleINBN
  
  
  class PyramidNet6INBNNoBlock(PyramidNet6INBN):
    NORM_TYPE = 'INBN'
  
  
  class PyramidNet8(PyramidNet):
    CHANNELS = [32, 64, 128, 128, 192, 192, 256, 256]
    TR_CHANNELS = [64, 128, 128, 192, 192, 192, 256, 256]
    DEPTHS = [1, 1, 1, 1, 1, 1, 1, 1]
  
  
  class PyramidNet8INBN(PyramidNet8):
    NORM_TYPE = 'INBN'
    PYRAMID_MODULE = PyramidModuleINBN
  