import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from pytorch_lightning.utilities.rank_zero import rank_zero_info
import models
from models.base import BaseModel
from models.utils import scale_anything, get_activation, cleanup, chunk_batch
from models.network_utils import get_encoding, get_mlp, get_encoding_with_network
from utils.misc import get_rank
from systems.utils import update_module_step
from nerfacc import ContractionType
def contract_to_unisphere(x, radius, contraction_type):
if contraction_type == ContractionType.AABB:
x = scale_anything(x, (-radius, radius), (0, 1))
elif contraction_type == ContractionType.UN_BOUNDED_SPHERE:
x = scale_anything(x, (-radius, radius), (0, 1))
x = x * 2 - 1 # aabb is at [-1, 1]
mag = x.norm(dim=-1, keepdim=True)
mask = mag.squeeze(-1) > 1
x[mask] = (2 - 1 / mag[mask]) * (x[mask] / mag[mask])
x = x / 4 + 0.5 # [-inf, inf] is at [0, 1]
else:
raise NotImplementedError
return x
class MarchingCubeHelper(nn.Module):
def __init__(self, resolution, use_torch=True):
super().__init__()
self.resolution = resolution
self.use_torch = use_torch
self.points_range = (0, 1)
if self.use_torch:
import torchmcubes
self.mc_func = torchmcubes.marching_cubes
else:
import mcubes
self.mc_func = mcubes.marching_cubes
self.verts = None
def grid_vertices(self):
if self.verts is None:
x, y, z = torch.linspace(*self.points_range, self.resolution), torch.linspace(*self.points_range, self.resolution), torch.linspace(*self.points_range, self.resolution)
x, y, z = torch.meshgrid(x, y, z, indexing='ij')
verts = torch.cat([x.reshape(-1, 1), y.reshape(-1, 1), z.reshape(-1, 1)], dim=-1).reshape(-1, 3)
self.verts = verts
return self.verts
def forward(self, level, threshold=0.):
level = level.float().view(self.resolution, self.resolution, self.resolution)
if self.use_torch:
verts, faces = self.mc_func(level.to(get_rank()), threshold)
verts, faces = verts.cpu(), faces.cpu().long()
else:
verts, faces = self.mc_func(-level.numpy(), threshold) # transform to numpy
verts, faces = torch.from_numpy(verts.astype(np.float32)), torch.from_numpy(faces.astype(np.int64)) # transform back to pytorch
verts = verts / (self.resolution - 1.)
return {
'v_pos': verts,
't_pos_idx': faces
}
class BaseImplicitGeometry(BaseModel):
def __init__(self, config):
super().__init__(config)
if self.config.isosurface is not None:
assert self.config.isosurface.method in ['mc', 'mc-torch']
if self.config.isosurface.method == 'mc-torch':
raise NotImplementedError("Please do not use mc-torch. It currently has some scaling issues I haven't fixed yet.")
self.helper = MarchingCubeHelper(self.config.isosurface.resolution, use_torch=self.config.isosurface.method=='mc-torch')
self.radius = self.config.radius
self.contraction_type = None # assigned in system
def forward_level(self, points):
raise NotImplementedError
def isosurface_(self, vmin, vmax):
def batch_func(x):
x = torch.stack([
scale_anything(x[...,0], (0, 1), (vmin[0], vmax[0])),
scale_anything(x[...,1], (0, 1), (vmin[1], vmax[1])),
scale_anything(x[...,2], (0, 1), (vmin[2], vmax[2])),
], dim=-1).to(self.rank)
rv = self.forward_level(x).cpu()
cleanup()
return rv
level = chunk_batch(batch_func, self.config.isosurface.chunk, True, self.helper.grid_vertices())
mesh = self.helper(level, threshold=self.config.isosurface.threshold)
mesh['v_pos'] = torch.stack([
scale_anything(mesh['v_pos'][...,0], (0, 1), (vmin[0], vmax[0])),
scale_anything(mesh['v_pos'][...,1], (0, 1), (vmin[1], vmax[1])),
scale_anything(mesh['v_pos'][...,2], (0, 1), (vmin[2], vmax[2]))
], dim=-1)
return mesh
@torch.no_grad()
def isosurface(self):
if self.config.isosurface is None:
raise NotImplementedError
mesh_coarse = self.isosurface_((-self.radius, -self.radius, -self.radius), (self.radius, self.radius, self.radius))
vmin, vmax = mesh_coarse['v_pos'].amin(dim=0), mesh_coarse['v_pos'].amax(dim=0)
vmin_ = (vmin - (vmax - vmin) * 0.1).clamp(-self.radius, self.radius)
vmax_ = (vmax + (vmax - vmin) * 0.1).clamp(-self.radius, self.radius)
mesh_fine = self.isosurface_(vmin_, vmax_)
return mesh_fine
@models.register('volume-density')
class VolumeDensity(BaseImplicitGeometry):
def setup(self):
self.n_input_dims = self.config.get('n_input_dims', 3)
self.n_output_dims = self.config.feature_dim
self.encoding_with_network = get_encoding_with_network(self.n_input_dims, self.n_output_dims, self.config.xyz_encoding_config, self.config.mlp_network_config)
def forward(self, points):
points = contract_to_unisphere(points, self.radius, self.contraction_type)
out = self.encoding_with_network(points.view(-1, self.n_input_dims)).view(*points.shape[:-1], self.n_output_dims).float()
density, feature = out[...,0], out
if 'density_activation' in self.config:
density = get_activation(self.config.density_activation)(density + float(self.config.density_bias))
if 'feature_activation' in self.config:
feature = get_activation(self.config.feature_activation)(feature)
return density, feature
def forward_level(self, points):
points = contract_to_unisphere(points, self.radius, self.contraction_type)
density = self.encoding_with_network(points.reshape(-1, self.n_input_dims)).reshape(*points.shape[:-1], self.n_output_dims)[...,0]
if 'density_activation' in self.config:
density = get_activation(self.config.density_activation)(density + float(self.config.density_bias))
return -density
def update_step(self, epoch, global_step):
update_module_step(self.encoding_with_network, epoch, global_step)
@models.register('volume-sdf')
class VolumeSDF(BaseImplicitGeometry):
def setup(self):
self.n_output_dims = self.config.feature_dim
encoding = get_encoding(3, self.config.xyz_encoding_config)
network = get_mlp(encoding.n_output_dims, self.n_output_dims, self.config.mlp_network_config)
self.encoding, self.network = encoding, network
self.grad_type = self.config.grad_type
self.finite_difference_eps = self.config.get('finite_difference_eps', 1e-3)
# the actual value used in training
# will update at certain steps if finite_difference_eps="progressive"
self._finite_difference_eps = None
if self.grad_type == 'finite_difference':
rank_zero_info(f"Using finite difference to compute gradients with eps={self.finite_difference_eps}")
def forward(self, points, with_grad=True, with_feature=True, with_laplace=False):
with torch.inference_mode(torch.is_inference_mode_enabled() and not (with_grad and self.grad_type == 'analytic')):
with torch.set_grad_enabled(self.training or (with_grad and self.grad_type == 'analytic')):
if with_grad and self.grad_type == 'analytic':
if not self.training:
points = points.clone() # points may be in inference mode, get a copy to enable grad
points.requires_grad_(True)
points_ = points # points in the original scale
points = contract_to_unisphere(points, self.radius, self.contraction_type) # points normalized to (0, 1)
out = self.network(self.encoding(points.view(-1, 3))).view(*points.shape[:-1], self.n_output_dims).float()
sdf, feature = out[...,0], out
if 'sdf_activation' in self.config:
sdf = get_activation(self.config.sdf_activation)(sdf + float(self.config.sdf_bias))
if 'feature_activation' in self.config:
feature = get_activation(self.config.feature_activation)(feature)
if with_grad:
if self.grad_type == 'analytic':
grad = torch.autograd.grad(
sdf, points_, grad_outputs=torch.ones_like(sdf),
create_graph=True, retain_graph=True, only_inputs=True
)[0]
elif self.grad_type == 'finite_difference':
eps = self._finite_difference_eps
offsets = torch.as_tensor(
[
[eps, 0.0, 0.0],
[-eps, 0.0, 0.0],
[0.0, eps, 0.0],
[0.0, -eps, 0.0],
[0.0, 0.0, eps],
[0.0, 0.0, -eps],
]
).to(points_)
points_d_ = (points_[...,None,:] + offsets).clamp(-self.radius, self.radius)
points_d = scale_anything(points_d_, (-self.radius, self.radius), (0, 1))
points_d_sdf = self.network(self.encoding(points_d.view(-1, 3)))[...,0].view(*points.shape[:-1], 6).float()
grad = 0.5 * (points_d_sdf[..., 0::2] - points_d_sdf[..., 1::2]) / eps
if with_laplace:
laplace = (points_d_sdf[..., 0::2] + points_d_sdf[..., 1::2] - 2 * sdf[..., None]).sum(-1) / (eps ** 2)
rv = [sdf]
if with_grad:
rv.append(grad)
if with_feature:
rv.append(feature)
if with_laplace:
assert self.config.grad_type == 'finite_difference', "Laplace computation is only supported with grad_type='finite_difference'"
rv.append(laplace)
rv = [v if self.training else v.detach() for v in rv]
return rv[0] if len(rv) == 1 else rv
def forward_level(self, points):
points = contract_to_unisphere(points, self.radius, self.contraction_type) # points normalized to (0, 1)
sdf = self.network(self.encoding(points.view(-1, 3))).view(*points.shape[:-1], self.n_output_dims)[...,0]
if 'sdf_activation' in self.config:
sdf = get_activation(self.config.sdf_activation)(sdf + float(self.config.sdf_bias))
return sdf
def update_step(self, epoch, global_step):
update_module_step(self.encoding, epoch, global_step)
update_module_step(self.network, epoch, global_step)
if self.grad_type == 'finite_difference':
if isinstance(self.finite_difference_eps, float):
self._finite_difference_eps = self.finite_difference_eps
elif self.finite_difference_eps == 'progressive':
hg_conf = self.config.xyz_encoding_config
assert hg_conf.otype == "ProgressiveBandHashGrid", "finite_difference_eps='progressive' only works with ProgressiveBandHashGrid"
current_level = min(
hg_conf.start_level + max(global_step - hg_conf.start_step, 0) // hg_conf.update_steps,
hg_conf.n_levels
)
grid_res = hg_conf.base_resolution * hg_conf.per_level_scale**(current_level - 1)
grid_size = 2 * self.config.radius / grid_res
if grid_size != self._finite_difference_eps:
rank_zero_info(f"Update finite_difference_eps to {grid_size}")
self._finite_difference_eps = grid_size
else:
raise ValueError(f"Unknown finite_difference_eps={self.finite_difference_eps}")