# This implementation is from NeRF

# bins: [B, T], old_z_vals

# weights: [B, T - 1], bin weights.

# return: [B, n_samples], new_z_vals

# Get pdf

weights = weights + 1e-5 # prevent nans

pdf = weights / jt.sum(weights, -1, keepdim=True)

cdf = jt.cumsum(pdf, -1)

cdf = jt.concat([jt.zeros_like(cdf[..., :1]), cdf], -1)

# Take uniform samples

if det:

u = jt.linspace(0. + 0.5 / n_samples, 1. - 0.5 / n_samples, steps=n_samples)

u = u.expand(list(cdf.shape[:-1]) + [n_samples])

else:

u = jt.rand(list(cdf.shape[:-1]) + [n_samples])

# Invert CDF

u = u.contiguous()

inds = jt.searchsorted(cdf, u, right=True)

# FIXME: check mix and max correctness

below = jt.max(jt.array([jt.zeros_like(inds - 1), inds - 1]), dim=0)

above = jt.min(jt.array([(cdf.shape[-1] - 1) * jt.ones_like(inds), inds]), dim=0)

inds_g = jt.stack([below, above], -1) # (B, n_samples, 2)

matched_shape = [inds_g.shape[0], inds_g.shape[1], cdf.shape[-1]]

cdf_g = jt.gather(cdf.unsqueeze(1).expand(matched_shape), 2, inds_g)

bins_g = jt.gather(bins.unsqueeze(1).expand(matched_shape), 2, inds_g)

denom = (cdf_g[..., 1] - cdf_g[..., 0])

denom = jt.where(denom < 1e-5, jt.ones_like(denom), denom)

t = (u - cdf_g[..., 0]) / denom

samples = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0])

# rays: [B, N, 3], [B, N, 3]

# bound: int, radius for ball or half-edge-length for cube

# return near [B, N, 1], far [B, N, 1]

radius = rays_o.norm(dim=-1, keepdim=True)

if type == 'sphere':

near = radius - bound # [B, N, 1]

far = radius + bound

elif type == 'cube':

tmin = (-bound - rays_o) / (rays_d + 1e-15) # [B, N, 3]

tmax = (bound - rays_o) / (rays_d + 1e-15)

near = jt.where(tmin < tmax, tmin, tmax).max(dim=-1, keepdim=True)[0]

far = jt.where(tmin > tmax, tmin, tmax).min(dim=-1, keepdim=True)[0]

# if far < near, means no intersection, set both near and far to inf (1e9 here)

mask = far < near

near[mask] = 1e9

far[mask] = 1e9

# restrict near to a minimal value

near = jt.clamp(near, min=min_near)

# pc: [N, 3]

# color: [N, 3/4]

print('[visualize points]', pc.shape, pc.dtype, pc.min(0), pc.max(0))

pc = trimesh.PointCloud(pc, color)

# axis

axes = trimesh.creation.axis(axis_length=4)

# sphere

sphere = trimesh.creation.icosphere(radius=1)

pcd = o3d.geometry.PointCloud()

pcd.points = o3d.utility.Vector3dVector(points)

if rgb is not None:

pcd.colors = o3d.utility.Vector3dVector(rgb)

def __init__(self, opt):

super().__init__()

self.opt = opt

self.bound = opt.bound

self.cascade = 1 + math.ceil(math.log2(opt.bound))

self.grid_size = 128

self.cuda_ray = opt.cuda_ray

self.min_near = opt.min_near

self.density_thresh = opt.density_thresh

self.bg_radius = opt.bg_radius

# prepare aabb with a 6D tensor (xmin, ymin, zmin, xmax, ymax, zmax)

# NOTE: aabb (can be rectangular) is only used to generate points, we still rely on bound (always cubic) to calculate density grid and hashing.

self.aabb_train = jt.array([-opt.bound, -opt.bound, -opt.bound, opt.bound, opt.bound, opt.bound], dtype=jt.float32)

self.aabb_infer = self.aabb_train.clone()

