content: import torch
  import torch.nn as nn
  import torchsparse
  import torchsparse.nn as spnn
  from torchsparse.tensor import PointTensor
  
  from tsparse.torchsparse_utils import *
  
  
  # __all__ = ['SPVCNN', 'SConv3d', 'SparseConvGRU']
  
  
  class ConvBnReLU(nn.Module):
      def __init__(self, in_channels, out_channels,
                   kernel_size=3, stride=1, pad=1):
          super(ConvBnReLU, self).__init__()
          self.conv = nn.Conv2d(in_channels, out_channels,
                                kernel_size, stride=stride, padding=pad, bias=False)
          self.bn = nn.BatchNorm2d(out_channels)
          self.activation = nn.ReLU(inplace=True)
  
      def forward(self, x):
          return self.activation(self.bn(self.conv(x)))
  
  
  class ConvBnReLU3D(nn.Module):
      def __init__(self, in_channels, out_channels,
                   kernel_size=3, stride=1, pad=1):
          super(ConvBnReLU3D, self).__init__()
          self.conv = nn.Conv3d(in_channels, out_channels,
                                kernel_size, stride=stride, padding=pad, bias=False)
          self.bn = nn.BatchNorm3d(out_channels)
          self.activation = nn.ReLU(inplace=True)
  
      def forward(self, x):
          return self.activation(self.bn(self.conv(x)))
  
  
  ###################################  feature net  ######################################
  class FeatureNet(nn.Module):
      """
      output 3 levels of features using a FPN structure
      """
  
      def __init__(self):
          super(FeatureNet, self).__init__()
  
          self.conv0 = nn.Sequential(
              ConvBnReLU(3, 8, 3, 1, 1),
              ConvBnReLU(8, 8, 3, 1, 1))
  
          self.conv1 = nn.Sequential(
              ConvBnReLU(8, 16, 5, 2, 2),
              ConvBnReLU(16, 16, 3, 1, 1),
              ConvBnReLU(16, 16, 3, 1, 1))
  
          self.conv2 = nn.Sequential(
              ConvBnReLU(16, 32, 5, 2, 2),
              ConvBnReLU(32, 32, 3, 1, 1),
              ConvBnReLU(32, 32, 3, 1, 1))
  
          self.toplayer = nn.Conv2d(32, 32, 1)
          self.lat1 = nn.Conv2d(16, 32, 1)
          self.lat0 = nn.Conv2d(8, 32, 1)
  
          # to reduce channel size of the outputs from FPN
          self.smooth1 = nn.Conv2d(32, 16, 3, padding=1)
          self.smooth0 = nn.Conv2d(32, 8, 3, padding=1)
  
      def _upsample_add(self, x, y):
          return torch.nn.functional.interpolate(x, scale_factor=2,
                                                 mode="bilinear", align_corners=True) + y
  
      def forward(self, x):
          # x: (B, 3, H, W)
          conv0 = self.conv0(x)  # (B, 8, H, W)
          conv1 = self.conv1(conv0)  # (B, 16, H//2, W//2)
          conv2 = self.conv2(conv1)  # (B, 32, H//4, W//4)
          feat2 = self.toplayer(conv2)  # (B, 32, H//4, W//4)
          feat1 = self._upsample_add(feat2, self.lat1(conv1))  # (B, 32, H//2, W//2)
          feat0 = self._upsample_add(feat1, self.lat0(conv0))  # (B, 32, H, W)
  
          # reduce output channels
          feat1 = self.smooth1(feat1)  # (B, 16, H//2, W//2)
          feat0 = self.smooth0(feat0)  # (B, 8, H, W)
  
          # feats = {"level_0": feat0,
          #          "level_1": feat1,
          #          "level_2": feat2}
  
          return [feat2, feat1, feat0]  # coarser to finer features
  
  
  class BasicSparseConvolutionBlock(nn.Module):
      def __init__(self, inc, outc, ks=3, stride=1, dilation=1):
          super().__init__()
          self.net = nn.Sequential(
              spnn.Conv3d(inc,
                          outc,
                          kernel_size=ks,
                          dilation=dilation,
                          stride=stride),
              spnn.BatchNorm(outc),
              spnn.ReLU(True))
  
      def forward(self, x):
          out = self.net(x)
          return out
  
  
  class BasicSparseDeconvolutionBlock(nn.Module):
      def __init__(self, inc, outc, ks=3, stride=1):
          super().__init__()
          self.net = nn.Sequential(
              spnn.Conv3d(inc,
                          outc,
                          kernel_size=ks,
                          stride=stride,
                          transposed=True),
              spnn.BatchNorm(outc),
              spnn.ReLU(True))
  
      def forward(self, x):
          return self.net(x)
  
  
  class SparseResidualBlock(nn.Module):
      def __init__(self, inc, outc, ks=3, stride=1, dilation=1):
          super().__init__()
          self.net = nn.Sequential(
              spnn.Conv3d(inc,
                          outc,
                          kernel_size=ks,
                          dilation=dilation,
                          stride=stride), spnn.BatchNorm(outc),
              spnn.ReLU(True),
              spnn.Conv3d(outc,
                          outc,
                          kernel_size=ks,
                          dilation=dilation,
                          stride=1), spnn.BatchNorm(outc))
  
          self.downsample = nn.Sequential() if (inc == outc and stride == 1) else \
              nn.Sequential(
                  spnn.Conv3d(inc, outc, kernel_size=1, dilation=1, stride=stride),
                  spnn.BatchNorm(outc)
              )
  
          self.relu = spnn.ReLU(True)
  
      def forward(self, x):
          out = self.relu(self.net(x) + self.downsample(x))
          return out
  
  
  class SPVCNN(nn.Module):
      def __init__(self, **kwargs):
          super().__init__()
  
          self.dropout = kwargs['dropout']
  
          cr = kwargs.get('cr', 1.0)
          cs = [32, 64, 128, 96, 96]
          cs = [int(cr * x) for x in cs]
  
          if 'pres' in kwargs and 'vres' in kwargs:
              self.pres = kwargs['pres']
              self.vres = kwargs['vres']
  
          self.stem = nn.Sequential(
              spnn.Conv3d(kwargs['in_channels'], cs[0], kernel_size=3, stride=1),
              spnn.BatchNorm(cs[0]), spnn.ReLU(True)
          )
  
          self.stage1 = nn.Sequential(
              BasicSparseConvolutionBlock(cs[0], cs[0], ks=2, stride=2, dilation=1),
              SparseResidualBlock(cs[0], cs[1], ks=3, stride=1, dilation=1),
              SparseResidualBlock(cs[1], cs[1], ks=3, stride=1, dilation=1),
          )
  
          self.stage2 = nn.Sequential(
              BasicSparseConvolutionBlock(cs[1], cs[1], ks=2, stride=2, dilation=1),
              SparseResidualBlock(cs[1], cs[2], ks=3, stride=1, dilation=1),
              SparseResidualBlock(cs[2], cs[2], ks=3, stride=1, dilation=1),
          )
  
          self.up1 = nn.ModuleList([
              BasicSparseDeconvolutionBlock(cs[2], cs[3], ks=2, stride=2),
              nn.Sequential(
                  SparseResidualBlock(cs[3] + cs[1], cs[3], ks=3, stride=1,
                                      dilation=1),
                  SparseResidualBlock(cs[3], cs[3], ks=3, stride=1, dilation=1),
              )
          ])
  
          self.up2 = nn.ModuleList([
              BasicSparseDeconvolutionBlock(cs[3], cs[4], ks=2, stride=2),
              nn.Sequential(
                  SparseResidualBlock(cs[4] + cs[0], cs[4], ks=3, stride=1,
                                      dilation=1),
                  SparseResidualBlock(cs[4], cs[4], ks=3, stride=1, dilation=1),
              )
          ])
  
          self.point_transforms = nn.ModuleList([
              nn.Sequential(
                  nn.Linear(cs[0], cs[2]),
                  nn.BatchNorm1d(cs[2]),
                  nn.ReLU(True),
              ),
              nn.Sequential(
                  nn.Linear(cs[2], cs[4]),
                  nn.BatchNorm1d(cs[4]),
                  nn.ReLU(True),
              )
          ])
  
          self.weight_initialization()
  
          if self.dropout:
              self.dropout = nn.Dropout(0.3, True)
  
      def weight_initialization(self):
          for m in self.modules():
              if isinstance(m, nn.BatchNorm1d):
                  nn.init.constant_(m.weight, 1)
                  nn.init.constant_(m.bias, 0)
  
      def forward(self, z):
          # x: SparseTensor z: PointTensor
          x0 = initial_voxelize(z, self.pres, self.vres)
  
          x0 = self.stem(x0)
          z0 = voxel_to_point(x0, z, nearest=False)
          z0.F = z0.F
  
          x1 = point_to_voxel(x0, z0)
          x1 = self.stage1(x1)
          x2 = self.stage2(x1)
          z1 = voxel_to_point(x2, z0)
          z1.F = z1.F + self.point_transforms[0](z0.F)
  
          y3 = point_to_voxel(x2, z1)
          if self.dropout:
              y3.F = self.dropout(y3.F)
          y3 = self.up1[0](y3)
          y3 = torchsparse.cat([y3, x1])
          y3 = self.up1[1](y3)
  
          y4 = self.up2[0](y3)
          y4 = torchsparse.cat([y4, x0])
          y4 = self.up2[1](y4)
          z3 = voxel_to_point(y4, z1)
          z3.F = z3.F + self.point_transforms[1](z1.F)
  
          return z3.F
  
  
  class SparseCostRegNet(nn.Module):
      """
      Sparse cost regularization network;
      require sparse tensors as input
      """
  
      def __init__(self, d_in, d_out=8):
          super(SparseCostRegNet, self).__init__()
          self.d_in = d_in
          self.d_out = d_out
  
          self.conv0 = BasicSparseConvolutionBlock(d_in, d_out)
  
          self.conv1 = BasicSparseConvolutionBlock(d_out, 16, stride=2)
          self.conv2 = BasicSparseConvolutionBlock(16, 16)
  
          self.conv3 = BasicSparseConvolutionBlock(16, 32, stride=2)
          self.conv4 = BasicSparseConvolutionBlock(32, 32)
  
          self.conv5 = BasicSparseConvolutionBlock(32, 64, stride=2)
          self.conv6 = BasicSparseConvolutionBlock(64, 64)
  
          self.conv7 = BasicSparseDeconvolutionBlock(64, 32, ks=3, stride=2)
  
          self.conv9 = BasicSparseDeconvolutionBlock(32, 16, ks=3, stride=2)
  
          self.conv11 = BasicSparseDeconvolutionBlock(16, d_out, ks=3, stride=2)
  
      def forward(self, x):
          """
  
          :param x: sparse tensor
          :return: sparse tensor
          """
          conv0 = self.conv0(x)
          conv2 = self.conv2(self.conv1(conv0))
          conv4 = self.conv4(self.conv3(conv2))
  
          x = self.conv6(self.conv5(conv4))
          x = conv4 + self.conv7(x)
          del conv4
          x = conv2 + self.conv9(x)
          del conv2
          x = conv0 + self.conv11(x)
          del conv0
          return x.F
  
  
  class SConv3d(nn.Module):
      def __init__(self, inc, outc, pres, vres, ks=3, stride=1, dilation=1):
          super().__init__()
          self.net = spnn.Conv3d(inc,
                                 outc,
                                 kernel_size=ks,
                                 dilation=dilation,
                                 stride=stride)
          self.point_transforms = nn.Sequential(
              nn.Linear(inc, outc),
          )
          self.pres = pres
          self.vres = vres
  
      def forward(self, z):
          x = initial_voxelize(z, self.pres, self.vres)
          x = self.net(x)
          out = voxel_to_point(x, z, nearest=False)
          out.F = out.F + self.point_transforms(z.F)
          return out
  