code: def render_hybrid_nerf(
      net,
      cap,
      posed_verts,
      faces,
      Ts,
      rays_per_batch=32768,
      samples_per_ray=64,
      importance_samples_per_ray=128,
      white_bkg=True,
      geo_threshold=DEFAULT_GEO_THRESH,
      return_depth=False,
  ):
      device = next(net.parameters()).device
  
      def build_batch(origins, dirs, near, far):
          if isinstance(origins, torch.Tensor):
              ray_batch = {
                  "origin": (origins).float().to(device),
                  "direction": (dirs).float().to(device),
                  "near": (near[..., None]).float().to(device),
                  "far": (far[..., None]).float().to(device),
              }
          else:
              ray_batch = {
                  "origin": torch.from_numpy(origins)
                  .float()
                  .to(device),
                  "direction": torch.from_numpy(dirs)
                  .float()
                  .to(device),
                  "near": torch.from_numpy(near[..., None])
                  .float()
                  .to(device),
                  "far": torch.from_numpy(far[..., None])
                  .float()
                  .to(device),
              }
          return ray_batch
  
      with torch.set_grad_enabled(False):
          coords = np.argwhere(np.ones(cap.shape))[:, ::-1]
          origins, dirs = ray_utils.shot_rays(cap, coords)
          total_rays = origins.shape[0]
          total_rgb_map = []
          total_depth_map = []
          total_acc_map = []
          for i in range(0, total_rays, rays_per_batch):
              print(f"{i} / {total_rays}")
              rgb_map = np.zeros_like(origins[i : i + rays_per_batch])
              depth_map = np.zeros_like(
                  origins[i : i + rays_per_batch, 0]
              )
              acc_map = np.zeros_like(
                  origins[i : i + rays_per_batch, 0]
              )
              bkg_ray_batch = build_batch(
                  origins[i : i + rays_per_batch],
                  dirs[i : i + rays_per_batch],
                  np.array(
                      [cap.near["bkg"]]
                      * origins[i : i + rays_per_batch].shape[0]
                  ),
                  np.array(
                      [cap.far["bkg"]]
                      * origins[i : i + rays_per_batch].shape[0]
                  ),
              )
              bkg_pts, bkg_dirs, bkg_z_vals = ray_utils.ray_to_samples(
                  bkg_ray_batch, samples_per_ray, device=device
              )
              bkg_out = net.coarse_bkg_net(bkg_pts, bkg_dirs)
              if net.fine_bkg_net is not None:
                  _, _, _, bkg_weights, _ = raw2outputs(
                      bkg_out,
                      bkg_z_vals,
                      bkg_dirs[:, 0, :],
                      white_bkg=white_bkg,
                  )
                  (
                      bkg_pts,
                      bkg_dirs,
                      bkg_z_vals,
                  ) = ray_utils.ray_to_importance_samples(
                      bkg_ray_batch,
                      bkg_z_vals,
                      bkg_weights,
                      importance_samples_per_ray,
                      device=device,
                  )
                  bkg_out = net.fine_bkg_net(bkg_pts, bkg_dirs)
              temp_near, temp_far = ray_utils.geometry_guided_near_far(
                  origins[i : i + rays_per_batch],
                  dirs[i : i + rays_per_batch],
                  posed_verts,
                  geo_threshold,
              )
              if (temp_near >= temp_far).any():
                  # no fuse
                  # render bkg colors
                  (
                      coarse_bkg_rgb_map,
                      _,
                      coarse_bkg_acc_map,
                      weights,
                      coarse_bkg_depth_map,
                  ) = raw2outputs(
                      bkg_out[temp_near >= temp_far],
                      bkg_z_vals[temp_near >= temp_far],
                      bkg_dirs[temp_near >= temp_far][:, 0, :],
                      white_bkg=white_bkg,
                  )
                  rgb_map[temp_near >= temp_far] = (
                      coarse_bkg_rgb_map.detach().cpu().numpy()
                  )
                  depth_map[temp_near >= temp_far] = (
                      coarse_bkg_depth_map.detach().cpu().numpy()
                  )
                  acc_map[temp_near >= temp_far] = 0
  
              if (temp_near < temp_far).any():
                  human_ray_batch = build_batch(
                      origins[i : i + rays_per_batch][
                          temp_near < temp_far
                      ],
                      dirs[i : i + rays_per_batch][
                          temp_near < temp_far
                      ],
                      temp_near[temp_near < temp_far],
                      temp_far[temp_near < temp_far],
                  )
                  (
                      human_pts,
                      human_dirs,
                      human_z_vals,
                  ) = ray_utils.ray_to_samples(
                      human_ray_batch, samples_per_ray, device=device
                  )
                  (
                      can_pts,
                      can_dirs,
                      _,
                  ) = ray_utils.warp_samples_to_canonical(
                      human_pts.cpu().numpy(), posed_verts, faces, Ts
                  )
                  can_pts = torch.from_numpy(can_pts).to(device).float()
                  can_dirs = (
                      torch.from_numpy(can_dirs).to(device).float()
                  )
                  human_out = net.coarse_human_net(can_pts, can_dirs)
                  coarse_total_zvals, coarse_order = torch.sort(
                      torch.cat(
                          [
                              bkg_z_vals[temp_near < temp_far],
                              human_z_vals,
                          ],
                          -1,
                      ),
                      -1,
                  )
                  coarse_total_out = torch.cat(
                      [bkg_out[temp_near < temp_far], human_out], 1
                  )
                  _b, _n, _c = coarse_total_out.shape
                  coarse_total_out = coarse_total_out[
                      torch.arange(_b).view(_b, 1, 1).repeat(1, _n, _c),
                      coarse_order.view(_b, _n, 1).repeat(1, 1, _c),
                      torch.arange(_c).view(1, 1, _c).repeat(_b, _n, 1),
                  ]
                  human_rgb_map, _, _, _, human_depth_map = raw2outputs(
                      coarse_total_out,
                      coarse_total_zvals,
                      human_dirs[:, 0, :],
                      white_bkg=white_bkg,
                  )
  
                  _, _, human_acc_map, _, _ = raw2outputs(
                      human_out,
                      human_z_vals,
                      human_dirs[:, 0, :],
                      white_bkg=white_bkg,
                  )
                  rgb_map[temp_near < temp_far] = (
                      human_rgb_map.detach().cpu().numpy()
                  )
                  depth_map[temp_near < temp_far] = (
                      human_depth_map.detach().cpu().numpy()
                  )
                  acc_map[temp_near < temp_far] = (
                      human_acc_map.detach().cpu().numpy()
                  )
              total_rgb_map.append(rgb_map)
              total_depth_map.append(depth_map)
              total_acc_map.append(acc_map)
          total_rgb_map = np.concatenate(total_rgb_map).reshape(
              *cap.shape, -1
          )
          total_depth_map = np.concatenate(total_depth_map).reshape(
              *cap.shape
          )
          total_acc_map = np.concatenate(total_acc_map).reshape(
              *cap.shape
          )
      if return_depth:
          return total_rgb_map, total_depth_map
      return total_rgb_map
  
file_path: ml-neuman/utils/render_utils.py
imports: import random
  import numpy
  import torch
  import torch.nn.functional
  from pytorch3d.structures import Meshes
  from pytorch3d.renderer import RasterizationSettings, MeshRenderer, MeshRasterizer, HardPhongShader, PointLights, TexturesVertex, PerspectiveCameras
  from PIL import Image
  from cameras.camera_pose import CameraPose
  from utils import ray_utils
  from utils.constant import DEFAULT_GEO_THRESH
  from geometry import transformations
label_desc: Use image node_id 6 for calling a global function with example usage: render_hybrid_nerf(net, cap, posed_verts, faces, Ts, rays_per_batch, samples_per_ray, importance_samples_per_ray, white_bkg, geo_threshold, return_depth) and returns: total_rgb_map, ray_batch, total_rgb_map,  total_depth_map