# self.register_buffer('aabb_train', aabb_train)

# self.register_buffer('aabb_infer', aabb_infer)

# extra state for cuda raymarching

if self.cuda_ray:

# density grid

self.density_grid = jt.zeros([self.cascade, self.grid_size ** 3]) # [CAS, H * H * H]

self.density_bitfield = jt.zeros(self.cascade * self.grid_size ** 3 // 8, dtype=jt.uint8) # [CAS * H * H * H // 8]

# self.register_buffer('density_grid', density_grid)

# self.register_buffer('density_bitfield', density_bitfield)

self.mean_density = 0

self.iter_density = 0

# step counter

self.step_counter = jt.zeros(16, 2, dtype=jt.int32) # 16 is hardcoded for averaging...

#self.register_buffer('step_counter', step_counter)

self.mean_count = 0

self.local_step = 0

def execute(self, *args, **kwargs):

raise NotImplementedError()

def density(self, x):

raise NotImplementedError()

def color(self, x, d, mask=None, **kwargs):

raise NotImplementedError()

def reset_extra_state(self):

if not self.cuda_ray:

return

# density grid

self.density_grid.zero_()

self.mean_density = 0

self.iter_density = 0

# step counter

self.step_counter.zero_()

self.mean_count = 0

self.local_step = 0

@jt.no_grad()

def export_mesh(self, path, resolution=None, S=128):

if resolution is None:

resolution = self.grid_size

if self.cuda_ray:

density_thresh = min(self.mean_density, self.density_thresh)

else:

density_thresh = self.density_thresh

sigmas = np.zeros([resolution, resolution, resolution], dtype=np.float32)

# query

X = jt.linspace(-1, 1, resolution).split(S)

Y = jt.linspace(-1, 1, resolution).split(S)

Z = jt.linspace(-1, 1, resolution).split(S)

for xi, xs in enumerate(X):

for yi, ys in enumerate(Y):

for zi, zs in enumerate(Z):

xx, yy, zz = custom_meshgrid(xs, ys, zs)

pts = jt.concat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1) # [S, 3]

val = self.density(pts.to(self.aabb_train.device), False)

sigmas[xi * S: xi * S + len(xs), yi * S: yi * S + len(ys), zi * S: zi * S + len(zs)] = val['sigma'].reshape(len(xs), len(ys), len(zs)).detach().cpu().numpy() # [S, 1] --> [x, y, z]

vertices, triangles = mcubes.marching_cubes(sigmas, density_thresh)

vertices = vertices / (resolution - 1.0) * 2 - 1

vertices = vertices.astype(np.float32)

triangles = triangles.astype(np.int32)

v = jt.array(vertices)

f = jt.array(triangles)

def _export(v, f, h0=2048, w0=2048, ssaa=1, name=''):

# v, f: torch Tensor

v_np = v.cpu().numpy() # [N, 3]

f_np = f.cpu().numpy() # [M, 3]

print(f'[INFO] running xatlas to unwrap UVs for mesh: v={v_np.shape} f={f_np.shape}')

# unwrap uvs

import xatlas

import nvdiffrast.torch as dr

from sklearn.neighbors import NearestNeighbors

from scipy.ndimage import binary_dilation, binary_erosion

glctx = dr.RasterizeCudaContext()

atlas = xatlas.Atlas()

atlas.add_mesh(v_np, f_np)

chart_options = xatlas.ChartOptions()

chart_options.max_iterations = 0 # disable merge_chart for faster unwrap...

atlas.generate(chart_options=chart_options)

vmapping, ft_np, vt_np = atlas[0] # [N], [M, 3], [N, 2]

# vmapping, ft_np, vt_np = xatlas.parametrize(v_np, f_np) # [N], [M, 3], [N, 2]

vt = jt.array(vt_np.astype(np.float32)).float()

ft = jt.array(ft_np.astype(np.int64)).int()

# render uv maps

uv = vt * 2.0 - 1.0 # uvs to range [-1, 1]

