content: import torch
  from ...modules.sparse import SparseTensor
  from easydict import EasyDict as edict
  from .utils_cube import *
  from .flexicube import FlexiCubes
  
  
  class MeshExtractResult:
      def __init__(self,
          vertices,
          faces,
          vertex_attrs=None,
          res=64
      ):
          self.vertices = vertices
          self.faces = faces.long()
          self.vertex_attrs = vertex_attrs
          self.face_normal = self.comput_face_normals(vertices, faces)
          self.res = res
          self.success = (vertices.shape[0] != 0 and faces.shape[0] != 0)
  
          # training only
          self.tsdf_v = None
          self.tsdf_s = None
          self.reg_loss = None
          
      def comput_face_normals(self, verts, faces):
          i0 = faces[..., 0].long()
          i1 = faces[..., 1].long()
          i2 = faces[..., 2].long()
  
          v0 = verts[i0, :]
          v1 = verts[i1, :]
          v2 = verts[i2, :]
          face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1)
          face_normals = torch.nn.functional.normalize(face_normals, dim=1)
          # print(face_normals.min(), face_normals.max(), face_normals.shape)
          return face_normals[:, None, :].repeat(1, 3, 1)
                  
      def comput_v_normals(self, verts, faces):
          i0 = faces[..., 0].long()
          i1 = faces[..., 1].long()
          i2 = faces[..., 2].long()
  
          v0 = verts[i0, :]
          v1 = verts[i1, :]
          v2 = verts[i2, :]
          face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1)
          v_normals = torch.zeros_like(verts)
          v_normals.scatter_add_(0, i0[..., None].repeat(1, 3), face_normals)
          v_normals.scatter_add_(0, i1[..., None].repeat(1, 3), face_normals)
          v_normals.scatter_add_(0, i2[..., None].repeat(1, 3), face_normals)
  
          v_normals = torch.nn.functional.normalize(v_normals, dim=1)
          return v_normals   
  
  
  class SparseFeatures2Mesh:
      def __init__(self, device="cuda", res=64, use_color=True):
          '''
          a model to generate a mesh from sparse features structures using flexicube
          '''
          super().__init__()
          self.device=device
          self.res = res
          self.mesh_extractor = FlexiCubes(device=device)
          self.sdf_bias = -1.0 / res
          verts, cube = construct_dense_grid(self.res, self.device)
          self.reg_c = cube.to(self.device)
          self.reg_v = verts.to(self.device)
          self.use_color = use_color
          self._calc_layout()
      
      def _calc_layout(self):
          LAYOUTS = {
              'sdf': {'shape': (8, 1), 'size': 8},
              'deform': {'shape': (8, 3), 'size': 8 * 3},
              'weights': {'shape': (21,), 'size': 21}
          }
          if self.use_color:
              '''
              6 channel color including normal map
              '''
              LAYOUTS['color'] = {'shape': (8, 6,), 'size': 8 * 6}
          self.layouts = edict(LAYOUTS)
          start = 0
          for k, v in self.layouts.items():
              v['range'] = (start, start + v['size'])
              start += v['size']
          self.feats_channels = start
          
      def get_layout(self, feats : torch.Tensor, name : str):
          if name not in self.layouts:
              return None
          return feats[:, self.layouts[name]['range'][0]:self.layouts[name]['range'][1]].reshape(-1, *self.layouts[name]['shape'])
      
      def __call__(self, cubefeats : SparseTensor, training=False):
          """
          Generates a mesh based on the specified sparse voxel structures.
          Args:
              cube_attrs [Nx21] : Sparse Tensor attrs about cube weights
              verts_attrs [Nx10] : [0:1] SDF [1:4] deform [4:7] color [7:10] normal 
          Returns:
              return the success tag and ni you loss, 
          """
          # add sdf bias to verts_attrs
          coords = cubefeats.coords[:, 1:]
          feats = cubefeats.feats
          
          sdf, deform, color, weights = [self.get_layout(feats, name) for name in ['sdf', 'deform', 'color', 'weights']]
          sdf += self.sdf_bias
          v_attrs = [sdf, deform, color] if self.use_color else [sdf, deform]
          v_pos, v_attrs, reg_loss = sparse_cube2verts(coords, torch.cat(v_attrs, dim=-1), training=training)
          v_attrs_d = get_dense_attrs(v_pos, v_attrs, res=self.res+1, sdf_init=True)
          weights_d = get_dense_attrs(coords, weights, res=self.res, sdf_init=False)
          if self.use_color:
              sdf_d, deform_d, colors_d = v_attrs_d[..., 0], v_attrs_d[..., 1:4], v_attrs_d[..., 4:]
          else:
              sdf_d, deform_d = v_attrs_d[..., 0], v_attrs_d[..., 1:4]
              colors_d = None
              
          x_nx3 = get_defomed_verts(self.reg_v, deform_d, self.res)
          
          vertices, faces, L_dev, colors = self.mesh_extractor(
              voxelgrid_vertices=x_nx3,
              scalar_field=sdf_d,
              cube_idx=self.reg_c,
              resolution=self.res,
              beta=weights_d[:, :12],
              alpha=weights_d[:, 12:20],
              gamma_f=weights_d[:, 20],
              voxelgrid_colors=colors_d,
              training=training)
          
          mesh = MeshExtractResult(vertices=vertices, faces=faces, vertex_attrs=colors, res=self.res)
          if training:
              if mesh.success:
                  reg_loss += L_dev.mean() * 0.5
              reg_loss += (weights[:,:20]).abs().mean() * 0.2
              mesh.reg_loss = reg_loss
              mesh.tsdf_v = get_defomed_verts(v_pos, v_attrs[:, 1:4], self.res)
              mesh.tsdf_s = v_attrs[:, 0]
          return mesh