content: """
  The codes are heavily borrowed from NeuS
  """
  
  import os
  import cv2 as cv
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import numpy as np
  import logging
  import mcubes
  from icecream import ic
  from models.render_utils import sample_pdf
  
  from models.projector import Projector
  from tsparse.torchsparse_utils import sparse_to_dense_channel
  
  from models.fast_renderer import FastRenderer
  
  from models.patch_projector import PatchProjector
  
  
  class SparseNeuSRenderer(nn.Module):
      """
      conditional neus render;
      optimize on normalized world space;
      warped by nn.Module to support DataParallel traning
      """
  
      def __init__(self,
                   rendering_network_outside,
                   sdf_network,
                   variance_network,
                   rendering_network,
                   n_samples,
                   n_importance,
                   n_outside,
                   perturb,
                   alpha_type='div',
                   conf=None
                   ):
          super(SparseNeuSRenderer, self).__init__()
  
          self.conf = conf
          self.base_exp_dir = conf['general.base_exp_dir']
  
          # network setups
          self.rendering_network_outside = rendering_network_outside
          self.sdf_network = sdf_network
          self.variance_network = variance_network
          self.rendering_network = rendering_network
  
          self.n_samples = n_samples
          self.n_importance = n_importance
          self.n_outside = n_outside
          self.perturb = perturb
          self.alpha_type = alpha_type
  
          self.rendering_projector = Projector()  # used to obtain features for generalized rendering
  
          self.h_patch_size = self.conf.get_int('model.h_patch_size', default=3)
          self.patch_projector = PatchProjector(self.h_patch_size)
  
          self.ray_tracer = FastRenderer()  # ray_tracer to extract depth maps from sdf_volume
  
          # - fitted rendering or general rendering
          try:
              self.if_fitted_rendering = self.sdf_network.if_fitted_rendering
          except:
              self.if_fitted_rendering = False
  
      def up_sample(self, rays_o, rays_d, z_vals, sdf, n_importance, inv_variance,
                    conditional_valid_mask_volume=None):
          device = rays_o.device
          batch_size, n_samples = z_vals.shape
          pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3
  
          if conditional_valid_mask_volume is not None:
              pts_mask = self.get_pts_mask_for_conditional_volume(pts.view(-1, 3), conditional_valid_mask_volume)
              pts_mask = pts_mask.reshape(batch_size, n_samples)
              pts_mask = pts_mask[:, :-1] * pts_mask[:, 1:]  # [batch_size, n_samples-1]
          else:
              pts_mask = torch.ones([batch_size, n_samples]).to(pts.device)
  
          sdf = sdf.reshape(batch_size, n_samples)
          prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
          prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
          mid_sdf = (prev_sdf + next_sdf) * 0.5
          dot_val = None
          if self.alpha_type == 'uniform':
              dot_val = torch.ones([batch_size, n_samples - 1]) * -1.0
          else:
              dot_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)
              prev_dot_val = torch.cat([torch.zeros([batch_size, 1]).to(device), dot_val[:, :-1]], dim=-1)
              dot_val = torch.stack([prev_dot_val, dot_val], dim=-1)
              dot_val, _ = torch.min(dot_val, dim=-1, keepdim=False)
              dot_val = dot_val.clip(-10.0, 0.0) * pts_mask
          dist = (next_z_vals - prev_z_vals)
          prev_esti_sdf = mid_sdf - dot_val * dist * 0.5
          next_esti_sdf = mid_sdf + dot_val * dist * 0.5
          prev_cdf = torch.sigmoid(prev_esti_sdf * inv_variance)
          next_cdf = torch.sigmoid(next_esti_sdf * inv_variance)
          alpha_sdf = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)
  
          alpha = alpha_sdf
  
          # - apply pts_mask
          alpha = pts_mask * alpha
  
          weights = alpha * torch.cumprod(
              torch.cat([torch.ones([batch_size, 1]).to(device), 1. - alpha + 1e-7], -1), -1)[:, :-1]
  
          z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()
          return z_samples
  
      def cat_z_vals(self, rays_o, rays_d, z_vals, new_z_vals, sdf, lod,
                     sdf_network, gru_fusion,
                     # * related to conditional feature
                     conditional_volume=None,
                     conditional_valid_mask_volume=None
                     ):
          device = rays_o.device
          batch_size, n_samples = z_vals.shape
          _, n_importance = new_z_vals.shape
          pts = rays_o[:, None, :] + rays_d[:, None, :] * new_z_vals[..., :, None]
  
          if conditional_valid_mask_volume is not None:
              pts_mask = self.get_pts_mask_for_conditional_volume(pts.view(-1, 3), conditional_valid_mask_volume)
              pts_mask = pts_mask.reshape(batch_size, n_importance)
              pts_mask_bool = (pts_mask > 0).view(-1)
          else:
              pts_mask = torch.ones([batch_size, n_importance]).to(pts.device)
  
          new_sdf = torch.ones([batch_size * n_importance, 1]).to(pts.dtype).to(device) * 100
  
          if torch.sum(pts_mask) > 1:
              new_outputs = sdf_network.sdf(pts.reshape(-1, 3)[pts_mask_bool], conditional_volume, lod=lod)
              new_sdf[pts_mask_bool] = new_outputs['sdf_pts_scale%d' % lod]  # .reshape(batch_size, n_importance)
  
          new_sdf = new_sdf.view(batch_size, n_importance)
  
          z_vals = torch.cat([z_vals, new_z_vals], dim=-1)
          sdf = torch.cat([sdf, new_sdf], dim=-1)
  
          z_vals, index = torch.sort(z_vals, dim=-1)
          xx = torch.arange(batch_size)[:, None].expand(batch_size, n_samples + n_importance).reshape(-1)
          index = index.reshape(-1)
          sdf = sdf[(xx, index)].reshape(batch_size, n_samples + n_importance)
  
          return z_vals, sdf
  
      @torch.no_grad()
      def get_pts_mask_for_conditional_volume(self, pts, mask_volume):
          """
  
          :param pts: [N, 3]
          :param mask_volume: [1, 1, X, Y, Z]
          :return:
          """
          num_pts = pts.shape[0]
          pts = pts.view(1, 1, 1, num_pts, 3)  # - should be in range (-1, 1)
  
          pts = torch.flip(pts, dims=[-1])
  
          pts_mask = F.grid_sample(mask_volume, pts, mode='nearest')  # [1, c, 1, 1, num_pts]
          pts_mask = pts_mask.view(-1, num_pts).permute(1, 0).contiguous()  # [num_pts, 1]
  
          return pts_mask
  
      def render_core(self,
                      rays_o,
                      rays_d,
                      z_vals,
                      sample_dist,
                      lod,
                      sdf_network,
                      rendering_network,
                      background_alpha=None,  # - no use here
                      background_sampled_color=None,  # - no use here
                      background_rgb=None,  # - no use here
                      alpha_inter_ratio=0.0,
                      # * related to conditional feature
                      conditional_volume=None,
                      conditional_valid_mask_volume=None,
                      # * 2d feature maps
                      feature_maps=None,
                      color_maps=None,
                      w2cs=None,
                      intrinsics=None,
                      img_wh=None,
                      query_c2w=None,  # - used for testing
                      if_general_rendering=True,
                      if_render_with_grad=True,
                      # * used for blending mlp rendering network
                      img_index=None,
                      rays_uv=None,
                      # * used for clear bg and fg
                      bg_num=0
                      ):
          device = rays_o.device
          N_rays = rays_o.shape[0]
          _, n_samples = z_vals.shape
          dists = z_vals[..., 1:] - z_vals[..., :-1]
          dists = torch.cat([dists, torch.Tensor([sample_dist]).expand(dists[..., :1].shape).to(device)], -1)
  
          mid_z_vals = z_vals + dists * 0.5
          mid_dists = mid_z_vals[..., 1:] - mid_z_vals[..., :-1]
  
          pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None]  # n_rays, n_samples, 3
          dirs = rays_d[:, None, :].expand(pts.shape)
  
          pts = pts.reshape(-1, 3)
          dirs = dirs.reshape(-1, 3)
  
          # * if conditional_volume is restored from sparse volume, need mask for pts
          if conditional_valid_mask_volume is not None:
              pts_mask = self.get_pts_mask_for_conditional_volume(pts, conditional_valid_mask_volume)
              pts_mask = pts_mask.reshape(N_rays, n_samples).float().detach()
              pts_mask_bool = (pts_mask > 0).view(-1)
  
              if torch.sum(pts_mask_bool.float()) < 1:  # ! when render out image, may meet this problem
                  pts_mask_bool[:100] = True
  
          else:
              pts_mask = torch.ones([N_rays, n_samples]).to(pts.device)
          # import ipdb; ipdb.set_trace()
          # pts_valid = pts[pts_mask_bool]
          sdf_nn_output = sdf_network.sdf(pts[pts_mask_bool], conditional_volume, lod=lod)
  
          sdf = torch.ones([N_rays * n_samples, 1]).to(pts.dtype).to(device) * 100
          sdf[pts_mask_bool] = sdf_nn_output['sdf_pts_scale%d' % lod]  # [N_rays*n_samples, 1]
          feature_vector_valid = sdf_nn_output['sdf_features_pts_scale%d' % lod]
          feature_vector = torch.zeros([N_rays * n_samples, feature_vector_valid.shape[1]]).to(pts.dtype).to(device)
          feature_vector[pts_mask_bool] = feature_vector_valid
  
          # * estimate alpha from sdf
          gradients = torch.zeros([N_rays * n_samples, 3]).to(pts.dtype).to(device)
          # import ipdb; ipdb.set_trace()
          gradients[pts_mask_bool] = sdf_network.gradient(
              pts[pts_mask_bool], conditional_volume, lod=lod).squeeze()
  
          sampled_color_mlp = None
          rendering_valid_mask_mlp = None
          sampled_color_patch = None
          rendering_patch_mask = None
  
          if self.if_fitted_rendering:  # used for fine-tuning
              position_latent = sdf_nn_output['sampled_latent_scale%d' % lod]
              sampled_color_mlp = torch.zeros([N_rays * n_samples, 3]).to(pts.dtype).to(device)
              sampled_color_mlp_mask = torch.zeros([N_rays * n_samples, 1]).to(pts.dtype).to(device)
  
              # - extract pixel
              pts_pixel_color, pts_pixel_mask = self.patch_projector.pixel_warp(
                  pts[pts_mask_bool][:, None, :], color_maps, intrinsics,
                  w2cs, img_wh=None)  # [N_rays * n_samples,1, N_views,  3] , [N_rays*n_samples, 1, N_views]
              pts_pixel_color = pts_pixel_color[:, 0, :, :]  # [N_rays * n_samples, N_views,  3]
              pts_pixel_mask = pts_pixel_mask[:, 0, :]  # [N_rays*n_samples, N_views]
  
              # - extract patch
              if_patch_blending = False if rays_uv is None else True
              pts_patch_color, pts_patch_mask = None, None
              if if_patch_blending:
                  pts_patch_color, pts_patch_mask = self.patch_projector.patch_warp(
                      pts.reshape([N_rays, n_samples, 3]),
                      rays_uv, gradients.reshape([N_rays, n_samples, 3]),
                      color_maps,
                      intrinsics[0], intrinsics,
                      query_c2w[0], torch.inverse(w2cs), img_wh=None
                  )  # (N_rays, n_samples, N_src, Npx, 3), (N_rays, n_samples, N_src, Npx)
                  N_src, Npx = pts_patch_mask.shape[2:]
                  pts_patch_color = pts_patch_color.view(N_rays * n_samples, N_src, Npx, 3)[pts_mask_bool]
                  pts_patch_mask = pts_patch_mask.view(N_rays * n_samples, N_src, Npx)[pts_mask_bool]
  
                  sampled_color_patch = torch.zeros([N_rays * n_samples, Npx, 3]).to(device)
                  sampled_color_patch_mask = torch.zeros([N_rays * n_samples, 1]).to(device)
  
              sampled_color_mlp_, sampled_color_mlp_mask_, \
              sampled_color_patch_, sampled_color_patch_mask_ = sdf_network.color_blend(
                  pts[pts_mask_bool],
                  position_latent,
                  gradients[pts_mask_bool],
                  dirs[pts_mask_bool],
                  feature_vector[pts_mask_bool],
                  img_index=img_index,
                  pts_pixel_color=pts_pixel_color,
                  pts_pixel_mask=pts_pixel_mask,
                  pts_patch_color=pts_patch_color,
                  pts_patch_mask=pts_patch_mask
  
              )  # [n, 3], [n, 1]
              sampled_color_mlp[pts_mask_bool] = sampled_color_mlp_
              sampled_color_mlp_mask[pts_mask_bool] = sampled_color_mlp_mask_.float()
              sampled_color_mlp = sampled_color_mlp.view(N_rays, n_samples, 3)
              sampled_color_mlp_mask = sampled_color_mlp_mask.view(N_rays, n_samples)
              rendering_valid_mask_mlp = torch.mean(pts_mask * sampled_color_mlp_mask, dim=-1, keepdim=True) > 0.5
  
              # patch blending
              if if_patch_blending:
                  sampled_color_patch[pts_mask_bool] = sampled_color_patch_
                  sampled_color_patch_mask[pts_mask_bool] = sampled_color_patch_mask_.float()
                  sampled_color_patch = sampled_color_patch.view(N_rays, n_samples, Npx, 3)
                  sampled_color_patch_mask = sampled_color_patch_mask.view(N_rays, n_samples)
                  rendering_patch_mask = torch.mean(pts_mask * sampled_color_patch_mask, dim=-1,
                                                    keepdim=True) > 0.5  # [N_rays, 1]
              else:
                  sampled_color_patch, rendering_patch_mask = None, None
  
          if if_general_rendering:  # used for general training
              # [512, 128, 16]; [4, 512, 128, 59]; [4, 512, 128, 4]
              ren_geo_feats, ren_rgb_feats, ren_ray_diff, ren_mask, _, _ = self.rendering_projector.compute(
                  pts.view(N_rays, n_samples, 3),
                  # * 3d geometry feature volumes
                  geometryVolume=conditional_volume[0],
                  geometryVolumeMask=conditional_valid_mask_volume[0],
                  # * 2d rendering feature maps
                  rendering_feature_maps=feature_maps, # [n_views, 56, 256, 256]
                  color_maps=color_maps,
                  w2cs=w2cs,
                  intrinsics=intrinsics,
                  img_wh=img_wh,
                  query_img_idx=0,  # the index of the N_views dim for rendering
                  query_c2w=query_c2w,
              )
  
              # (N_rays, n_samples, 3)
              if if_render_with_grad:
                  # import ipdb; ipdb.set_trace()
                  # [nrays, 3] [nrays, 1]
                  sampled_color, rendering_valid_mask = rendering_network(
                      ren_geo_feats, ren_rgb_feats, ren_ray_diff, ren_mask)
                  # import ipdb; ipdb.set_trace()
              else:
                  with torch.no_grad():
                      sampled_color, rendering_valid_mask = rendering_network(
                          ren_geo_feats, ren_rgb_feats, ren_ray_diff, ren_mask)
          else:
              sampled_color, rendering_valid_mask = None, None
  
          inv_variance = self.variance_network(feature_vector)[:, :1].clip(1e-6, 1e6)
  
          true_dot_val = (dirs * gradients).sum(-1, keepdim=True)  # * calculate
  
          iter_cos = -(F.relu(-true_dot_val * 0.5 + 0.5) * (1.0 - alpha_inter_ratio) + F.relu(
              -true_dot_val) * alpha_inter_ratio)  # always non-positive
  
          iter_cos = iter_cos * pts_mask.view(-1, 1)
  
          true_estimate_sdf_half_next = sdf + iter_cos.clip(-10.0, 10.0) * dists.reshape(-1, 1) * 0.5
          true_estimate_sdf_half_prev = sdf - iter_cos.clip(-10.0, 10.0) * dists.reshape(-1, 1) * 0.5
  
          prev_cdf = torch.sigmoid(true_estimate_sdf_half_prev * inv_variance)
          next_cdf = torch.sigmoid(true_estimate_sdf_half_next * inv_variance)
  
          p = prev_cdf - next_cdf
          c = prev_cdf
  
          if self.alpha_type == 'div':
              alpha_sdf = ((p + 1e-5) / (c + 1e-5)).reshape(N_rays, n_samples).clip(0.0, 1.0)
          elif self.alpha_type == 'uniform':
              uniform_estimate_sdf_half_next = sdf - dists.reshape(-1, 1) * 0.5
              uniform_estimate_sdf_half_prev = sdf + dists.reshape(-1, 1) * 0.5
              uniform_prev_cdf = torch.sigmoid(uniform_estimate_sdf_half_prev * inv_variance)
              uniform_next_cdf = torch.sigmoid(uniform_estimate_sdf_half_next * inv_variance)
              uniform_alpha = F.relu(
                  (uniform_prev_cdf - uniform_next_cdf + 1e-5) / (uniform_prev_cdf + 1e-5)).reshape(
                  N_rays, n_samples).clip(0.0, 1.0)
              alpha_sdf = uniform_alpha
          else:
              assert False
  
          alpha = alpha_sdf
  
          # - apply pts_mask
          alpha = alpha * pts_mask
  
          # pts_radius = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=True).reshape(N_rays, n_samples)
          # inside_sphere = (pts_radius < 1.0).float().detach()
          # relax_inside_sphere = (pts_radius < 1.2).float().detach()
          inside_sphere = pts_mask
          relax_inside_sphere = pts_mask
  
          weights = alpha * torch.cumprod(torch.cat([torch.ones([N_rays, 1]).to(device), 1. - alpha + 1e-7], -1), -1)[:,
                            :-1]  # n_rays, n_samples
          weights_sum = weights.sum(dim=-1, keepdim=True)
          alpha_sum = alpha.sum(dim=-1, keepdim=True)
  
          if bg_num > 0:
              weights_sum_fg = weights[:, :-bg_num].sum(dim=-1, keepdim=True)
          else:
              weights_sum_fg = weights_sum
  
          if sampled_color is not None:
              color = (sampled_color * weights[:, :, None]).sum(dim=1)
          else:
              color = None
          # import ipdb; ipdb.set_trace()
  
          if background_rgb is not None and color is not None:
              color = color + background_rgb * (1.0 - weights_sum)
              # print("color device:" + str(color.device))
          # if color is not None:
          #     # import ipdb; ipdb.set_trace()
          #     color = color + (1.0 - weights_sum)
  
  
          ###################*  mlp color rendering  #####################
          color_mlp = None
          # import ipdb; ipdb.set_trace()
          if sampled_color_mlp is not None:
              color_mlp = (sampled_color_mlp * weights[:, :, None]).sum(dim=1)
  
          if background_rgb is not None and color_mlp is not None:
              color_mlp = color_mlp + background_rgb * (1.0 - weights_sum)
  
          ############################ *  patch blending  ################
          blended_color_patch = None
          if sampled_color_patch is not None:
              blended_color_patch = (sampled_color_patch * weights[:, :, None, None]).sum(dim=1)  # [N_rays, Npx, 3]
  
          ######################################################
  
          gradient_error = (torch.linalg.norm(gradients.reshape(N_rays, n_samples, 3), ord=2,
                                              dim=-1) - 1.0) ** 2
          # ! the gradient normal should be masked out, the pts out of the bounding box should also be penalized
          gradient_error = (pts_mask * gradient_error).sum() / (
                  (pts_mask).sum() + 1e-5)
  
          depth = (mid_z_vals * weights[:, :n_samples]).sum(dim=1, keepdim=True)
          # print("[TEST]: weights_sum in render_core", weights_sum.mean())
          # print("[TEST]: weights_sum in render_core NAN number", weights_sum.isnan().sum())
          # if weights_sum.isnan().sum() > 0:
          #     import ipdb; ipdb.set_trace()
          return {
              'color': color,
              'color_mask': rendering_valid_mask,  # (N_rays, 1)
              'color_mlp': color_mlp,
              'color_mlp_mask': rendering_valid_mask_mlp,
              'sdf': sdf,  # (N_rays, n_samples)
              'depth': depth,  # (N_rays, 1)
              'dists': dists,
              'gradients': gradients.reshape(N_rays, n_samples, 3),
              'variance': 1.0 / inv_variance,
              'mid_z_vals': mid_z_vals,
              'weights': weights,
              'weights_sum': weights_sum,
              'alpha_sum': alpha_sum,
              'alpha_mean': alpha.mean(),
              'cdf': c.reshape(N_rays, n_samples),
              'gradient_error': gradient_error,
              'inside_sphere': inside_sphere,
              'blended_color_patch': blended_color_patch,
              'blended_color_patch_mask': rendering_patch_mask,
              'weights_sum_fg': weights_sum_fg
          }
  
      def render(self, rays_o, rays_d, near, far, sdf_network, rendering_network,
                 perturb_overwrite=-1,
                 background_rgb=None,
                 alpha_inter_ratio=0.0,
                 # * related to conditional feature
                 lod=None,
                 conditional_volume=None,
                 conditional_valid_mask_volume=None,
                 # * 2d feature maps
                 feature_maps=None,
                 color_maps=None,
                 w2cs=None,
                 intrinsics=None,
                 img_wh=None,
                 query_c2w=None,  # -used for testing
                 if_general_rendering=True,
                 if_render_with_grad=True,
                 # * used for blending mlp rendering network
                 img_index=None,
                 rays_uv=None,
                 # * importance sample for second lod network
                 pre_sample=False,  # no use here
                 # * for clear foreground
                 bg_ratio=0.0
                 ):
          device = rays_o.device
          N_rays = len(rays_o)
          # sample_dist = 2.0 / self.n_samples
          sample_dist = ((far - near) / self.n_samples).mean().item()
          z_vals = torch.linspace(0.0, 1.0, self.n_samples).to(device)
          z_vals = near + (far - near) * z_vals[None, :]
  
          bg_num = int(self.n_samples * bg_ratio)
  
          if z_vals.shape[0] == 1:
              z_vals = z_vals.repeat(N_rays, 1)
  
          if bg_num > 0:
              z_vals_bg = z_vals[:, self.n_samples - bg_num:]
              z_vals = z_vals[:, :self.n_samples - bg_num]
  
          n_samples = self.n_samples - bg_num
          perturb = self.perturb
  
          # - significantly speed up training, for the second lod network
          if pre_sample:
              z_vals = self.sample_z_vals_from_maskVolume(rays_o, rays_d, near, far,
                                                          conditional_valid_mask_volume)
  
          if perturb_overwrite >= 0:
              perturb = perturb_overwrite
          if perturb > 0:
              # get intervals between samples
              mids = .5 * (z_vals[..., 1:] + z_vals[..., :-1])
              upper = torch.cat([mids, z_vals[..., -1:]], -1)
              lower = torch.cat([z_vals[..., :1], mids], -1)
              # stratified samples in those intervals
              t_rand = torch.rand(z_vals.shape).to(device)
              z_vals = lower + (upper - lower) * t_rand
  
          background_alpha = None
          background_sampled_color = None
          z_val_before = z_vals.clone()
          # Up sample
          if self.n_importance > 0:
              with torch.no_grad():
                  pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]
  
                  sdf_outputs = sdf_network.sdf(
                      pts.reshape(-1, 3), conditional_volume, lod=lod)
                  # pdb.set_trace()
                  sdf = sdf_outputs['sdf_pts_scale%d' % lod].reshape(N_rays, self.n_samples - bg_num)
  
                  n_steps = 4
                  for i in range(n_steps):
                      new_z_vals = self.up_sample(rays_o, rays_d, z_vals, sdf, self.n_importance // n_steps,
                                                  64 * 2 ** i,
                                                  conditional_valid_mask_volume=conditional_valid_mask_volume,
                                                  )
  
                      # if new_z_vals.isnan().sum() > 0:
                      #     import ipdb; ipdb.set_trace()
  
                      z_vals, sdf = self.cat_z_vals(
                          rays_o, rays_d, z_vals, new_z_vals, sdf, lod,
                          sdf_network, gru_fusion=False,
                          conditional_volume=conditional_volume,
                          conditional_valid_mask_volume=conditional_valid_mask_volume,
                      )
  
                  del sdf
  
              n_samples = self.n_samples + self.n_importance
  
          # Background
          ret_outside = None
  
          # Render
          if bg_num > 0:
              z_vals = torch.cat([z_vals, z_vals_bg], dim=1)
          # if z_vals.isnan().sum() > 0:
          #     import ipdb; ipdb.set_trace()
          ret_fine = self.render_core(rays_o,
                                      rays_d,
                                      z_vals,
                                      sample_dist,
                                      lod,
                                      sdf_network,
                                      rendering_network,
                                      background_rgb=background_rgb,
                                      background_alpha=background_alpha,
                                      background_sampled_color=background_sampled_color,
                                      alpha_inter_ratio=alpha_inter_ratio,
                                      # * related to conditional feature
                                      conditional_volume=conditional_volume,
                                      conditional_valid_mask_volume=conditional_valid_mask_volume,
                                      # * 2d feature maps
                                      feature_maps=feature_maps,
                                      color_maps=color_maps,
                                      w2cs=w2cs,
                                      intrinsics=intrinsics,
                                      img_wh=img_wh,
                                      query_c2w=query_c2w,
                                      if_general_rendering=if_general_rendering,
                                      if_render_with_grad=if_render_with_grad,
                                      # * used for blending mlp rendering network
                                      img_index=img_index,
                                      rays_uv=rays_uv
                                      )
  
          color_fine = ret_fine['color']
  
          if self.n_outside > 0:
              color_fine_mask = torch.logical_or(ret_fine['color_mask'], ret_outside['color_mask'])
          else:
              color_fine_mask = ret_fine['color_mask']
  
          weights = ret_fine['weights']
          weights_sum = ret_fine['weights_sum']
  
          gradients = ret_fine['gradients']
          mid_z_vals = ret_fine['mid_z_vals']
  
          # depth = (mid_z_vals * weights[:, :n_samples]).sum(dim=1, keepdim=True)
          depth = ret_fine['depth']
          depth_varaince = ((mid_z_vals - depth) ** 2 * weights[:, :n_samples]).sum(dim=-1, keepdim=True)
          variance = ret_fine['variance'].reshape(N_rays, n_samples).mean(dim=-1, keepdim=True)
  
          # - randomly sample points from the volume, and maximize the sdf
          pts_random = torch.rand([1024, 3]).float().to(device) * 2 - 1  # normalized to (-1, 1)
          sdf_random = sdf_network.sdf(pts_random, conditional_volume, lod=lod)['sdf_pts_scale%d' % lod]
  
          result = {
              'depth': depth,
              'color_fine': color_fine,
              'color_fine_mask': color_fine_mask,
              'color_outside': ret_outside['color'] if ret_outside is not None else None,
              'color_outside_mask': ret_outside['color_mask'] if ret_outside is not None else None,
              'color_mlp': ret_fine['color_mlp'],
              'color_mlp_mask': ret_fine['color_mlp_mask'],
              'variance': variance.mean(),
              'cdf_fine': ret_fine['cdf'],
              'depth_variance': depth_varaince,
              'weights_sum': weights_sum,
              'weights_max': torch.max(weights, dim=-1, keepdim=True)[0],
              'alpha_sum': ret_fine['alpha_sum'].mean(),
              'alpha_mean': ret_fine['alpha_mean'],
              'gradients': gradients,
              'weights': weights,
              'gradient_error_fine': ret_fine['gradient_error'],
              'inside_sphere': ret_fine['inside_sphere'],
              'sdf': ret_fine['sdf'],
              'sdf_random': sdf_random,
              'blended_color_patch': ret_fine['blended_color_patch'],
              'blended_color_patch_mask': ret_fine['blended_color_patch_mask'],
              'weights_sum_fg': ret_fine['weights_sum_fg']
          }
  
          return result
  
      @torch.no_grad()
      def sample_z_vals_from_sdfVolume(self, rays_o, rays_d, near, far, sdf_volume, mask_volume):
          # ? based on sdf to do importance sampling, seems that too biased on pre-estimation
          device = rays_o.device
          N_rays = len(rays_o)
          n_samples = self.n_samples * 2
  
          z_vals = torch.linspace(0.0, 1.0, n_samples).to(device)
          z_vals = near + (far - near) * z_vals[None, :]
  
          if z_vals.shape[0] == 1:
              z_vals = z_vals.repeat(N_rays, 1)
  
          pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]
  
          sdf = self.get_pts_mask_for_conditional_volume(pts.view(-1, 3), sdf_volume).reshape([N_rays, n_samples])
  
          new_z_vals = self.up_sample(rays_o, rays_d, z_vals, sdf, self.n_samples,
                                      200,
                                      conditional_valid_mask_volume=mask_volume,
                                      )
          return new_z_vals
  
      @torch.no_grad()
      def sample_z_vals_from_maskVolume(self, rays_o, rays_d, near, far, mask_volume):  # don't use
          device = rays_o.device
          N_rays = len(rays_o)
          n_samples = self.n_samples * 2
  
          z_vals = torch.linspace(0.0, 1.0, n_samples).to(device)
          z_vals = near + (far - near) * z_vals[None, :]
  
          if z_vals.shape[0] == 1:
              z_vals = z_vals.repeat(N_rays, 1)
  
          mid_z_vals = (z_vals[:, 1:] + z_vals[:, :-1]) * 0.5
  
          pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None]
  
          pts_mask = self.get_pts_mask_for_conditional_volume(pts.view(-1, 3), mask_volume).reshape(
              [N_rays, n_samples - 1])
  
          # empty voxel set to 0.1, non-empty voxel set to 1
          weights = torch.where(pts_mask > 0, torch.ones_like(pts_mask).to(device),
                                0.1 * torch.ones_like(pts_mask).to(device))
  
          # sample more pts in non-empty voxels
          z_samples = sample_pdf(z_vals, weights, self.n_samples, det=True).detach()
          return z_samples
  
      @torch.no_grad()
      def filter_pts_by_depthmaps(self, coords, pred_depth_maps, proj_matrices,
                                  partial_vol_origin, voxel_size,
                                  near, far, depth_interval, d_plane_nums):
          """
          Use the pred_depthmaps to remove redundant pts (pruned by sdf, sdf always have two sides, the back side is useless)
          :param coords: [n, 3]  int coords
          :param pred_depth_maps: [N_views, 1, h, w]
          :param proj_matrices: [N_views, 4, 4]
          :param partial_vol_origin: [3]
          :param voxel_size: 1
          :param near: 1
          :param far: 1
          :param depth_interval: 1
          :param d_plane_nums: 1
          :return:
          """
          device = pred_depth_maps.device
          n_views, _, sizeH, sizeW = pred_depth_maps.shape
  
          if len(partial_vol_origin.shape) == 1:
              partial_vol_origin = partial_vol_origin[None, :]
          pts = coords * voxel_size + partial_vol_origin
  
          rs_grid = pts.unsqueeze(0).expand(n_views, -1, -1)
          rs_grid = rs_grid.permute(0, 2, 1).contiguous()  # [n_views, 3, n_pts]
          nV = rs_grid.shape[-1]
          rs_grid = torch.cat([rs_grid, torch.ones([n_views, 1, nV]).to(device)], dim=1)  # [n_views, 4, n_pts]
  
          # Project grid
          im_p = proj_matrices @ rs_grid  # - transform world pts to image UV space   # [n_views, 4, n_pts]
          im_x, im_y, im_z = im_p[:, 0], im_p[:, 1], im_p[:, 2]
          im_x = im_x / im_z
          im_y = im_y / im_z
  
          im_grid = torch.stack([2 * im_x / (sizeW - 1) - 1, 2 * im_y / (sizeH - 1) - 1], dim=-1)
  
          im_grid = im_grid.view(n_views, 1, -1, 2)
          sampled_depths = torch.nn.functional.grid_sample(pred_depth_maps, im_grid, mode='bilinear',
                                                           padding_mode='zeros',
                                                           align_corners=True)[:, 0, 0, :]  # [n_views, n_pts]
          sampled_depths_valid = (sampled_depths > 0.5 * near).float()
          valid_d_min = (sampled_depths - d_plane_nums * depth_interval).clamp(near.item(),
                                                                               far.item()) * sampled_depths_valid
          valid_d_max = (sampled_depths + d_plane_nums * depth_interval).clamp(near.item(),
                                                                               far.item()) * sampled_depths_valid
  
          mask = im_grid.abs() <= 1
  