uv = jt.concat((uv, jt.zeros_like(uv[..., :1]), jt.ones_like(uv[..., :1])), dim=-1) # [N, 4]

if ssaa > 1:

h = int(h0 * ssaa)

w = int(w0 * ssaa)

else:

h, w = h0, w0

rast, _ = dr.rasterize(glctx, uv.unsqueeze(0), ft, (h, w)) # [1, h, w, 4]

xyzs, _ = dr.interpolate(v.unsqueeze(0), rast, f) # [1, h, w, 3]

mask, _ = dr.interpolate(jt.ones_like(v[:, :1]).unsqueeze(0), rast, f) # [1, h, w, 1]

# masked query

xyzs = xyzs.view(-1, 3)

mask = (mask > 0).view(-1)

sigmas = jt.zeros(h * w, dtype=jt.float32)

feats = jt.zeros(h * w, 3, dtype=jt.float32)

if mask.any():

xyzs = xyzs[mask] # [M, 3]

# batched inference to avoid OOM

all_sigmas = []

all_feats = []

head = 0

while head < xyzs.shape[0]:

tail = min(head + 640000, xyzs.shape[0])

results_ = self.density(xyzs[head:tail])

all_sigmas.append(results_['sigma'].float())

all_feats.append(results_['albedo'].float())

head += 640000

sigmas[mask] = jt.concat(all_sigmas, dim=0)

feats[mask] = jt.concat(all_feats, dim=0)

sigmas = sigmas.view(h, w, 1)

feats = feats.view(h, w, -1)

mask = mask.view(h, w)

### alpha mask

# quantize [0.0, 1.0] to [0, 255]

feats = feats.cpu().numpy()

feats = (feats * 255).astype(np.uint8)

# alphas = alphas.cpu().numpy()

# alphas = (alphas * 255).astype(np.uint8)

### NN search as an antialiasing ...

mask = mask.cpu().numpy()

inpaint_region = binary_dilation(mask, iterations=3)

inpaint_region[mask] = 0

search_region = mask.copy()

not_search_region = binary_erosion(search_region, iterations=2)

search_region[not_search_region] = 0

search_coords = np.stack(np.nonzero(search_region), axis=-1)

inpaint_coords = np.stack(np.nonzero(inpaint_region), axis=-1)

knn = NearestNeighbors(n_neighbors=1, algorithm='kd_tree').fit(search_coords)

_, indices = knn.kneighbors(inpaint_coords)

feats[tuple(inpaint_coords.T)] = feats[tuple(search_coords[indices[:, 0]].T)]

# do ssaa after the NN search, in numpy

feats = cv2.cvtColor(feats, cv2.COLOR_RGB2BGR)

if ssaa > 1:

# alphas = cv2.resize(alphas, (w0, h0), interpolation=cv2.INTER_NEAREST)

feats = cv2.resize(feats, (w0, h0), interpolation=cv2.INTER_LINEAR)

# cv2.imwrite(os.path.join(path, f'alpha.png'), alphas)

cv2.imwrite(os.path.join(path, f'{name}albedo.png'), feats)

# save obj (v, vt, f /)

obj_file = os.path.join(path, f'{name}mesh.obj')

mtl_file = os.path.join(path, f'{name}mesh.mtl')

print(f'[INFO] writing obj mesh to {obj_file}')

with open(obj_file, "w") as fp:

fp.write(f'mtllib {name}mesh.mtl \n')

print(f'[INFO] writing vertices {v_np.shape}')

for v in v_np:

fp.write(f'v {v[0]} {v[1]} {v[2]} \n')

print(f'[INFO] writing vertices texture coords {vt_np.shape}')

for v in vt_np:

fp.write(f'vt {v[0]} {1 - v[1]} \n')

print(f'[INFO] writing faces {f_np.shape}')

fp.write(f'usemtl mat0 \n')

for i in range(len(f_np)):

fp.write(f"f {f_np[i, 0] + 1}/{ft_np[i, 0] + 1} {f_np[i, 1] + 1}/{ft_np[i, 1] + 1} {f_np[i, 2] + 1}/{ft_np[i, 2] + 1} \n")

