content: """
  decouple the trainer with the renderer
  """
  import os
  import cv2 as cv
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  import numpy as np
  
  import trimesh
  
  from utils.misc_utils import visualize_depth_numpy
  
  from utils.training_utils import numpy2tensor
  
  from loss.depth_loss import DepthLoss, DepthSmoothLoss
  
  from models.sparse_neus_renderer import SparseNeuSRenderer
  
  
  class GenericTrainer(nn.Module):
      def __init__(self,
                   rendering_network_outside,
                   pyramid_feature_network_lod0,
                   pyramid_feature_network_lod1,
                   sdf_network_lod0,
                   sdf_network_lod1,
                   variance_network_lod0,
                   variance_network_lod1,
                   rendering_network_lod0,
                   rendering_network_lod1,
                   n_samples_lod0,
                   n_importance_lod0,
                   n_samples_lod1,
                   n_importance_lod1,
                   n_outside,
                   perturb,
                   alpha_type='div',
                   conf=None,
                   timestamp="",
                   mode='train',
                   base_exp_dir=None,
                   ):
          super(GenericTrainer, self).__init__()
  
          self.conf = conf
          self.timestamp = timestamp
  
          
          self.base_exp_dir = base_exp_dir 
          
  
          self.anneal_start = self.conf.get_float('train.anneal_start', default=0.0)
          self.anneal_end = self.conf.get_float('train.anneal_end', default=0.0)
          self.anneal_start_lod1 = self.conf.get_float('train.anneal_start_lod1', default=0.0)
          self.anneal_end_lod1 = self.conf.get_float('train.anneal_end_lod1', default=0.0)
  
          # network setups
          self.rendering_network_outside = rendering_network_outside
          self.pyramid_feature_network_geometry_lod0 = pyramid_feature_network_lod0  # 2D pyramid feature network for geometry
          self.pyramid_feature_network_geometry_lod1 = pyramid_feature_network_lod1  # use differnet networks for the two lods
  
          # when num_lods==2, may consume too much memeory
          self.sdf_network_lod0 = sdf_network_lod0
          self.sdf_network_lod1 = sdf_network_lod1
  
          # - warpped by ModuleList to support DataParallel
          self.variance_network_lod0 = variance_network_lod0
          self.variance_network_lod1 = variance_network_lod1
  
          self.rendering_network_lod0 = rendering_network_lod0
          self.rendering_network_lod1 = rendering_network_lod1
  
          self.n_samples_lod0 = n_samples_lod0
          self.n_importance_lod0 = n_importance_lod0
          self.n_samples_lod1 = n_samples_lod1
          self.n_importance_lod1 = n_importance_lod1
          self.n_outside = n_outside
          self.num_lods = conf.get_int('model.num_lods')  # the number of octree lods
          self.perturb = perturb
          self.alpha_type = alpha_type
  
          # - the two renderers
          self.sdf_renderer_lod0 = SparseNeuSRenderer(
              self.rendering_network_outside,
              self.sdf_network_lod0,
              self.variance_network_lod0,
              self.rendering_network_lod0,
              self.n_samples_lod0,
              self.n_importance_lod0,
              self.n_outside,
              self.perturb,
              alpha_type='div',
              conf=self.conf)
  
          self.sdf_renderer_lod1 = SparseNeuSRenderer(
              self.rendering_network_outside,
              self.sdf_network_lod1,
              self.variance_network_lod1,
              self.rendering_network_lod1,
              self.n_samples_lod1,
              self.n_importance_lod1,
              self.n_outside,
              self.perturb,
              alpha_type='div',
              conf=self.conf)
  
          self.if_fix_lod0_networks = self.conf.get_bool('train.if_fix_lod0_networks')
  
          # sdf network weights
          self.sdf_igr_weight = self.conf.get_float('train.sdf_igr_weight')
          self.sdf_sparse_weight = self.conf.get_float('train.sdf_sparse_weight', default=0)
          self.sdf_decay_param = self.conf.get_float('train.sdf_decay_param', default=100)
          self.fg_bg_weight = self.conf.get_float('train.fg_bg_weight', default=0.00)
          self.bg_ratio = self.conf.get_float('train.bg_ratio', default=0.0)
  
          self.depth_loss_weight = self.conf.get_float('train.depth_loss_weight', default=1.00)
  
          print("depth_loss_weight: ", self.depth_loss_weight)
          self.depth_criterion = DepthLoss()
  
          # - DataParallel mode, cannot modify attributes in forward()
          # self.iter_step = 0
          self.val_mesh_freq = self.conf.get_int('train.val_mesh_freq')
  
          # - True for finetuning; False for general training
          self.if_fitted_rendering = self.conf.get_bool('train.if_fitted_rendering', default=False)
  
          self.prune_depth_filter = self.conf.get_bool('model.prune_depth_filter', default=False)
  
      def get_trainable_params(self):
          # set trainable params
  
          self.params_to_train = []
  
          if not self.if_fix_lod0_networks:
              #  load pretrained featurenet
              self.params_to_train += list(self.pyramid_feature_network_geometry_lod0.parameters())
              self.params_to_train += list(self.sdf_network_lod0.parameters())
              self.params_to_train += list(self.variance_network_lod0.parameters())
  
              if self.rendering_network_lod0 is not None:
                  self.params_to_train += list(self.rendering_network_lod0.parameters())
  
          if self.sdf_network_lod1 is not None:
              #  load pretrained featurenet
              self.params_to_train += list(self.pyramid_feature_network_geometry_lod1.parameters())
  
              self.params_to_train += list(self.sdf_network_lod1.parameters())
              self.params_to_train += list(self.variance_network_lod1.parameters())
              if self.rendering_network_lod1 is not None:
                  self.params_to_train += list(self.rendering_network_lod1.parameters())
  
          return self.params_to_train
  
      def train_step(self, sample,
                     perturb_overwrite=-1,
                     background_rgb=None,
                     alpha_inter_ratio_lod0=0.0,
                     alpha_inter_ratio_lod1=0.0,
                     iter_step=0,
                     ):
          # * only support batch_size==1
          # ! attention: the list of string cannot be splitted in DataParallel
          batch_idx = sample['batch_idx'][0]
          meta = sample['meta'][batch_idx]  # the scan lighting ref_view info
  
          sizeW = sample['img_wh'][0][0]
          sizeH = sample['img_wh'][0][1]
          partial_vol_origin = sample['partial_vol_origin']  # [B, 3]
          near, far = sample['near_fars'][0, 0, :1], sample['near_fars'][0, 0, 1:]
  
          # the full-size ray variables
          sample_rays = sample['rays']
          rays_o = sample_rays['rays_o'][0]
          rays_d = sample_rays['rays_v'][0]
  
          imgs = sample['images'][0]
          intrinsics = sample['intrinsics'][0]
          intrinsics_l_4x = intrinsics.clone()
          intrinsics_l_4x[:, :2] *= 0.25
          w2cs = sample['w2cs'][0]
          c2ws = sample['c2ws'][0]
          proj_matrices = sample['affine_mats']
          scale_mat = sample['scale_mat']
          trans_mat = sample['trans_mat']
  
          # ***********************     Lod==0     ***********************
          if not self.if_fix_lod0_networks:
              geometry_feature_maps = self.obtain_pyramid_feature_maps(imgs)
  
              conditional_features_lod0 = self.sdf_network_lod0.get_conditional_volume(
                  feature_maps=geometry_feature_maps[None, 1:, :, :, :],
                  partial_vol_origin=partial_vol_origin,
                  proj_mats=proj_matrices[:,1:],
                  # proj_mats=proj_matrices,
                  sizeH=sizeH,
                  sizeW=sizeW,
                  lod=0,
              )
  
          else:
              with torch.no_grad():
                  geometry_feature_maps = self.obtain_pyramid_feature_maps(imgs, lod=0)
                  conditional_features_lod0 = self.sdf_network_lod0.get_conditional_volume(
                      feature_maps=geometry_feature_maps[None, 1:, :, :, :],
                      partial_vol_origin=partial_vol_origin,
                      proj_mats=proj_matrices[:,1:],
                      # proj_mats=proj_matrices,
                      sizeH=sizeH,
                      sizeW=sizeW,
                      lod=0,
                  )
  
          con_volume_lod0 = conditional_features_lod0['dense_volume_scale0']
  
          con_valid_mask_volume_lod0 = conditional_features_lod0['valid_mask_volume_scale0']
  
          coords_lod0 = conditional_features_lod0['coords_scale0']  # [1,3,wX,wY,wZ]
  
          # * extract depth maps for all the images
          depth_maps_lod0, depth_masks_lod0 = None, None
          if self.num_lods > 1:
              sdf_volume_lod0 = self.sdf_network_lod0.get_sdf_volume(
                  con_volume_lod0, con_valid_mask_volume_lod0,
                  coords_lod0, partial_vol_origin)  # [1, 1, dX, dY, dZ]
  
          if self.prune_depth_filter:
              depth_maps_lod0_l4x, depth_masks_lod0_l4x = self.sdf_renderer_lod0.extract_depth_maps(
                  self.sdf_network_lod0, sdf_volume_lod0, intrinsics_l_4x, c2ws,
                  sizeH // 4, sizeW // 4, near * 1.5, far)
              depth_maps_lod0 = F.interpolate(depth_maps_lod0_l4x, size=(sizeH, sizeW), mode='bilinear',
                                              align_corners=True)
  
          # *************** losses
          loss_lod0, losses_lod0, depth_statis_lod0 = None, None, None
  
          if not self.if_fix_lod0_networks:
  
              render_out = self.sdf_renderer_lod0.render(
                  rays_o, rays_d, near, far,
                  self.sdf_network_lod0,
                  self.rendering_network_lod0,
                  background_rgb=background_rgb,
                  alpha_inter_ratio=alpha_inter_ratio_lod0,
                  # * related to conditional feature
                  lod=0,
                  conditional_volume=con_volume_lod0,
                  conditional_valid_mask_volume=con_valid_mask_volume_lod0,
                  # * 2d feature maps
                  feature_maps=geometry_feature_maps,
                  color_maps=imgs,
                  w2cs=w2cs,
                  intrinsics=intrinsics,
                  img_wh=[sizeW, sizeH],
                  if_general_rendering=True,
                  if_render_with_grad=True,
              )
  
              loss_lod0, losses_lod0, depth_statis_lod0 = self.cal_losses_sdf(render_out, sample_rays,
                                                                                       iter_step, lod=0)
  
          # ***********************     Lod==1     ***********************
  
          loss_lod1, losses_lod1, depth_statis_lod1 = None, None, None
  
          if self.num_lods > 1:
              geometry_feature_maps_lod1 = self.obtain_pyramid_feature_maps(imgs, lod=1)
              # geometry_feature_maps_lod1 = self.obtain_pyramid_feature_maps(imgs, lod=1)
              if self.prune_depth_filter:
                  pre_coords, pre_feats = self.sdf_renderer_lod0.get_valid_sparse_coords_by_sdf_depthfilter(
                      sdf_volume_lod0[0], coords_lod0[0], con_valid_mask_volume_lod0[0], con_volume_lod0[0],
                      depth_maps_lod0, proj_matrices[0],
                      partial_vol_origin, self.sdf_network_lod0.voxel_size,
                      near, far, self.sdf_network_lod0.voxel_size, 12)
              else:
                  pre_coords, pre_feats = self.sdf_renderer_lod0.get_valid_sparse_coords_by_sdf(
                      sdf_volume_lod0[0], coords_lod0[0], con_valid_mask_volume_lod0[0], con_volume_lod0[0])
  
              pre_coords[:, 1:] = pre_coords[:, 1:] * 2
  
              # ? It seems that training gru_fusion, this part should be trainable too
              conditional_features_lod1 = self.sdf_network_lod1.get_conditional_volume(
                  feature_maps=geometry_feature_maps_lod1[None, 1:, :, :, :],
                  partial_vol_origin=partial_vol_origin,
                  proj_mats=proj_matrices[:,1:],
                  # proj_mats=proj_matrices,
                  sizeH=sizeH,
                  sizeW=sizeW,
                  pre_coords=pre_coords,
                  pre_feats=pre_feats,
              )
  
              con_volume_lod1 = conditional_features_lod1['dense_volume_scale1']
              con_valid_mask_volume_lod1 = conditional_features_lod1['valid_mask_volume_scale1']
  
              # if not self.if_gru_fusion_lod1:
              render_out_lod1 = self.sdf_renderer_lod1.render(
                  rays_o, rays_d, near, far,
                  self.sdf_network_lod1,
                  self.rendering_network_lod1,
                  background_rgb=background_rgb,
                  alpha_inter_ratio=alpha_inter_ratio_lod1,
                  # * related to conditional feature
                  lod=1,
                  conditional_volume=con_volume_lod1,
                  conditional_valid_mask_volume=con_valid_mask_volume_lod1,
                  # * 2d feature maps
                  feature_maps=geometry_feature_maps_lod1,
                  color_maps=imgs,
                  w2cs=w2cs,
                  intrinsics=intrinsics,
                  img_wh=[sizeW, sizeH],
                  bg_ratio=self.bg_ratio,
              )
              loss_lod1, losses_lod1, depth_statis_lod1 = self.cal_losses_sdf(render_out_lod1, sample_rays,
                                                                                       iter_step, lod=1)
  
  
          # # - extract mesh
          if iter_step % self.val_mesh_freq == 0:
              torch.cuda.empty_cache()
              self.validate_mesh(self.sdf_network_lod0,
                                 self.sdf_renderer_lod0.extract_geometry,
                                 conditional_volume=con_volume_lod0, lod=0,
                                 threshold=0,
                                 # occupancy_mask=con_valid_mask_volume_lod0[0, 0],
                                 mode='train_bg', meta=meta,
                                 iter_step=iter_step, scale_mat=scale_mat,
                                 trans_mat=trans_mat)
              torch.cuda.empty_cache()
  
              if self.num_lods > 1:
                  self.validate_mesh(self.sdf_network_lod1,
                                     self.sdf_renderer_lod1.extract_geometry,
                                     conditional_volume=con_volume_lod1, lod=1,
                                     # occupancy_mask=con_valid_mask_volume_lod1[0, 0].detach(),
                                     mode='train_bg', meta=meta,
                                     iter_step=iter_step, scale_mat=scale_mat,
                                     trans_mat=trans_mat)
  
          losses = {
              # - lod 0
              'loss_lod0': loss_lod0,
              'losses_lod0': losses_lod0,
              'depth_statis_lod0': depth_statis_lod0,
  
              # - lod 1
              'loss_lod1': loss_lod1,
              'losses_lod1': losses_lod1,
              'depth_statis_lod1': depth_statis_lod1,
  
          }
  
          return losses
  
      def val_step(self, sample,
                   perturb_overwrite=-1,
                   background_rgb=None,
                   alpha_inter_ratio_lod0=0.0,
                   alpha_inter_ratio_lod1=0.0,
                   iter_step=0,
                   chunk_size=512,
                   save_vis=False,
                   ):
          # * only support batch_size==1
          # ! attention: the list of string cannot be splitted in DataParallel
          batch_idx = sample['batch_idx'][0]
          meta = sample['meta'][batch_idx]  # the scan lighting ref_view info
  
          sizeW = sample['img_wh'][0][0]
          sizeH = sample['img_wh'][0][1]
          H, W = sizeH, sizeW
  
          partial_vol_origin = sample['partial_vol_origin']  # [B, 3]
          near, far = sample['query_near_far'][0, :1], sample['query_near_far'][0, 1:]
  
          # the ray variables
          sample_rays = sample['rays']
          rays_o = sample_rays['rays_o'][0]
          rays_d = sample_rays['rays_v'][0]
          rays_ndc_uv = sample_rays['rays_ndc_uv'][0]
  
          imgs = sample['images'][0]
          intrinsics = sample['intrinsics'][0]
          intrinsics_l_4x = intrinsics.clone()
          intrinsics_l_4x[:, :2] *= 0.25
          w2cs = sample['w2cs'][0]
          c2ws = sample['c2ws'][0]
          proj_matrices = sample['affine_mats']
  
          # render_img_idx = sample['render_img_idx'][0]
          # true_img = sample['images'][0][render_img_idx]
  
          # - the image to render
          scale_mat = sample['scale_mat']  # [1,4,4]  used to convert mesh into true scale
          trans_mat = sample['trans_mat']
          query_c2w = sample['query_c2w']  # [1,4,4]
          query_w2c = sample['query_w2c']  # [1,4,4]
          true_img = sample['query_image'][0]
          true_img = np.uint8(true_img.permute(1, 2, 0).cpu().numpy() * 255)
  
          depth_min, depth_max = near.cpu().numpy(), far.cpu().numpy()
  
          scale_factor = sample['scale_factor'][0].cpu().numpy()
          true_depth = sample['query_depth'] if 'query_depth' in sample.keys() else None
          if true_depth is not None:
              true_depth = true_depth[0].cpu().numpy()
              true_depth_colored = visualize_depth_numpy(true_depth, [depth_min, depth_max])[0]
          else:
              true_depth_colored = None
  
          rays_o = rays_o.reshape(-1, 3).split(chunk_size)
          rays_d = rays_d.reshape(-1, 3).split(chunk_size)
  
          # - obtain conditional features
          with torch.no_grad():
              # - obtain conditional features
              geometry_feature_maps = self.obtain_pyramid_feature_maps(imgs, lod=0)
  
              # - lod 0
              conditional_features_lod0 = self.sdf_network_lod0.get_conditional_volume(
                  feature_maps=geometry_feature_maps[None, :, :, :, :],
                  partial_vol_origin=partial_vol_origin,
                  proj_mats=proj_matrices,
                  sizeH=sizeH,
                  sizeW=sizeW,
                  lod=0,
              )
  
          con_volume_lod0 = conditional_features_lod0['dense_volume_scale0']
          con_valid_mask_volume_lod0 = conditional_features_lod0['valid_mask_volume_scale0']
          coords_lod0 = conditional_features_lod0['coords_scale0']  # [1,3,wX,wY,wZ]
  
          if self.num_lods > 1:
              sdf_volume_lod0 = self.sdf_network_lod0.get_sdf_volume(
                  con_volume_lod0, con_valid_mask_volume_lod0,
                  coords_lod0, partial_vol_origin)  # [1, 1, dX, dY, dZ]
  
          depth_maps_lod0, depth_masks_lod0 = None, None
          if self.prune_depth_filter:
              depth_maps_lod0_l4x, depth_masks_lod0_l4x = self.sdf_renderer_lod0.extract_depth_maps(
                  self.sdf_network_lod0, sdf_volume_lod0,
                  intrinsics_l_4x, c2ws,
                  sizeH // 4, sizeW // 4, near * 1.5, far)  # - near*1.5 is a experienced number
              depth_maps_lod0 = F.interpolate(depth_maps_lod0_l4x, size=(sizeH, sizeW), mode='bilinear',
                                              align_corners=True)
              depth_masks_lod0 = F.interpolate(depth_masks_lod0_l4x.float(), size=(sizeH, sizeW), mode='nearest')
  
              #### visualize the depth_maps_lod0 for checking
              colored_depth_maps_lod0 = []
              for i in range(depth_maps_lod0.shape[0]):
                  colored_depth_maps_lod0.append(
                      visualize_depth_numpy(depth_maps_lod0[i, 0].cpu().numpy(), [depth_min, depth_max])[0])
  
              colored_depth_maps_lod0 = np.concatenate(colored_depth_maps_lod0, axis=0).astype(np.uint8)
              os.makedirs(os.path.join(self.base_exp_dir, 'depth_maps_lod0'), exist_ok=True)
              cv.imwrite(os.path.join(self.base_exp_dir, 'depth_maps_lod0',
                                      '{:0>8d}_{}.png'.format(iter_step, meta)),
                         colored_depth_maps_lod0[:, :, ::-1])
  
          if self.num_lods > 1:
              geometry_feature_maps_lod1 = self.obtain_pyramid_feature_maps(imgs, lod=1)
  
              if self.prune_depth_filter:
                  pre_coords, pre_feats = self.sdf_renderer_lod0.get_valid_sparse_coords_by_sdf_depthfilter(
                      sdf_volume_lod0[0], coords_lod0[0], con_valid_mask_volume_lod0[0], con_volume_lod0[0],
                      depth_maps_lod0, proj_matrices[0],
                      partial_vol_origin, self.sdf_network_lod0.voxel_size,
                      near, far, self.sdf_network_lod0.voxel_size, 12)
              else:
                  pre_coords, pre_feats = self.sdf_renderer_lod0.get_valid_sparse_coords_by_sdf(
                      sdf_volume_lod0[0], coords_lod0[0], con_valid_mask_volume_lod0[0], con_volume_lod0[0])
  
              pre_coords[:, 1:] = pre_coords[:, 1:] * 2
  
              with torch.no_grad():
                  conditional_features_lod1 = self.sdf_network_lod1.get_conditional_volume(
                      feature_maps=geometry_feature_maps_lod1[None, :, :, :, :],
                      partial_vol_origin=partial_vol_origin,
                      proj_mats=proj_matrices,
                      sizeH=sizeH,
                      sizeW=sizeW,
                      pre_coords=pre_coords,
                      pre_feats=pre_feats,
                  )
  
              con_volume_lod1 = conditional_features_lod1['dense_volume_scale1']
              con_valid_mask_volume_lod1 = conditional_features_lod1['valid_mask_volume_scale1']
  
          out_rgb_fine = []
          out_normal_fine = []
          out_depth_fine = []
  
          out_rgb_fine_lod1 = []
          out_normal_fine_lod1 = []
          out_depth_fine_lod1 = []
  
          # out_depth_fine_explicit = []
          if save_vis:
              for rays_o_batch, rays_d_batch in zip(rays_o, rays_d):
  
                  # ****** lod 0 ****
                  render_out = self.sdf_renderer_lod0.render(
                      rays_o_batch, rays_d_batch, near, far,
                      self.sdf_network_lod0,
                      self.rendering_network_lod0,
                      background_rgb=background_rgb,
                      alpha_inter_ratio=alpha_inter_ratio_lod0,
                      # * related to conditional feature
                      lod=0,
                      conditional_volume=con_volume_lod0,
                      conditional_valid_mask_volume=con_valid_mask_volume_lod0,
                      # * 2d feature maps
                      feature_maps=geometry_feature_maps,
                      color_maps=imgs,
                      w2cs=w2cs,
                      intrinsics=intrinsics,
                      img_wh=[sizeW, sizeH],
                      query_c2w=query_c2w,
                      if_render_with_grad=False,
                  )
  
                  feasible = lambda key: ((key in render_out) and (render_out[key] is not None))
  
                  if feasible('depth'):
                      out_depth_fine.append(render_out['depth'].detach().cpu().numpy())
  
                  # if render_out['color_coarse'] is not None:
                  if feasible('color_fine'):
                      out_rgb_fine.append(render_out['color_fine'].detach().cpu().numpy())
                  if feasible('gradients') and feasible('weights'):
                      if render_out['inside_sphere'] is not None:
                          out_normal_fine.append((render_out['gradients'] * render_out['weights'][:,
                                                                            :self.n_samples_lod0 + self.n_importance_lod0,
                                                                            None] * render_out['inside_sphere'][
                                                      ..., None]).sum(dim=1).detach().cpu().numpy())
                      else:
                          out_normal_fine.append((render_out['gradients'] * render_out['weights'][:,
                                                                            :self.n_samples_lod0 + self.n_importance_lod0,
                                                                            None]).sum(dim=1).detach().cpu().numpy())
                  del render_out
  
              # ****************** lod 1 **************************
              if self.num_lods > 1:
                  for rays_o_batch, rays_d_batch in zip(rays_o, rays_d):
                      render_out_lod1 = self.sdf_renderer_lod1.render(
                          rays_o_batch, rays_d_batch, near, far,
                          self.sdf_network_lod1,
                          self.rendering_network_lod1,
                          background_rgb=background_rgb,
                          alpha_inter_ratio=alpha_inter_ratio_lod1,
                          # * related to conditional feature
                          lod=1,
                          conditional_volume=con_volume_lod1,
                          conditional_valid_mask_volume=con_valid_mask_volume_lod1,
                          # * 2d feature maps
                          feature_maps=geometry_feature_maps_lod1,
                          color_maps=imgs,
                          w2cs=w2cs,
                          intrinsics=intrinsics,
                          img_wh=[sizeW, sizeH],
                          query_c2w=query_c2w,
                          if_render_with_grad=False,
                      )
  
                      feasible = lambda key: ((key in render_out_lod1) and (render_out_lod1[key] is not None))
  
                      if feasible('depth'):
                          out_depth_fine_lod1.append(render_out_lod1['depth'].detach().cpu().numpy())
  
                      # if render_out['color_coarse'] is not None:
                      if feasible('color_fine'):
                          out_rgb_fine_lod1.append(render_out_lod1['color_fine'].detach().cpu().numpy())
                      if feasible('gradients') and feasible('weights'):
                          if render_out_lod1['inside_sphere'] is not None:
                              out_normal_fine_lod1.append((render_out_lod1['gradients'] * render_out_lod1['weights'][:,
                                                                                          :self.n_samples_lod1 + self.n_importance_lod1,
                                                                                          None] *
                                                           render_out_lod1['inside_sphere'][
                                                               ..., None]).sum(dim=1).detach().cpu().numpy())
                          else:
                              out_normal_fine_lod1.append((render_out_lod1['gradients'] * render_out_lod1['weights'][:,
                                                                                          :self.n_samples_lod1 + self.n_importance_lod1,
                                                                                          None]).sum(
                                  dim=1).detach().cpu().numpy())
                      del render_out_lod1
  
              # - save visualization of lod 0
  
              self.save_visualization(true_img, true_depth_colored, out_depth_fine, out_normal_fine,
                                      query_w2c[0], out_rgb_fine, H, W,
                                      depth_min, depth_max, iter_step, meta, "val_lod0", true_depth=true_depth, scale_factor=scale_factor)
  
              if self.num_lods > 1:
                  self.save_visualization(true_img, true_depth_colored, out_depth_fine_lod1, out_normal_fine_lod1,
                                          query_w2c[0], out_rgb_fine_lod1, H, W,
                                          depth_min, depth_max, iter_step, meta, "val_lod1", true_depth=true_depth, scale_factor=scale_factor)
  
          # - extract mesh
          if (iter_step % self.val_mesh_freq == 0):
              torch.cuda.empty_cache()
              self.validate_mesh(self.sdf_network_lod0,
                                 self.sdf_renderer_lod0.extract_geometry,
                                 conditional_volume=con_volume_lod0, lod=0,
                                 threshold=0,
                                 # occupancy_mask=con_valid_mask_volume_lod0[0, 0],
                                 mode='val_bg', meta=meta,
                                 iter_step=iter_step, scale_mat=scale_mat, trans_mat=trans_mat)
              torch.cuda.empty_cache()
  
              if self.num_lods > 1:
                  self.validate_mesh(self.sdf_network_lod1,
                                     self.sdf_renderer_lod1.extract_geometry,
                                     conditional_volume=con_volume_lod1, lod=1,
                                     # occupancy_mask=con_valid_mask_volume_lod1[0, 0].detach(),
                                     mode='val_bg', meta=meta,
                                     iter_step=iter_step, scale_mat=scale_mat, trans_mat=trans_mat)
  
              torch.cuda.empty_cache()
  
      @torch.no_grad()
      def get_metrics_step(self, sample,
                     perturb_overwrite=-1,
                     background_rgb=None,
                     alpha_inter_ratio_lod0=0.0,
                     alpha_inter_ratio_lod1=0.0,
                     iter_step=0,
                     ):
          # * only support batch_size==1
          # ! attention: the list of string cannot be splitted in DataParallel
          batch_idx = sample['batch_idx'][0]
          meta = sample['meta'][batch_idx]  # the scan lighting ref_view info
  
          sizeW = sample['img_wh'][0][0]
          sizeH = sample['img_wh'][0][1]
          partial_vol_origin = sample['partial_vol_origin']  # [B, 3]
          near, far = sample['near_fars'][0, 0, :1], sample['near_fars'][0, 0, 1:]
  
          # the full-size ray variables
          sample_rays = sample['rays']
          rays_o = sample_rays['rays_o'][0]
          rays_d = sample_rays['rays_v'][0]
  
          imgs = sample['images'][0]
          intrinsics = sample['intrinsics'][0]
          intrinsics_l_4x = intrinsics.clone()
          intrinsics_l_4x[:, :2] *= 0.25
          w2cs = sample['w2cs'][0]
          c2ws = sample['c2ws'][0]
          proj_matrices = sample['affine_mats']
          scale_mat = sample['scale_mat']
          trans_mat = sample['trans_mat']
  
          # ***********************     Lod==0     ***********************
          if not self.if_fix_lod0_networks:
              geometry_feature_maps = self.obtain_pyramid_feature_maps(imgs)
  
              conditional_features_lod0 = self.sdf_network_lod0.get_conditional_volume(
                  feature_maps=geometry_feature_maps[None, 1:, :, :, :],
                  partial_vol_origin=partial_vol_origin,
                  proj_mats=proj_matrices[:,1:],
                  # proj_mats=proj_matrices,
                  sizeH=sizeH,
                  sizeW=sizeW,
                  lod=0,
              )
  
          else:
              with torch.no_grad():
                  geometry_feature_maps = self.obtain_pyramid_feature_maps(imgs, lod=0)
                  # geometry_feature_maps = self.obtain_pyramid_feature_maps(imgs, lod=0)
                  conditional_features_lod0 = self.sdf_network_lod0.get_conditional_volume(
                      feature_maps=geometry_feature_maps[None, 1:, :, :, :],
                      partial_vol_origin=partial_vol_origin,
                      proj_mats=proj_matrices[:,1:],
                      # proj_mats=proj_matrices,
                      sizeH=sizeH,
                      sizeW=sizeW,
                      lod=0,
                  )
          con_volume_lod0 = conditional_features_lod0['dense_volume_scale0']
  
          con_valid_mask_volume_lod0 = conditional_features_lod0['valid_mask_volume_scale0']
          coords_lod0 = conditional_features_lod0['coords_scale0']  # [1,3,wX,wY,wZ]
  
          # * extract depth maps for all the images
          depth_maps_lod0, depth_masks_lod0 = None, None
          if self.num_lods > 1:
              sdf_volume_lod0 = self.sdf_network_lod0.get_sdf_volume(
                  con_volume_lod0, con_valid_mask_volume_lod0,
                  coords_lod0, partial_vol_origin)  # [1, 1, dX, dY, dZ]
  
          if self.prune_depth_filter:
              depth_maps_lod0_l4x, depth_masks_lod0_l4x = self.sdf_renderer_lod0.extract_depth_maps(
                  self.sdf_network_lod0, sdf_volume_lod0, intrinsics_l_4x, c2ws,
                  sizeH // 4, sizeW // 4, near * 1.5, far)
              depth_maps_lod0 = F.interpolate(depth_maps_lod0_l4x, size=(sizeH, sizeW), mode='bilinear',
                                              align_corners=True)
              depth_masks_lod0 = F.interpolate(depth_masks_lod0_l4x.float(), size=(sizeH, sizeW), mode='nearest')
  
          # *************** losses
          loss_lod0, losses_lod0, depth_statis_lod0 = None, None, None
  
          if not self.if_fix_lod0_networks:
  
              render_out = self.sdf_renderer_lod0.render(
                  rays_o, rays_d, near, far,
                  self.sdf_network_lod0,
                  self.rendering_network_lod0,
                  background_rgb=background_rgb,
                  alpha_inter_ratio=alpha_inter_ratio_lod0,
                  # * related to conditional feature
                  lod=0,
                  conditional_volume=con_volume_lod0,
  