432 lines
19 KiB
Python
432 lines
19 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
import mcubes
|
|
|
|
def extract_fields(bound_min, bound_max, resolution, query_func):
|
|
N = 64
|
|
X = torch.linspace(bound_min[0], bound_max[0], resolution).split(N)
|
|
Y = torch.linspace(bound_min[1], bound_max[1], resolution).split(N)
|
|
Z = torch.linspace(bound_min[2], bound_max[2], resolution).split(N)
|
|
|
|
u = np.zeros([resolution, resolution, resolution], dtype=np.float32)
|
|
with torch.no_grad():
|
|
for xi, xs in enumerate(X):
|
|
for yi, ys in enumerate(Y):
|
|
for zi, zs in enumerate(Z):
|
|
xx, yy, zz = torch.meshgrid(xs, ys, zs)
|
|
pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1)
|
|
val = query_func(pts.cuda()).reshape(len(xs), len(ys), len(zs)).detach().cpu().numpy()
|
|
u[xi * N: xi * N + len(xs), yi * N: yi * N + len(ys), zi * N: zi * N + len(zs)] = val
|
|
return u
|
|
|
|
|
|
def extract_geometry(bound_min, bound_max, resolution, threshold, query_func, color_func):
|
|
u = extract_fields(bound_min, bound_max, resolution, query_func)
|
|
vertices, triangles = mcubes.marching_cubes(u, threshold)
|
|
b_max_np = bound_max.detach().cpu().numpy()
|
|
b_min_np = bound_min.detach().cpu().numpy()
|
|
|
|
vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
|
|
|
|
vertices_color = color_func(vertices)
|
|
|
|
return vertices, triangles, vertices_color
|
|
|
|
|
|
def sample_pdf(bins, weights, n_samples, det=False):
|
|
# This implementation is from NeRF
|
|
# Get pdf
|
|
device = weights.device
|
|
weights = weights + 1e-5 # prevent nans
|
|
pdf = weights / torch.sum(weights, -1, keepdim=True)
|
|
cdf = torch.cumsum(pdf, -1)
|
|
cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], -1)
|
|
# Take uniform samples
|
|
if det:
|
|
u = torch.linspace(0. + 0.5 / n_samples, 1. - 0.5 / n_samples, steps=n_samples).to(device)
|
|
u = u.expand(list(cdf.shape[:-1]) + [n_samples])
|
|
else:
|
|
u = torch.rand(list(cdf.shape[:-1]) + [n_samples]).to(device)
|
|
|
|
# Invert CDF
|
|
u = u.contiguous()
|
|
inds = torch.searchsorted(cdf, u, right=True)
|
|
below = torch.max(torch.zeros_like(inds - 1).to(device), inds - 1)
|
|
above = torch.min((cdf.shape[-1] - 1) * torch.ones_like(inds).to(device), inds)
|
|
inds_g = torch.stack([below, above], -1) # (batch, N_samples, 2)
|
|
|
|
matched_shape = [inds_g.shape[0], inds_g.shape[1], cdf.shape[-1]]
|
|
cdf_g = torch.gather(cdf.unsqueeze(1).expand(matched_shape), 2, inds_g)
|
|
bins_g = torch.gather(bins.unsqueeze(1).expand(matched_shape), 2, inds_g)
|
|
|
|
denom = (cdf_g[..., 1] - cdf_g[..., 0])
|
|
denom = torch.where(denom < 1e-5, torch.ones_like(denom).to(device), denom)
|
|
t = (u - cdf_g[..., 0]) / denom
|
|
samples = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0])
|
|
|
|
return samples
|
|
|
|
|
|
class NeuSRenderer:
|
|
def __init__(self,
|
|
nerf,
|
|
sdf_network,
|
|
deviation_network,
|
|
color_network,
|
|
n_samples,
|
|
n_importance,
|
|
n_outside,
|
|
up_sample_steps,
|
|
perturb,
|
|
sdf_decay_param):
|
|
self.nerf = nerf
|
|
self.sdf_network = sdf_network
|
|
self.deviation_network = deviation_network
|
|
self.color_network = color_network
|
|
self.n_samples = n_samples
|
|
self.n_importance = n_importance
|
|
self.n_outside = n_outside
|
|
self.up_sample_steps = up_sample_steps
|
|
self.perturb = perturb
|
|
|
|
self.sdf_decay_param = sdf_decay_param
|
|
|
|
def render_core_outside(self, rays_o, rays_d, z_vals, sample_dist, nerf, background_rgb=None):
|
|
"""
|
|
Render background
|
|
"""
|
|
batch_size, n_samples = z_vals.shape
|
|
|
|
# Section length
|
|
dists = z_vals[..., 1:] - z_vals[..., :-1]
|
|
dists = torch.cat([dists, torch.Tensor([sample_dist]).expand(dists[..., :1].shape)], -1)
|
|
mid_z_vals = z_vals + dists * 0.5
|
|
|
|
# Section midpoints
|
|
pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None] # batch_size, n_samples, 3
|
|
|
|
dis_to_center = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=True).clip(1.0, 1e10)
|
|
pts = torch.cat([pts / dis_to_center, 1.0 / dis_to_center], dim=-1) # batch_size, n_samples, 4
|
|
|
|
dirs = rays_d[:, None, :].expand(batch_size, n_samples, 3)
|
|
|
|
pts = pts.reshape(-1, 3 + int(self.n_outside > 0))
|
|
dirs = dirs.reshape(-1, 3)
|
|
|
|
density, sampled_color = nerf(pts, dirs)
|
|
sampled_color = torch.sigmoid(sampled_color)
|
|
alpha = 1.0 - torch.exp(-F.softplus(density.reshape(batch_size, n_samples)) * dists)
|
|
alpha = alpha.reshape(batch_size, n_samples)
|
|
weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]
|
|
sampled_color = sampled_color.reshape(batch_size, n_samples, 3)
|
|
color = (weights[:, :, None] * sampled_color).sum(dim=1)
|
|
if background_rgb is not None:
|
|
color = color + background_rgb * (1.0 - weights.sum(dim=-1, keepdim=True))
|
|
|
|
return {
|
|
'color': color,
|
|
'sampled_color': sampled_color,
|
|
'alpha': alpha,
|
|
'weights': weights,
|
|
}
|
|
|
|
def up_sample(self, rays_o, rays_d, z_vals, sdf, n_importance, inv_s):
|
|
"""
|
|
Up sampling give a fixed inv_s
|
|
"""
|
|
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
|
|
radius = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=False)
|
|
inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
|
|
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
|
|
cos_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)
|
|
|
|
# ----------------------------------------------------------------------------------------------------------
|
|
# Use min value of [ cos, prev_cos ]
|
|
# Though it makes the sampling (not rendering) a little bit biased, this strategy can make the sampling more
|
|
# robust when meeting situations like below:
|
|
#
|
|
# SDF
|
|
# ^
|
|
# |\ -----x----...
|
|
# | \ /
|
|
# | x x
|
|
# |---\----/-------------> 0 level
|
|
# | \ /
|
|
# | \/
|
|
# |
|
|
# ----------------------------------------------------------------------------------------------------------
|
|
prev_cos_val = torch.cat([torch.zeros([batch_size, 1]).to(device), cos_val[:, :-1]], dim=-1)
|
|
cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
|
|
cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
|
|
cos_val = cos_val.clip(-1e3, 0.0) * inside_sphere
|
|
|
|
dist = (next_z_vals - prev_z_vals)
|
|
prev_esti_sdf = mid_sdf - cos_val * dist * 0.5
|
|
next_esti_sdf = mid_sdf + cos_val * dist * 0.5
|
|
prev_cdf = torch.sigmoid(prev_esti_sdf * inv_s)
|
|
next_cdf = torch.sigmoid(next_esti_sdf * inv_s)
|
|
alpha = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)
|
|
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, last=False):
|
|
batch_size, n_samples = z_vals.shape
|
|
_, n_importance = new_z_vals.shape
|
|
pts = rays_o[:, None, :] + rays_d[:, None, :] * new_z_vals[..., :, None]
|
|
z_vals = torch.cat([z_vals, new_z_vals], dim=-1)
|
|
z_vals, index = torch.sort(z_vals, dim=-1)
|
|
|
|
if not last:
|
|
new_sdf = self.sdf_network.sdf(pts.reshape(-1, 3)).reshape(batch_size, n_importance)
|
|
sdf = torch.cat([sdf, new_sdf], 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
|
|
|
|
def render_core(self,
|
|
rays_o,
|
|
rays_d,
|
|
z_vals,
|
|
sample_dist,
|
|
sdf_network,
|
|
deviation_network,
|
|
color_network,
|
|
background_alpha=None,
|
|
background_sampled_color=None,
|
|
background_rgb=None,
|
|
cos_anneal_ratio=0.0):
|
|
batch_size, n_samples = z_vals.shape
|
|
device = rays_o.device
|
|
# Section length
|
|
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
|
|
|
|
# Section midpoints
|
|
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)
|
|
|
|
sdf_nn_output = sdf_network(pts)
|
|
sdf = sdf_nn_output[:, :1]
|
|
feature_vector = sdf_nn_output[:, 1:]
|
|
|
|
gradients = sdf_network.gradient(pts).squeeze()
|
|
sampled_color = color_network(pts, gradients, dirs, feature_vector).reshape(batch_size, n_samples, 3)
|
|
|
|
inv_s = deviation_network(torch.zeros([1, 3]).to(device))[:, :1].clip(1e-6, 1e6) # Single parameter
|
|
inv_s = inv_s.expand(batch_size * n_samples, 1)
|
|
|
|
true_cos = (dirs * gradients).sum(-1, keepdim=True)
|
|
|
|
# "cos_anneal_ratio" grows from 0 to 1 in the beginning training iterations. The anneal strategy below makes
|
|
# the cos value "not dead" at the beginning training iterations, for better convergence.
|
|
iter_cos = -(F.relu(-true_cos * 0.5 + 0.5) * (1.0 - cos_anneal_ratio) +
|
|
F.relu(-true_cos) * cos_anneal_ratio) # always non-positive
|
|
|
|
# Estimate signed distances at section points
|
|
estimated_next_sdf = sdf + iter_cos * dists.reshape(-1, 1) * 0.5
|
|
estimated_prev_sdf = sdf - iter_cos * dists.reshape(-1, 1) * 0.5
|
|
|
|
prev_cdf = torch.sigmoid(estimated_prev_sdf * inv_s)
|
|
next_cdf = torch.sigmoid(estimated_next_sdf * inv_s)
|
|
|
|
p = prev_cdf - next_cdf
|
|
c = prev_cdf
|
|
|
|
alpha = ((p + 1e-5) / (c + 1e-5)).reshape(batch_size, n_samples).clip(0.0, 1.0)
|
|
|
|
pts_norm = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=True).reshape(batch_size, n_samples)
|
|
inside_sphere = (pts_norm < 1.0).float().detach()
|
|
relax_inside_sphere = (pts_norm < 1.2).float().detach()
|
|
|
|
# Render with background
|
|
if background_alpha is not None:
|
|
alpha = alpha * inside_sphere + background_alpha[:, :n_samples] * (1.0 - inside_sphere)
|
|
alpha = torch.cat([alpha, background_alpha[:, n_samples:]], dim=-1)
|
|
sampled_color = sampled_color * inside_sphere[:, :, None] +\
|
|
background_sampled_color[:, :n_samples] * (1.0 - inside_sphere)[:, :, None]
|
|
sampled_color = torch.cat([sampled_color, background_sampled_color[:, n_samples:]], dim=1)
|
|
|
|
weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]).to(device), 1. - alpha + 1e-7], -1), -1)[:, :-1]
|
|
weights_sum = weights.sum(dim=-1, keepdim=True)
|
|
|
|
color = (sampled_color * weights[:, :, None]).sum(dim=1)
|
|
if background_rgb is not None: # Fixed background, usually black
|
|
color = color + background_rgb.to(device) * (1.0 - weights_sum)
|
|
|
|
# Eikonal loss
|
|
gradient_error = (torch.linalg.norm(gradients.reshape(batch_size, n_samples, 3), ord=2,
|
|
dim=-1) - 1.0) ** 2
|
|
gradient_error = (relax_inside_sphere * gradient_error).sum() / (relax_inside_sphere.sum() + 1e-5)
|
|
|
|
depth = (mid_z_vals * weights[:, :n_samples]).sum(dim=1, keepdim=True)
|
|
|
|
return {
|
|
'color': color,
|
|
'sdf': sdf,
|
|
'dists': dists,
|
|
'gradients': gradients.reshape(batch_size, n_samples, 3),
|
|
's_val': 1.0 / inv_s,
|
|
'mid_z_vals': mid_z_vals,
|
|
'weights': weights,
|
|
'cdf': c.reshape(batch_size, n_samples),
|
|
'gradient_error': gradient_error,
|
|
'inside_sphere': inside_sphere,
|
|
'depth': depth,
|
|
}
|
|
|
|
def render(self, rays_o, rays_d, near, far, perturb_overwrite=-1, background_rgb=None, cos_anneal_ratio=0.0):
|
|
batch_size = len(rays_o)
|
|
device = rays_o.device
|
|
sample_dist = 2.0 / self.n_samples # Assuming the region of interest is a unit sphere
|
|
z_vals = torch.linspace(0.0, 1.0, self.n_samples)
|
|
z_vals = near + (far - near) * z_vals[None, :]
|
|
|
|
z_vals_outside = None
|
|
if self.n_outside > 0:
|
|
z_vals_outside = torch.linspace(1e-3, 1.0 - 1.0 / (self.n_outside + 1.0), self.n_outside)
|
|
|
|
n_samples = self.n_samples
|
|
perturb = self.perturb
|
|
|
|
if perturb_overwrite >= 0:
|
|
perturb = perturb_overwrite
|
|
if perturb > 0:
|
|
t_rand = (torch.rand([batch_size, 1]) - 0.5)
|
|
z_vals = z_vals + t_rand * 2.0 / self.n_samples
|
|
z_vals = z_vals.to(device)
|
|
|
|
if self.n_outside > 0:
|
|
mids = .5 * (z_vals_outside[..., 1:] + z_vals_outside[..., :-1])
|
|
upper = torch.cat([mids, z_vals_outside[..., -1:]], -1)
|
|
lower = torch.cat([z_vals_outside[..., :1], mids], -1)
|
|
t_rand = torch.rand([batch_size, z_vals_outside.shape[-1]])
|
|
z_vals_outside = lower[None, :] + (upper - lower)[None, :] * t_rand
|
|
|
|
if self.n_outside > 0:
|
|
z_vals_outside = far / torch.flip(z_vals_outside, dims=[-1]) + 1.0 / self.n_samples
|
|
|
|
background_alpha = None
|
|
background_sampled_color = None
|
|
|
|
# Up sample
|
|
if self.n_importance > 0:
|
|
with torch.no_grad():
|
|
pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]
|
|
sdf = self.sdf_network.sdf(pts.reshape(-1, 3)).reshape(batch_size, self.n_samples)
|
|
|
|
for i in range(self.up_sample_steps):
|
|
new_z_vals = self.up_sample(rays_o,
|
|
rays_d,
|
|
z_vals,
|
|
sdf,
|
|
self.n_importance // self.up_sample_steps,
|
|
64 * 2**i)
|
|
z_vals, sdf = self.cat_z_vals(rays_o,
|
|
rays_d,
|
|
z_vals,
|
|
new_z_vals,
|
|
sdf,
|
|
last=(i + 1 == self.up_sample_steps))
|
|
|
|
n_samples = self.n_samples + self.n_importance
|
|
|
|
# Background model
|
|
if self.n_outside > 0:
|
|
z_vals_feed = torch.cat([z_vals, z_vals_outside], dim=-1)
|
|
z_vals_feed, _ = torch.sort(z_vals_feed, dim=-1)
|
|
ret_outside = self.render_core_outside(rays_o, rays_d, z_vals_feed, sample_dist, self.nerf)
|
|
|
|
background_sampled_color = ret_outside['sampled_color']
|
|
background_alpha = ret_outside['alpha']
|
|
|
|
# Render core
|
|
ret_fine = self.render_core(rays_o,
|
|
rays_d,
|
|
z_vals,
|
|
sample_dist,
|
|
self.sdf_network,
|
|
self.deviation_network,
|
|
self.color_network,
|
|
background_rgb=background_rgb,
|
|
background_alpha=background_alpha,
|
|
background_sampled_color=background_sampled_color,
|
|
cos_anneal_ratio=cos_anneal_ratio)
|
|
|
|
color_fine = ret_fine['color']
|
|
weights = ret_fine['weights']
|
|
weights_sum = weights.sum(dim=-1, keepdim=True)
|
|
gradients = ret_fine['gradients']
|
|
s_val = ret_fine['s_val'].reshape(batch_size, n_samples).mean(dim=-1, keepdim=True)
|
|
|
|
|
|
# - randomly sample points from the volume, and maximize the sdf
|
|
pts_random = torch.rand([1024, 3]).float().cuda() * 2 - 1 # normalized to (-1, 1)
|
|
sdf_random = self.sdf_network(pts_random)[:, :1]
|
|
|
|
sparse_loss_1 = torch.exp(
|
|
-1 * torch.abs(sdf_random) * self.sdf_decay_param).mean() # - should equal
|
|
sparse_loss_2 = torch.exp(-1 * torch.abs(ret_fine['sdf']) * self.sdf_decay_param).mean()
|
|
sparse_loss = (sparse_loss_1 + sparse_loss_2) / 2
|
|
|
|
return {
|
|
'color_fine': color_fine,
|
|
'depth': ret_fine['depth'],
|
|
's_val': s_val,
|
|
'sparse_loss': sparse_loss,
|
|
'cdf_fine': ret_fine['cdf'],
|
|
'weight_sum': weights_sum,
|
|
'weight_max': torch.max(weights, dim=-1, keepdim=True)[0],
|
|
'gradients': gradients,
|
|
'weights': weights,
|
|
'gradient_error': ret_fine['gradient_error'],
|
|
'inside_sphere': ret_fine['inside_sphere']
|
|
}
|
|
|
|
|
|
def get_vertex_colors(self, vertices):
|
|
"""
|
|
@param vertices: n,3
|
|
@return:
|
|
"""
|
|
V = vertices.shape[0]
|
|
bn = 20480
|
|
verts_colors = []
|
|
with torch.no_grad():
|
|
for vi in range(0, V, bn):
|
|
verts = torch.from_numpy(vertices[vi:vi+bn].astype(np.float32)).cuda()
|
|
feats = self.sdf_network(verts)[..., 1:]
|
|
# pdb.set_trace()
|
|
gradients = self.sdf_network.gradient(verts) # ...,3
|
|
gradients = F.normalize(gradients, dim=-1).squeeze(1)
|
|
# pdb.set_trace()
|
|
colors = self.color_network(verts, gradients, gradients, feats)
|
|
colors = torch.clamp(colors,min=0,max=1).cpu().numpy()
|
|
verts_colors.append(colors)
|
|
|
|
verts_colors = (np.concatenate(verts_colors, 0)*255).astype(np.uint8)
|
|
return verts_colors[:, ::-1]
|
|
|
|
def extract_geometry(self, bound_min, bound_max, resolution, threshold=0.0):
|
|
return extract_geometry(bound_min,
|
|
bound_max,
|
|
resolution=resolution,
|
|
threshold=threshold,
|
|
query_func=lambda pts: -self.sdf_network.sdf(pts),
|
|
color_func=lambda pts: self.get_vertex_colors(pts))
|