with open(mtl_file, "w") as fp:

fp.write(f'newmtl mat0 \n')

fp.write(f'Ka 1.000000 1.000000 1.000000 \n')

fp.write(f'Kd 1.000000 1.000000 1.000000 \n')

fp.write(f'Ks 0.000000 0.000000 0.000000 \n')

fp.write(f'Tr 1.000000 \n')

fp.write(f'illum 1 \n')

fp.write(f'Ns 0.000000 \n')

fp.write(f'map_Kd {name}albedo.png \n')

_export(v, f)

def run(self, rays_o, rays_d, ref_bg=None, num_steps=128, upsample_steps=128, light_d=None, ambient_ratio=1.0, shading='albedo', bg_color=None, perturb=False, **kwargs):

# rays_o, rays_d: [B, N, 3], assumes B == 1

# bg_color: [BN, 3] in range [0, 1]

# return: image: [B, N, 3], depth: [B, N]

prefix = rays_o.shape[:-1]

rays_o = rays_o.contiguous().view(-1, 3)

rays_d = rays_d.contiguous().view(-1, 3)

N = rays_o.shape[0] # N = B * N, in fact

results = {}

# choose aabb

aabb = self.aabb_train if self.training else self.aabb_infer

# sample steps

nears, fars = near_far_from_bound(rays_o, rays_d, self.bound, type='sphere', min_near=self.min_near)

# random sample light_d if not provided

if light_d is None:

# gaussian noise around the ray origin, so the light always face the view dir (avoid dark face)

light_d = (rays_o[0] + jt.randn(3, dtype=jt.float32))

light_d = safe_normalize(light_d)

z_vals = jt.linspace(0.0, 1.0, num_steps).unsqueeze(0) # [1, T]

z_vals = z_vals.expand((N, num_steps)) # [N, T]

z_vals = nears + (fars - nears) * z_vals # [N, T], in [nears, fars]

# perturb z_vals

sample_dist = (fars - nears) / num_steps

if perturb:

z_vals = z_vals + (jt.rand(z_vals.shape) - 0.5) * sample_dist

# generate xyzs

xyzs = rays_o.unsqueeze(-2) + rays_d.unsqueeze(-2) * z_vals.unsqueeze(-1) # [N, 1, 3] * [N, T, 1] -> [N, T, 3]

xyzs = jt.min(jt.array([jt.max(jt.array([xyzs, aabb[:3]]), dim=0), aabb[3:]]), dim=0) # a manual clip.

# xyzs=jt.clamp(xyzs, -1, 1)

#plot_pointcloud(xyzs.reshape(-1, 3).detach().cpu().numpy())

# query SDF and RGB

density_outputs = self.density(xyzs.reshape(-1, 3))

#sigmas = density_outputs['sigma'].view(N, num_steps) # [N, T]

for k, v in density_outputs.items():

density_outputs[k] = v.view(N, num_steps, -1)

# upsample z_vals (nerf-like)

if upsample_steps > 0:

with jt.no_grad():

deltas = z_vals[..., 1:] - z_vals[..., :-1] # [N, T-1]

deltas = jt.concat([deltas, sample_dist * jt.ones_like(deltas[..., :1])], dim=-1)

alphas = 1 - jt.exp(-deltas * density_outputs['sigma'].squeeze(-1)) # [N, T]

alphas_shifted = jt.concat([jt.ones_like(alphas[..., :1]), 1 - alphas + 1e-15], dim=-1) # [N, T+1]

weights = alphas * jt.cumprod(alphas_shifted, dim=-1)[..., :-1] # [N, T]

# sample new z_vals

z_vals_mid = (z_vals[..., :-1] + 0.5 * deltas[..., :-1]) # [N, T-1]

new_z_vals = sample_pdf(z_vals_mid, weights[:, 1:-1], upsample_steps, det=not self.training).detach() # [N, t]

new_xyzs = rays_o.unsqueeze(-2) + rays_d.unsqueeze(-2) * new_z_vals.unsqueeze(-1) # [N, 1, 3] * [N, t, 1] -> [N, t, 3]

# FIXME: need to check correstness

new_xyzs = jt.min(jt.array([jt.max(jt.array([new_xyzs, aabb[:3]]), dim=0), aabb[3:]]), dim=0) # a manual clip.

# new_xyzs = jt.clamp(new_xyzs, -1, 1)

# only forward new points to save computation

new_density_outputs = self.density(new_xyzs.reshape(-1, 3))

#new_sigmas = new_density_outputs['sigma'].view(N, upsample_steps) # [N, t]

for k, v in new_density_outputs.items():

new_density_outputs[k] = v.view(N, upsample_steps, -1)

# re-order

z_vals = jt.concat([z_vals, new_z_vals], dim=1) # [N, T+t]

z_vals, z_index = jt.sort(z_vals, dim=1)

xyzs = jt.concat([xyzs, new_xyzs], dim=1) # [N, T+t, 3]

xyzs = jt.gather(xyzs, dim=1, index=z_index.unsqueeze(-1).expand_as(xyzs))

for k in density_outputs:

tmp_output = jt.concat([density_outputs[k], new_density_outputs[k]], dim=1)

density_outputs[k] = jt.gather(tmp_output, dim=1, index=z_index.unsqueeze(-1).expand_as(tmp_output))

deltas = z_vals[..., 1:] - z_vals[..., :-1] # [N, T+t-1]

deltas = jt.concat([deltas, sample_dist * jt.ones_like(deltas[..., :1])], dim=-1)

alphas = 1 - jt.exp(-deltas * density_outputs['sigma'].squeeze(-1)) # [N, T+t]

alphas_shifted = jt.concat([jt.ones_like(alphas[..., :1]), 1 - alphas + 1e-15], dim=-1) # [N, T+t+1]

weights = alphas * jt.cumprod(alphas_shifted, dim=-1)[..., :-1] # [N, T+t]

dirs = rays_d.view(-1, 1, 3).expand_as(xyzs)

for k, v in density_outputs.items():

density_outputs[k] = v.view(-1, v.shape[-1])

sigmas, rgbs, normals = self(xyzs.reshape(-1, 3), dirs.reshape(-1, 3), light_d, ratio=ambient_ratio, shading=shading)

rgbs = rgbs.view(N, -1, 3) # [N, T+t, 3]

if normals is not None:

# calculate normal

normal_map = normals.reshape(N, -1, 3) # [N, T, 3]

normal_map = jt.sum(normal_map * weights[:, :, None], dim=1)

# orientation loss

normals = normals.view(N, -1, 3)

loss_orient = weights.detach() * (normals * dirs).sum(-1).clamp(min=0) ** 2

results['loss_orient'] = loss_orient.sum(-1).mean()

# surface normal smoothness

if self.opt.lambda_smooth > 0:

normals_perturb = self.normal(xyzs.reshape(-1, 3) + jt.randn_like(xyzs).reshape(-1, 3) * 1e-2).view(N, -1, 3)

loss_smooth = (normals - normals_perturb).abs()

results['loss_smooth'] = loss_smooth.mean()

# calculate weight_sum (mask)

weights_sum = weights.sum(dim=-1) # [N]

# calculate depth

depth = jt.sum(weights * z_vals, dim=-1)

# calculate color

image = jt.sum(weights.unsqueeze(-1) * rgbs, dim=-2) # [N, 3], in [0, 1]

# mix background color

if self.bg_radius > 0:

bg_color = self.background(rays_d.reshape(-1, 3)) # [N, 3]

elif bg_color is None:

bg_color = 1

fg_image = image

bg_image = bg_color.view(*prefix, 3)

fg_image = fg_image.view(*prefix, 3)

image = image + (1 - weights_sum).unsqueeze(-1) * bg_color

image = image.view(*prefix, 3)

depth = depth.view(*prefix, 1)

mask = (nears < fars).reshape(*prefix)

results['image'] = image

results['depth'] = depth

results['weights_sum'] = weights_sum

results['mask'] = mask

results['normal'] = normal_map

if self.bg_radius > 0:

results['bg'] = bg_image

return results

def run_cuda(self, rays_o, rays_d, depth_scale=None, bg_color=None, dt_gamma=0, light_d=None, ambient_ratio=1.0, shading='albedo', perturb=False, force_all_rays=False, max_steps=1024, T_thresh=1e-4, **kwargs):

# rays_o, rays_d: [B, N, 3], assumes B == 1

# return: image: [B, N, 3], depth: [B, N]

B = rays_o.shape[0]

prefix = rays_o.shape[:-1]

rays_o = rays_o.contiguous().view(-1, 3)

rays_d = rays_d.contiguous().view(-1, 3)

N = rays_o.shape[0] # N = B * N, in fact

# pre-calculate near far

nears, fars = raymarching.near_far_from_aabb(rays_o, rays_d, self.aabb_train if self.training else self.aabb_infer)

# random sample light_d if not provided

if light_d is None:

# gaussian noise around the ray origin, so the light always face the view dir (avoid dark face)

light_d = (rays_o[0] + jt.randn(3, dtype=jt.float32))

light_d = safe_normalize(light_d)

results = {}

if self.training:

# setup counter

counter = self.step_counter[self.local_step % 16]

counter.zero_() # set to 0

xyzs, dirs, deltas, rays = raymarching.march_rays_train(rays_o, rays_d, self.bound, self.density_bitfield, self.cascade, self.grid_size, nears, fars, counter, self.mean_count, perturb, 128, force_all_rays, dt_gamma, max_steps)

self.step_counter[self.local_step % 16] = counter

self.local_step += 1

sigmas, rgbs, normals = self(xyzs, dirs, light_d, ratio=ambient_ratio, shading=shading)

weights_sum, depth, image = raymarching.composite_rays_train(sigmas, rgbs, deltas, rays, T_thresh)

# normals related regularizations

if normals is not None:

# orientation loss (not very exact in cuda ray mode)

weights = 1 - jt.exp(-sigmas)

loss_orient = weights.detach() * (normals * dirs).sum(-1).clamp(min_v=0) ** 2

results['loss_orient'] = loss_orient.mean()

# surface normal smoothness

if self.opt.lambda_smooth > 0:

normals_perturb = self.normal(xyzs + jt.randn_like(xyzs) * 1e-2)

loss_smooth = (normals - normals_perturb).abs()

results['loss_smooth'] = loss_smooth.mean()

else:

# allocate outputs

dtype = jt.float32

weights_sum = jt.zeros(N, dtype=dtype)

depth = jt.zeros(N, dtype=dtype)

image = jt.zeros(N, 3, dtype=dtype)

normal = jt.zeros(N, 3, dtype=dtype)

n_alive = N

rays_alive = jt.arange(n_alive, dtype=jt.int32) # [N]

rays_t = nears.clone() # [N]

step = 0

while step < max_steps: # hard coded max step

# count alive rays

n_alive = rays_alive.shape[0]

# exit loop

if n_alive <= 0:

break

# decide compact_steps

n_step = max(min(N // n_alive, 8), 1)

xyzs, dirs, deltas = raymarching.march_rays(n_alive, n_step, rays_alive, rays_t, rays_o, rays_d, self.bound, self.density_bitfield, self.cascade, self.grid_size, nears, fars, 128, perturb if step == 0 else False, dt_gamma, max_steps)

sigmas, rgbs, normals = self(xyzs, dirs, light_d, is_grad=False, ratio=ambient_ratio, shading=shading) # syh: 问题莫不是在这儿？

normals = (normals + 1) / 2

raymarching.composite_rays(n_alive, n_step, rays_alive, rays_t, sigmas, rgbs, normals, deltas, weights_sum, depth, image, normal, T_thresh)

rays_alive=rays_alive[rays_alive>=0]

step += n_step

# mix background color

if bg_color is None:

bg_color = 1

image = image + (1 - weights_sum).unsqueeze(-1) * bg_color

image = image.view(*prefix, 3)

if not self.training:

normal = normal + (1 - weights_sum).unsqueeze(-1) * bg_color

normal = normal.view(*prefix, 3)

bg_depth = self.opt.max_depth

depth = depth + (1 - weights_sum) * bg_depth

if depth_scale is not None:

depth = depth.view(*prefix, 1) * depth_scale.view(*prefix, 1)

else:

depth = depth.view(*prefix, 1)

weights_sum = weights_sum.reshape(*prefix)

mask = (nears < fars).reshape(*prefix)

results['image'] = image

results['depth'] = depth

results['weights_sum'] = weights_sum

results['mask'] = mask

if not self.training:

results['normal'] = normal

return results

@jt.no_grad()

def update_extra_state(self, decay=0.95, S=128):

# call before each epoch to update extra states.

if not self.cuda_ray:

return

### update density grid

tmp_grid = - jt.ones_like(self.density_grid)

X = jt.arange(self.grid_size, dtype=jt.int32).split(S)

Y = jt.arange(self.grid_size, dtype=jt.int32).split(S)

Z = jt.arange(self.grid_size, dtype=jt.int32).split(S)

for xs in X:

for ys in Y:

for zs in Z:

# construct points

xx, yy, zz = custom_meshgrid(xs, ys, zs)

coords = jt.concat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1) # [N, 3], in [0, 128)

indices = raymarching.morton3D(coords) # [N]

xyzs = 2 * coords.float() / (self.grid_size - 1) - 1 # [N, 3] in [-1, 1]

# cascading

for cas in range(self.cascade):

bound = min(2 ** cas, self.bound)

half_grid_size = bound / self.grid_size

# scale to current cascade's resolution

cas_xyzs = xyzs * (bound - half_grid_size)

# add noise in [-hgs, hgs]

cas_xyzs += (jt.rand_like(cas_xyzs) * 2 - 1) * half_grid_size

# query density

sigmas = self.density(cas_xyzs, False)['sigma'].reshape(-1).detach()

# assign

tmp_grid[cas, indices] = sigmas

# ema update

valid_mask = self.density_grid >= 0

self.density_grid[valid_mask] = jt.maximum(self.density_grid[valid_mask] * decay, tmp_grid[valid_mask])

density_grid = self.density_grid[valid_mask].numpy()

self.mean_density = density_grid.mean()

self.iter_density += 1

# convert to bitfield

density_thresh = min(self.mean_density, self.density_thresh)

self.density_bitfield = raymarching.packbits(self.density_grid, density_thresh, self.density_bitfield)

### update step counter

total_step = min(16, self.local_step)

if total_step > 0:

self.mean_count = int(self.step_counter[:total_step, 0].sum().item() / total_step)

self.local_step = 0

def render(self, rays_o, rays_d, depth_scale=None, staged=False, max_ray_batch=4096, **kwargs):

# rays_o, rays_d: [B, N, 3], assumes B == 1

# return: pred_rgb: [B, N, 3]

if self.cuda_ray:

_run = self.run_cuda

else:

_run = self.run

B, N = rays_o.shape[:2]

# never stage when cuda_ray

if staged and not self.cuda_ray:

depth = jt.empty((B, N, 1))

image = jt.empty((B, N, 3))

weights_sum = jt.empty((B, N))

for b in range(B):

head = 0

while head < N:

tail = min(head + max_ray_batch, N)

results_ = _run(rays_o[b:b+1, head:tail], rays_d[b:b+1, head:tail], **kwargs)

depth[b:b+1, head:tail] = results_['depth']

weights_sum[b:b+1, head:tail] = results_['weights_sum']

image[b:b+1, head:tail] = results_['image']

head += max_ray_batch

results = {}

results['depth'] = depth

results['image'] = image

results['weights_sum'] = weights_sum

else:

results = _run(rays_o, rays_d, depth_scale, **kwargs)