Add Deep Marching Tetrahedra (DMTet)
This commit is contained in:
MrForExample
@@ -0,0 +1,8 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
grid_res = 128
|
||||
p = os.path.join(os.path.dirname(__file__), f'../data/tets/{grid_res}_tets.npz')
|
||||
print(p)
|
||||
|
||||
@@ -0,0 +1,148 @@
|
||||
import os
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import nvdiffrast.torch as dr
|
||||
|
||||
from ..mesh_processer.mesh import safe_normalize
|
||||
|
||||
def scale_img_nhwc(x, size, mag='bilinear', min='bilinear'):
|
||||
assert (x.shape[1] >= size[0] and x.shape[2] >= size[1]) or (x.shape[1] < size[0] and x.shape[2] < size[1]), "Trying to magnify image in one dimension and minify in the other"
|
||||
y = x.permute(0, 3, 1, 2) # NHWC -> NCHW
|
||||
if x.shape[1] > size[0] and x.shape[2] > size[1]: # Minification, previous size was bigger
|
||||
y = torch.nn.functional.interpolate(y, size, mode=min)
|
||||
else: # Magnification
|
||||
if mag == 'bilinear' or mag == 'bicubic':
|
||||
y = torch.nn.functional.interpolate(y, size, mode=mag, align_corners=True)
|
||||
else:
|
||||
y = torch.nn.functional.interpolate(y, size, mode=mag)
|
||||
return y.permute(0, 2, 3, 1).contiguous() # NCHW -> NHWC
|
||||
|
||||
def scale_img_hwc(x, size, mag='bilinear', min='bilinear'):
|
||||
return scale_img_nhwc(x[None, ...], size, mag, min)[0]
|
||||
|
||||
def scale_img_nhw(x, size, mag='bilinear', min='bilinear'):
|
||||
return scale_img_nhwc(x[..., None], size, mag, min)[..., 0]
|
||||
|
||||
def scale_img_hw(x, size, mag='bilinear', min='bilinear'):
|
||||
return scale_img_nhwc(x[None, ..., None], size, mag, min)[0, ..., 0]
|
||||
|
||||
def trunc_rev_sigmoid(x, eps=1e-6):
|
||||
x = x.clamp(eps, 1 - eps)
|
||||
return torch.log(x / (1 - x))
|
||||
|
||||
def make_divisible(x, m=8):
|
||||
return int(math.ceil(x / m) * m)
|
||||
|
||||
class Renderer(nn.Module):
|
||||
def __init__(self, mesh, force_cuda_rast):
|
||||
|
||||
super().__init__()
|
||||
|
||||
self.mesh = mesh
|
||||
|
||||
if force_cuda_rast or os.name != 'nt':
|
||||
self.glctx = dr.RasterizeCudaContext()
|
||||
else:
|
||||
self.glctx = dr.RasterizeGLContext()
|
||||
|
||||
# extract trainable parameters
|
||||
self.v_offsets = nn.Parameter(torch.zeros_like(self.mesh.v), requires_grad=True)
|
||||
self.raw_albedo = nn.Parameter(trunc_rev_sigmoid(self.mesh.albedo), requires_grad=True)
|
||||
|
||||
|
||||
def get_params(self, texture_lr, train_geo, geom_lr):
|
||||
|
||||
params = [
|
||||
{'params': self.raw_albedo, 'lr': texture_lr},
|
||||
]
|
||||
|
||||
self.train_geo = train_geo
|
||||
if train_geo:
|
||||
params.append({'params': self.v_offsets, 'lr': geom_lr})
|
||||
|
||||
return params
|
||||
|
||||
def update_mesh(self):
|
||||
self.mesh.v = (self.mesh.v + self.v_offsets).detach()
|
||||
self.mesh.albedo = torch.sigmoid(self.raw_albedo.detach())
|
||||
|
||||
def render(self, pose, proj, h0, w0, ssaa=1, bg_color=1, texture_filter='linear-mipmap-linear'):
|
||||
|
||||
# do super-sampling
|
||||
if ssaa != 1:
|
||||
h = make_divisible(h0 * ssaa, 8)
|
||||
w = make_divisible(w0 * ssaa, 8)
|
||||
else:
|
||||
h, w = h0, w0
|
||||
|
||||
results = {}
|
||||
|
||||
# get v
|
||||
if self.train_geo:
|
||||
v = self.mesh.v + self.v_offsets # [N, 3]
|
||||
else:
|
||||
v = self.mesh.v
|
||||
|
||||
pose = torch.from_numpy(pose.astype(np.float32)).to(v.device)
|
||||
proj = torch.from_numpy(proj.astype(np.float32)).to(v.device)
|
||||
|
||||
# get v_clip and render rgb
|
||||
v_cam = torch.matmul(F.pad(v, pad=(0, 1), mode='constant', value=1.0), torch.inverse(pose).T).float().unsqueeze(0)
|
||||
v_clip = v_cam @ proj.T
|
||||
|
||||
rast, rast_db = dr.rasterize(self.glctx, v_clip, self.mesh.f, (h, w))
|
||||
|
||||
alpha = (rast[0, ..., 3:] > 0).float()
|
||||
depth, _ = dr.interpolate(-v_cam[..., [2]], rast, self.mesh.f) # [1, H, W, 1]
|
||||
depth = depth.squeeze(0) # [H, W, 1]
|
||||
|
||||
texc, texc_db = dr.interpolate(self.mesh.vt.unsqueeze(0).contiguous(), rast, self.mesh.ft, rast_db=rast_db, diff_attrs='all')
|
||||
albedo = dr.texture(self.raw_albedo.unsqueeze(0), texc, uv_da=texc_db, filter_mode=texture_filter) # [1, H, W, 3]
|
||||
albedo = torch.sigmoid(albedo)
|
||||
# get vn and render normal
|
||||
if self.train_geo:
|
||||
i0, i1, i2 = self.mesh.f[:, 0].long(), self.mesh.f[:, 1].long(), self.mesh.f[:, 2].long()
|
||||
v0, v1, v2 = v[i0, :], v[i1, :], v[i2, :]
|
||||
|
||||
face_normals = torch.cross(v1 - v0, v2 - v0)
|
||||
face_normals = safe_normalize(face_normals)
|
||||
|
||||
vn = torch.zeros_like(v)
|
||||
vn.scatter_add_(0, i0[:, None].repeat(1,3), face_normals)
|
||||
vn.scatter_add_(0, i1[:, None].repeat(1,3), face_normals)
|
||||
vn.scatter_add_(0, i2[:, None].repeat(1,3), face_normals)
|
||||
|
||||
vn = torch.where(torch.sum(vn * vn, -1, keepdim=True) > 1e-20, vn, torch.tensor([0.0, 0.0, 1.0], dtype=torch.float32, device=vn.device))
|
||||
else:
|
||||
vn = self.mesh.vn
|
||||
|
||||
normal, _ = dr.interpolate(vn.unsqueeze(0).contiguous(), rast, self.mesh.fn)
|
||||
normal = safe_normalize(normal[0])
|
||||
|
||||
# rotated normal (where [0, 0, 1] always faces camera)
|
||||
rot_normal = normal @ pose[:3, :3]
|
||||
viewcos = rot_normal[..., [2]]
|
||||
|
||||
# antialias
|
||||
albedo = dr.antialias(albedo, rast, v_clip, self.mesh.f).squeeze(0) # [H, W, 3]
|
||||
albedo = alpha * albedo + (1 - alpha) * bg_color
|
||||
|
||||
# ssaa
|
||||
if ssaa != 1:
|
||||
albedo = scale_img_hwc(albedo, (h0, w0))
|
||||
alpha = scale_img_hwc(alpha, (h0, w0))
|
||||
depth = scale_img_hwc(depth, (h0, w0))
|
||||
normal = scale_img_hwc(normal, (h0, w0))
|
||||
viewcos = scale_img_hwc(viewcos, (h0, w0))
|
||||
|
||||
results['image'] = albedo.clamp(0, 1)
|
||||
results['alpha'] = alpha
|
||||
results['depth'] = depth
|
||||
results['normal'] = (normal + 1) / 2
|
||||
results['viewcos'] = viewcos
|
||||
|
||||
return results
|
||||
@@ -0,0 +1,120 @@
|
||||
import random
|
||||
import tqdm
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
from pytorch_msssim import SSIM, MS_SSIM
|
||||
|
||||
from .diff_mesh_renderer import Renderer
|
||||
from ..shared_utils.camera_utils import orbit_camera, OrbitCamera
|
||||
|
||||
class DiffTextureBaker:
|
||||
|
||||
def __init__(self, mesh, training_iterations, batch_size, texture_learning_rate, train_mesh_geometry, geometry_learning_rate, ms_ssim_loss_weight, force_cuda_rasterize):
|
||||
self.device = torch.device("cuda")
|
||||
|
||||
# prepare main components for optimization
|
||||
self.renderer = Renderer(mesh, force_cuda_rasterize).to(self.device)
|
||||
|
||||
self.optimizer = torch.optim.Adam(self.renderer.get_params(texture_learning_rate, train_mesh_geometry, geometry_learning_rate))
|
||||
#self.ssim_loss = SSIM(data_range=1, size_average=True, channel=3)
|
||||
self.ms_ssim_loss = MS_SSIM(data_range=1, size_average=True, channel=3)
|
||||
self.lambda_ssim = ms_ssim_loss_weight
|
||||
|
||||
self.training_iterations = training_iterations
|
||||
|
||||
self.batch_size = batch_size
|
||||
|
||||
def prepare_img(self, img):
|
||||
img_new = img.permute(2, 0, 1).unsqueeze(0).to(self.device)
|
||||
img_new = F.interpolate(img_new, (self.ref_size_H, self.ref_size_W), mode="bilinear", align_corners=False).contiguous()
|
||||
return img_new
|
||||
|
||||
def prepare_training(self, reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy):
|
||||
self.ref_imgs_num = len(reference_images)
|
||||
|
||||
self.all_ref_cam_poses = reference_orbit_camera_poses
|
||||
self.ref_cam_fovy = reference_orbit_camera_fovy
|
||||
|
||||
self.ref_size_H = reference_images[0].shape[0]
|
||||
self.ref_size_W = reference_images[0].shape[1]
|
||||
|
||||
self.cam = OrbitCamera(self.ref_size_W, self.ref_size_H, fovy=reference_orbit_camera_fovy)
|
||||
|
||||
# prepare reference images and masks
|
||||
ref_imgs_torch_list = []
|
||||
ref_masks_torch_list = []
|
||||
for i in range(self.ref_imgs_num):
|
||||
ref_imgs_torch_list.append(self.prepare_img(reference_images[i]))
|
||||
ref_masks_torch_list.append(self.prepare_img(reference_masks[i].unsqueeze(2)))
|
||||
|
||||
self.ref_imgs_torch = torch.cat(ref_imgs_torch_list, dim=0)
|
||||
self.ref_masks_torch = torch.cat(ref_masks_torch_list, dim=0)
|
||||
|
||||
def training(self):
|
||||
starter = torch.cuda.Event(enable_timing=True)
|
||||
ender = torch.cuda.Event(enable_timing=True)
|
||||
starter.record()
|
||||
|
||||
ref_imgs_masked = []
|
||||
for i in range(self.ref_imgs_num):
|
||||
ref_imgs_masked.append((self.ref_imgs_torch[i] * self.ref_masks_torch[i]).unsqueeze(0))
|
||||
|
||||
ref_imgs_num_minus_1 = self.ref_imgs_num-1
|
||||
|
||||
for step in tqdm.trange(self.training_iterations):
|
||||
|
||||
### calculate loss between reference and rendered image from known view
|
||||
loss = 0
|
||||
masked_rendered_img_batch = []
|
||||
masked_ref_img_batch = []
|
||||
for _ in range(self.batch_size):
|
||||
|
||||
i = random.randint(0, ref_imgs_num_minus_1)
|
||||
|
||||
radius, elevation, azimuth, center_X, center_Y, center_Z = self.all_ref_cam_poses[i]
|
||||
|
||||
# render output
|
||||
orbit_target = np.array([center_X, center_Y, center_Z], dtype=np.float32)
|
||||
ref_pose = orbit_camera(elevation, azimuth, radius, target=orbit_target)
|
||||
ref_cam = (ref_pose, self.cam.perspective)
|
||||
out = self.renderer.render(*ref_cam, self.ref_size_H, self.ref_size_W, ssaa=1) #ssaa = min(2.0, max(0.125, 2 * np.random.random()))
|
||||
|
||||
image = out["image"] # [H, W, 3] in [0, 1]
|
||||
image = image.permute(2, 0, 1).contiguous() # [3, H, W] in [0, 1]
|
||||
|
||||
image_masked = (image * self.ref_masks_torch[i]).unsqueeze(0)
|
||||
|
||||
masked_rendered_img_batch.append(image_masked)
|
||||
masked_ref_img_batch.append(ref_imgs_masked[i])
|
||||
|
||||
masked_rendered_img_batch_torch = torch.cat(masked_rendered_img_batch, dim=0)
|
||||
masked_ref_img_batch_torch = torch.cat(masked_ref_img_batch, dim=0)
|
||||
|
||||
# rgb loss
|
||||
loss += (1 - self.lambda_ssim) * F.mse_loss(masked_rendered_img_batch_torch, masked_ref_img_batch_torch)
|
||||
|
||||
# D-SSIM loss
|
||||
# [1, 3, H, W] in [0, 1]
|
||||
#loss += self.lambda_ssim * (1 - self.ssim_loss(X, Y))
|
||||
loss += self.lambda_ssim * (1 - self.ms_ssim_loss(masked_ref_img_batch_torch, masked_rendered_img_batch_torch))
|
||||
|
||||
# optimize step
|
||||
loss.backward()
|
||||
self.optimizer.step()
|
||||
self.optimizer.zero_grad()
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
self.need_update = True
|
||||
|
||||
print(f"Step: {step}")
|
||||
|
||||
self.renderer.update_mesh()
|
||||
|
||||
ender.record()
|
||||
#t = starter.elapsed_time(ender)
|
||||
|
||||
def get_mesh_and_texture(self):
|
||||
return (self.renderer.mesh, self.renderer.mesh.albedo, )
|
||||
@@ -0,0 +1,271 @@
|
||||
import os
|
||||
import tqdm
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from pytorch3d.loss import chamfer_distance, mesh_laplacian_smoothing
|
||||
from pytorch3d.structures.meshes import Meshes
|
||||
|
||||
from ..mesh_processer.mesh import Mesh
|
||||
from ..mesh_processer.mesh_utils import sample_points
|
||||
from .dmtet_network import Decoder
|
||||
|
||||
# Deep Marching Tetrahedrons implementation, adapted from
|
||||
# https://github.com/NVIDIAGameWorks/kaolin/blob/master/kaolin/ops/conversions/tetmesh.py
|
||||
class DMTet:
|
||||
def __init__(self, device='cuda'):
|
||||
self.device = device
|
||||
self.triangle_table = torch.tensor([
|
||||
[-1, -1, -1, -1, -1, -1],
|
||||
[ 1, 0, 2, -1, -1, -1],
|
||||
[ 4, 0, 3, -1, -1, -1],
|
||||
[ 1, 4, 2, 1, 3, 4],
|
||||
[ 3, 1, 5, -1, -1, -1],
|
||||
[ 2, 3, 0, 2, 5, 3],
|
||||
[ 1, 4, 0, 1, 5, 4],
|
||||
[ 4, 2, 5, -1, -1, -1],
|
||||
[ 4, 5, 2, -1, -1, -1],
|
||||
[ 4, 1, 0, 4, 5, 1],
|
||||
[ 3, 2, 0, 3, 5, 2],
|
||||
[ 1, 3, 5, -1, -1, -1],
|
||||
[ 4, 1, 2, 4, 3, 1],
|
||||
[ 3, 0, 4, -1, -1, -1],
|
||||
[ 2, 0, 1, -1, -1, -1],
|
||||
[-1, -1, -1, -1, -1, -1]
|
||||
], dtype=torch.long, device=self.device)
|
||||
|
||||
self.num_triangles_table = torch.tensor([0,1,1,2,1,2,2,1,1,2,2,1,2,1,1,0], dtype=torch.long, device=self.device)
|
||||
self.base_tet_edges = torch.tensor([0,1,0,2,0,3,1,2,1,3,2,3], dtype=torch.long, device=self.device)
|
||||
|
||||
# Utility functions
|
||||
def sort_edges(self, edges_ex2):
|
||||
with torch.no_grad():
|
||||
order = (edges_ex2[:,0] > edges_ex2[:,1]).long()
|
||||
order = order.unsqueeze(dim=1)
|
||||
|
||||
a = torch.gather(input=edges_ex2, index=order, dim=1)
|
||||
b = torch.gather(input=edges_ex2, index=1-order, dim=1)
|
||||
|
||||
return torch.stack([a, b],-1)
|
||||
|
||||
def map_uv(self, faces, face_gidx, max_idx):
|
||||
N = int(np.ceil(np.sqrt((max_idx+1)//2)))
|
||||
tex_y, tex_x = torch.meshgrid(
|
||||
torch.linspace(0, 1 - (1 / N), N, dtype=torch.float32, device="cuda"),
|
||||
torch.linspace(0, 1 - (1 / N), N, dtype=torch.float32, device="cuda"),
|
||||
indexing='ij'
|
||||
)
|
||||
|
||||
pad = 0.9 / N
|
||||
|
||||
uvs = torch.stack([
|
||||
tex_x , tex_y,
|
||||
tex_x + pad, tex_y,
|
||||
tex_x + pad, tex_y + pad,
|
||||
tex_x , tex_y + pad
|
||||
], dim=-1).view(-1, 2)
|
||||
|
||||
def _idx(tet_idx, N):
|
||||
x = tet_idx % N
|
||||
y = torch.div(tet_idx, N, rounding_mode='trunc')
|
||||
return y * N + x
|
||||
|
||||
tet_idx = _idx(torch.div(face_gidx, 2, rounding_mode='trunc'), N)
|
||||
tri_idx = face_gidx % 2
|
||||
|
||||
uv_idx = torch.stack((
|
||||
tet_idx * 4, tet_idx * 4 + tri_idx + 1, tet_idx * 4 + tri_idx + 2
|
||||
), dim = -1). view(-1, 3)
|
||||
|
||||
return uvs, uv_idx
|
||||
|
||||
# Marching tets implementation
|
||||
def __call__(self, vertices, tet, sdf, with_uv=False):
|
||||
|
||||
r"""Convert discrete signed distance fields encoded on tetrahedral grids to triangle
|
||||
meshes using marching tetrahedra algorithm as described in `An efficient method of
|
||||
triangulating equi-valued surfaces by using tetrahedral cells`_. The output surface is differentiable with respect to
|
||||
input vertex positions and the SDF values. For more details and example usage in learning, see
|
||||
`Deep Marching Tetrahedra\: a Hybrid Representation for High-Resolution 3D Shape Synthesis`_ NeurIPS 2021.
|
||||
|
||||
|
||||
Args:
|
||||
vertices (torch.tensor): batched vertices of tetrahedral meshes, of shape
|
||||
: (num_vertices, 3).
|
||||
tets (torch.tensor): unbatched tetrahedral mesh topology, of shape
|
||||
: (num_tetrahedrons, 4).
|
||||
sdf (torch.tensor): batched SDFs which specify the SDF value of each vertex, of shape
|
||||
: (num_vertices).
|
||||
return_tet_idx (optional, bool): if True, return index of tetrahedron
|
||||
where each face is extracted. Default: False.
|
||||
|
||||
Returns:
|
||||
(list[torch.Tensor], list[torch.LongTensor], (optional) list[torch.LongTensor]):
|
||||
|
||||
- the list of vertices for mesh converted from each tetrahedral grid.
|
||||
- the list of faces for mesh converted from each tetrahedral grid.
|
||||
- the list of indices that correspond to tetrahedra where faces are extracted.
|
||||
|
||||
Example:
|
||||
>>> vertices = torch.tensor([[[0, 0, 0],
|
||||
... [1, 0, 0],
|
||||
... [0, 1, 0],
|
||||
... [0, 0, 1]]], dtype=torch.float)
|
||||
>>> tets = torch.tensor([[0, 1, 2, 3]], dtype=torch.long)
|
||||
>>> sdf = torch.tensor([[-1., -1., 0.5, 0.5]], dtype=torch.float)
|
||||
>>> verts_list, faces_list, tet_idx_list = marching_tetrahedra(vertices, tets, sdf, True)
|
||||
>>> verts_list[0]
|
||||
tensor([[0.0000, 0.6667, 0.0000],
|
||||
[0.0000, 0.0000, 0.6667],
|
||||
[0.3333, 0.6667, 0.0000],
|
||||
[0.3333, 0.0000, 0.6667]])
|
||||
>>> faces_list[0]
|
||||
tensor([[3, 0, 1],
|
||||
[3, 2, 0]])
|
||||
>>> tet_idx_list[0]
|
||||
tensor([0, 0])
|
||||
|
||||
.. _An efficient method of triangulating equi-valued surfaces by using tetrahedral cells:
|
||||
https://search.ieice.org/bin/summary.php?id=e74-d_1_214
|
||||
|
||||
.. _Deep Marching Tetrahedra\: a Hybrid Representation for High-Resolution 3D Shape Synthesis:
|
||||
https://arxiv.org/abs/2111.04276
|
||||
"""
|
||||
|
||||
with torch.no_grad():
|
||||
occ_n = sdf > 0
|
||||
occ_fx4 = occ_n[tet.reshape(-1)].reshape(-1,4)
|
||||
occ_sum = torch.sum(occ_fx4, -1)
|
||||
valid_tets = (occ_sum>0) & (occ_sum<4)
|
||||
occ_sum = occ_sum[valid_tets]
|
||||
|
||||
# find all vertices
|
||||
all_edges = tet[valid_tets][:,self.base_tet_edges].reshape(-1,2)
|
||||
all_edges = self.sort_edges(all_edges)
|
||||
unique_edges, idx_map = torch.unique(all_edges,dim=0, return_inverse=True)
|
||||
|
||||
unique_edges = unique_edges.long()
|
||||
mask_edges = occ_n[unique_edges.reshape(-1)].reshape(-1,2).sum(-1) == 1
|
||||
mapping = torch.ones((unique_edges.shape[0]), dtype=torch.long, device="cuda") * -1
|
||||
mapping[mask_edges] = torch.arange(mask_edges.sum(), dtype=torch.long,device="cuda")
|
||||
idx_map = mapping[idx_map] # map edges to verts
|
||||
|
||||
interp_v = unique_edges[mask_edges]
|
||||
edges_to_interp = vertices[interp_v.reshape(-1)].reshape(-1,2,3)
|
||||
edges_to_interp_sdf = sdf[interp_v.reshape(-1)].reshape(-1,2,1)
|
||||
edges_to_interp_sdf[:,-1] *= -1
|
||||
|
||||
denominator = edges_to_interp_sdf.sum(1,keepdim = True)
|
||||
|
||||
edges_to_interp_sdf = torch.flip(edges_to_interp_sdf, [1])/denominator
|
||||
verts = (edges_to_interp * edges_to_interp_sdf).sum(1)
|
||||
|
||||
idx_map = idx_map.reshape(-1,6)
|
||||
|
||||
v_id = torch.pow(2, torch.arange(4, dtype=torch.long, device="cuda"))
|
||||
tetindex = (occ_fx4[valid_tets] * v_id.unsqueeze(0)).sum(-1)
|
||||
num_triangles = self.num_triangles_table[tetindex]
|
||||
|
||||
# Generate triangle indices
|
||||
faces = torch.cat((
|
||||
torch.gather(input=idx_map[num_triangles == 1], dim=1, index=self.triangle_table[tetindex[num_triangles == 1]][:, :3]).reshape(-1,3),
|
||||
torch.gather(input=idx_map[num_triangles == 2], dim=1, index=self.triangle_table[tetindex[num_triangles == 2]][:, :6]).reshape(-1,3),
|
||||
), dim=0)
|
||||
|
||||
if with_uv:
|
||||
# Get global face index (static, does not depend on topology)
|
||||
num_tets = tet.shape[0]
|
||||
tet_gidx = torch.arange(num_tets, dtype=torch.long, device="cuda")[valid_tets]
|
||||
face_gidx = torch.cat((
|
||||
tet_gidx[num_triangles == 1]*2,
|
||||
torch.stack((tet_gidx[num_triangles == 2]*2, tet_gidx[num_triangles == 2]*2 + 1), dim=-1).view(-1)
|
||||
), dim=0)
|
||||
|
||||
uvs, uv_idx = self.map_uv(faces, face_gidx, num_tets*2)
|
||||
|
||||
return verts, faces, uvs, uv_idx
|
||||
else:
|
||||
return verts, faces
|
||||
|
||||
# Geometry interface
|
||||
class DMTetMesh:
|
||||
def __init__(self, iterations=5000, chamfer_distance_weight=1, laplacian_weight=0.1, chamfer_faces_sample_scale=1, mesh_scale=1.0, grid_res=128, geom_lr=0.0001, multires=2, mlp_internal_dims=128, mlp_hidden_layer_num=5, device='cuda'):
|
||||
|
||||
self.device = torch.device(device)
|
||||
self.chamfer_distance_weight = chamfer_distance_weight
|
||||
self.laplacian_weight = laplacian_weight
|
||||
self.chamfer_faces_sample_scale = chamfer_faces_sample_scale
|
||||
self.iterations = iterations
|
||||
self.grid_res = grid_res
|
||||
|
||||
self.marching_tets = DMTet(device=self.device)
|
||||
|
||||
self.model = Decoder(multires=multires, internal_dims=mlp_internal_dims, hidden=mlp_hidden_layer_num).to(self.device)
|
||||
self.model.pre_train_sphere()
|
||||
|
||||
self.geom_optimizer = torch.optim.Adam(self.model.parameters(), lr=geom_lr)
|
||||
self.geom_scheduler = torch.optim.lr_scheduler.LambdaLR(self.geom_optimizer, lr_lambda=lambda x: self.lr_schedule(x))
|
||||
|
||||
# Load tetrahedral tile, for grid_res==128: verts: torch.Size([277410, 3]), indices: torch.Size([1524684, 4])
|
||||
# Get pre-generated tetrahedral grid from https://github.com/NVlabs/nvdiffrec/tree/main/data/tets
|
||||
tets_path = os.path.join(os.path.dirname(__file__), f'../data/tets/{self.grid_res}_tets.npz')
|
||||
tets = np.load(tets_path)
|
||||
self.verts = torch.tensor(tets['vertices'], dtype=torch.float32, device=self.device) * mesh_scale
|
||||
self.indices = torch.tensor(tets['indices'], dtype=torch.long, device=self.device)
|
||||
|
||||
self.generate_edges()
|
||||
|
||||
def lr_schedule(self, iter):
|
||||
return max(0.0, 10**(-iter*0.0002)) # Exponential falloff from [1.0, 0.1] over 5k epochs.
|
||||
|
||||
def generate_edges(self):
|
||||
with torch.no_grad():
|
||||
edges = torch.tensor([0,1,0,2,0,3,1,2,1,3,2,3], dtype = torch.long, device = "cuda")
|
||||
all_edges = self.indices[:,edges].reshape(-1,2)
|
||||
all_edges_sorted = torch.sort(all_edges, dim=1)[0]
|
||||
self.all_edges = torch.unique(all_edges_sorted, dim=0)
|
||||
|
||||
def get_mesh(self):
|
||||
verts, faces = self.pred_verts_faces()
|
||||
mesh = Mesh(v=verts, f=faces)
|
||||
mesh.auto_normal()
|
||||
mesh.auto_uv()
|
||||
return mesh
|
||||
|
||||
def pred_verts_faces(self):
|
||||
pred = self.model(self.verts) # predict SDF and per-vertex deformation
|
||||
sdf, deform = self.model.get_sdf_and_deform(pred)
|
||||
v_deformed = self.verts + torch.tanh(deform) / self.grid_res # constraint deformation to avoid flipping tets
|
||||
return self.marching_tets(v_deformed, self.indices, sdf)
|
||||
|
||||
def training(self, pcd, ref_imgs, ref_masks):
|
||||
|
||||
allow_ref_points_loss = pcd is not None
|
||||
if allow_ref_points_loss:
|
||||
ref_points_batched = torch.from_numpy(np.asarray(pcd.points)).float().to(self.device).unsqueeze(0)
|
||||
|
||||
allow_imgs_loss = ref_imgs is not None
|
||||
allow_masks_loss = ref_masks is not None
|
||||
|
||||
for step in tqdm.trange(self.iterations):
|
||||
|
||||
loss = 0
|
||||
|
||||
verts, faces = self.pred_verts_faces()
|
||||
|
||||
# chamfer distance loss
|
||||
if allow_ref_points_loss:
|
||||
num_faces_samples = int(self.chamfer_faces_sample_scale * faces.shape[0])
|
||||
sampled_points_per_faces = sample_points(verts.unsqueeze(0), faces, num_faces_samples)[0]
|
||||
loss += self.chamfer_distance_weight * chamfer_distance(sampled_points_per_faces, ref_points_batched)[0].mean()
|
||||
|
||||
# laplacian smoothing regularization loss
|
||||
mesh = Meshes([verts], [faces])
|
||||
loss += self.laplacian_weight * mesh_laplacian_smoothing(mesh)
|
||||
|
||||
# TODO: add images based diffrast loss, i.e (RGB, Alpha, Normal; with L1, SSIM loss)
|
||||
|
||||
loss.backward()
|
||||
self.geom_optimizer.step()
|
||||
self.geom_scheduler.step()
|
||||
self.geom_optimizer.zero_grad()
|
||||
@@ -0,0 +1,93 @@
|
||||
# This script is modified based on: https://github.com/NVIDIAGameWorks/kaolin/blob/5406915c3a020596384b1a36679ff68afd04d455/examples/tutorial/dmtet_network.py
|
||||
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
|
||||
# MLP + Positional Encoding
|
||||
class Decoder(torch.nn.Module):
|
||||
def __init__(self, input_dims = 3, internal_dims = 128, output_dims = 4, hidden = 5, multires = 2):
|
||||
super().__init__()
|
||||
self.embed_fn = None
|
||||
if multires > 0:
|
||||
embed_fn, input_ch = get_embedder(multires)
|
||||
self.embed_fn = embed_fn
|
||||
input_dims = input_ch
|
||||
|
||||
net = (torch.nn.Linear(input_dims, internal_dims, bias=False), torch.nn.ReLU())
|
||||
for i in range(hidden-1):
|
||||
net = net + (torch.nn.Linear(internal_dims, internal_dims, bias=False), torch.nn.ReLU())
|
||||
net = net + (torch.nn.Linear(internal_dims, output_dims, bias=False),)
|
||||
self.net = torch.nn.Sequential(*net)
|
||||
|
||||
def forward(self, p):
|
||||
if self.embed_fn is not None:
|
||||
p = self.embed_fn(p)
|
||||
out = self.net(p)
|
||||
return out
|
||||
|
||||
def get_sdf_and_deform(self, out):
|
||||
return out[:,0], out[:,1:] # sdf, deform
|
||||
|
||||
def pre_train_sphere(self, iter=1000):
|
||||
print ("Initialize SDF to sphere")
|
||||
loss_fn = torch.nn.MSELoss()
|
||||
optimizer = torch.optim.Adam(list(self.parameters()), lr=1e-4)
|
||||
|
||||
for i in tqdm(range(iter)):
|
||||
p = torch.rand((1024,3), device='cuda') - 0.5
|
||||
ref_value = torch.sqrt((p**2).sum(-1)) - 0.3
|
||||
output = self(p)
|
||||
loss = loss_fn(output[...,0], ref_value)
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
print("Pre-trained MLP", loss.item())
|
||||
|
||||
|
||||
# Positional Encoding from https://github.com/yenchenlin/nerf-pytorch/blob/1f064835d2cca26e4df2d7d130daa39a8cee1795/run_nerf_helpers.py
|
||||
class Embedder:
|
||||
def __init__(self, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
self.create_embedding_fn()
|
||||
|
||||
def create_embedding_fn(self):
|
||||
embed_fns = []
|
||||
d = self.kwargs['input_dims']
|
||||
out_dim = 0
|
||||
if self.kwargs['include_input']:
|
||||
embed_fns.append(lambda x : x)
|
||||
out_dim += d
|
||||
|
||||
max_freq = self.kwargs['max_freq_log2']
|
||||
N_freqs = self.kwargs['num_freqs']
|
||||
|
||||
if self.kwargs['log_sampling']:
|
||||
freq_bands = 2.**torch.linspace(0., max_freq, steps=N_freqs)
|
||||
else:
|
||||
freq_bands = torch.linspace(2.**0., 2.**max_freq, steps=N_freqs)
|
||||
|
||||
for freq in freq_bands:
|
||||
for p_fn in self.kwargs['periodic_fns']:
|
||||
embed_fns.append(lambda x, p_fn=p_fn, freq=freq : p_fn(x * freq))
|
||||
out_dim += d
|
||||
|
||||
self.embed_fns = embed_fns
|
||||
self.out_dim = out_dim
|
||||
|
||||
def embed(self, inputs):
|
||||
return torch.cat([fn(inputs) for fn in self.embed_fns], -1)
|
||||
|
||||
def get_embedder(multires):
|
||||
embed_kwargs = {
|
||||
'include_input' : True,
|
||||
'input_dims' : 3,
|
||||
'max_freq_log2' : multires-1,
|
||||
'num_freqs' : multires,
|
||||
'log_sampling' : True,
|
||||
'periodic_fns' : [torch.sin, torch.cos],
|
||||
}
|
||||
|
||||
embedder_obj = Embedder(**embed_kwargs)
|
||||
embed = lambda x, eo=embedder_obj : eo.embed(x)
|
||||
return embed, embedder_obj.out_dim
|
||||
@@ -0,0 +1,229 @@
|
||||
import random
|
||||
import tqdm
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
from pytorch_msssim import SSIM, MS_SSIM
|
||||
|
||||
from .main_3DGS_renderer import Renderer
|
||||
from ..shared_utils.camera_utils import orbit_camera, OrbitCamera, MiniCam, calculate_fovX, get_projection_matrix
|
||||
|
||||
class GSParams:
|
||||
def __init__(self,
|
||||
training_iterations=3000,
|
||||
batch_size=1,
|
||||
loss_scale=10000,
|
||||
lambda_ssim=0.2,
|
||||
lambda_alpha=3,
|
||||
lambda_offset=0,
|
||||
lambda_offset_opacity=0,
|
||||
invert_bg_prob = 0.5,
|
||||
feature_lr=0.01,
|
||||
opacity_lr=0.05,
|
||||
scaling_lr=0.005,
|
||||
rotation_lr=0.005,
|
||||
position_lr_init=0.001,
|
||||
position_lr_final=0.00002,
|
||||
position_lr_delay_mult=0.02,
|
||||
position_lr_max_steps=500,
|
||||
num_pts=5000,
|
||||
K=3,
|
||||
percent_dense=0.01,
|
||||
density_start_iter=100,
|
||||
density_end_iter=100000,
|
||||
densification_interval=100,
|
||||
opacity_reset_interval=700,
|
||||
densify_grad_threshold=0.01,
|
||||
sh_degree=0):
|
||||
|
||||
# training params
|
||||
self.training_iterations = training_iterations
|
||||
self.batch_size = batch_size
|
||||
self.loss_scale = loss_scale
|
||||
self.lambda_ssim = lambda_ssim
|
||||
self.lambda_alpha = lambda_alpha
|
||||
self.lambda_offset = lambda_offset
|
||||
self.lambda_offset_opacity = lambda_offset_opacity
|
||||
self.invert_bg_prob = invert_bg_prob
|
||||
|
||||
# learning params
|
||||
self.feature_lr = feature_lr
|
||||
self.opacity_lr = opacity_lr
|
||||
self.scaling_lr = scaling_lr
|
||||
self.rotation_lr = rotation_lr
|
||||
self.position_lr_init = position_lr_init
|
||||
self.position_lr_final = position_lr_final
|
||||
self.position_lr_delay_mult = position_lr_delay_mult
|
||||
self.position_lr_max_steps = position_lr_max_steps
|
||||
|
||||
# densify and prune params
|
||||
self.num_pts = num_pts
|
||||
self.K = K
|
||||
self.percent_dense = percent_dense
|
||||
self.density_start_iter=density_start_iter
|
||||
self.density_end_iter=density_end_iter
|
||||
self.densification_interval=densification_interval
|
||||
self.opacity_reset_interval=opacity_reset_interval
|
||||
self.densify_grad_threshold=densify_grad_threshold
|
||||
|
||||
# other gaussian params
|
||||
self.sh_degree = sh_degree
|
||||
|
||||
|
||||
class GaussianSplatting:
|
||||
|
||||
def __init__(self, gs_params=None, init_input=None, device='cuda'):
|
||||
self.device = torch.device(device)
|
||||
|
||||
# prepare renderer for optimization
|
||||
if gs_params is None:
|
||||
gs_params = GSParams()
|
||||
|
||||
self.renderer = Renderer(sh_degree=gs_params.sh_degree)
|
||||
self.gaussain_scale_factor = 1
|
||||
self.renderer.initialize(init_input, num_pts=gs_params.num_pts)
|
||||
|
||||
# setup training
|
||||
self.renderer.gaussians.training_setup(gs_params)
|
||||
# do not do progressive sh-level
|
||||
self.renderer.gaussians.active_sh_degree = self.renderer.gaussians.max_sh_degree
|
||||
self.optimizer = self.renderer.gaussians.optimizer
|
||||
#self.ssim_loss = SSIM(data_range=1, size_average=True, channel=3)
|
||||
self.ms_ssim_loss = MS_SSIM(data_range=1, size_average=True, channel=3)
|
||||
|
||||
self.gs_params = gs_params
|
||||
|
||||
def prepare_img(self, img):
|
||||
img_new = img.permute(2, 0, 1).unsqueeze(0).to(self.device)
|
||||
img_new = F.interpolate(img_new, (self.ref_size_H, self.ref_size_W), mode="bilinear", align_corners=False).contiguous()
|
||||
return img_new
|
||||
|
||||
def prepare_training(self, reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy):
|
||||
self.ref_imgs_num = len(reference_images)
|
||||
|
||||
self.all_ref_cam_poses = reference_orbit_camera_poses
|
||||
self.ref_cam_fovy = reference_orbit_camera_fovy
|
||||
|
||||
self.ref_size_H = reference_images[0].shape[0]
|
||||
self.ref_size_W = reference_images[0].shape[1]
|
||||
|
||||
# default camera settings
|
||||
self.cam = OrbitCamera(self.ref_size_W, self.ref_size_H, fovy=reference_orbit_camera_fovy)
|
||||
self.projection_matrix = get_projection_matrix(self.cam.near, self.cam.far, self.cam.fovx, self.cam.fovy).transpose(0, 1).cuda()
|
||||
|
||||
# prepare reference images and masks
|
||||
ref_imgs_torch_list = []
|
||||
ref_masks_torch_list = []
|
||||
for i in range(self.ref_imgs_num):
|
||||
ref_imgs_torch_list.append(self.prepare_img(reference_images[i]))
|
||||
ref_masks_torch_list.append(self.prepare_img(reference_masks[i].unsqueeze(2)))
|
||||
|
||||
self.ref_imgs_torch = torch.cat(ref_imgs_torch_list, dim=0)
|
||||
self.ref_masks_torch = torch.cat(ref_masks_torch_list, dim=0)
|
||||
|
||||
def training(self):
|
||||
starter = torch.cuda.Event(enable_timing=True)
|
||||
ender = torch.cuda.Event(enable_timing=True)
|
||||
starter.record()
|
||||
|
||||
ref_imgs_masked = []
|
||||
for i in range(self.ref_imgs_num):
|
||||
ref_imgs_masked.append((self.ref_imgs_torch[i] * self.ref_masks_torch[i]).unsqueeze(0))
|
||||
|
||||
ref_imgs_num_minus_1 = self.ref_imgs_num-1
|
||||
|
||||
for step in tqdm.trange(self.gs_params.training_iterations):
|
||||
|
||||
### calculate loss between reference and rendered image from known view
|
||||
loss = 0
|
||||
masked_rendered_img_batch = []
|
||||
masked_ref_img_batch = []
|
||||
masks_batch = []
|
||||
ref_masks_batch = []
|
||||
|
||||
step_ratio = min(1, step / self.gs_params.training_iterations)
|
||||
|
||||
# update lr
|
||||
self.renderer.gaussians.update_learning_rate(step)
|
||||
|
||||
loss = 0
|
||||
|
||||
for _ in range(self.gs_params.batch_size):
|
||||
### calculate loss between reference and rendered image from known view
|
||||
|
||||
i = random.randint(0, ref_imgs_num_minus_1)
|
||||
|
||||
radius, elevation, azimuth, center_X, center_Y, center_Z = self.all_ref_cam_poses[i]
|
||||
|
||||
orbit_target = np.array([center_X, center_Y, center_Z], dtype=np.float32)
|
||||
ref_pose = orbit_camera(elevation, azimuth, radius, target=orbit_target)
|
||||
ref_cam = MiniCam(ref_pose, self.ref_size_W, self.ref_size_H, self.cam.fovy, self.cam.fovx, self.cam.near, self.cam.far, self.projection_matrix)
|
||||
|
||||
bg_color = torch.tensor([1, 1, 1] if np.random.rand() > self.gs_params.invert_bg_prob else [0, 0, 0], dtype=torch.float32, device=self.device)
|
||||
out = self.renderer.render(ref_cam, bg_color=bg_color)
|
||||
|
||||
image = out["image"] # [3, H, W] in [0, 1]
|
||||
mask = out["alpha"] # [1, H, W] in [0, 1]
|
||||
|
||||
ref_mask = self.ref_masks_torch[i]
|
||||
image_masked = (image * ref_mask).unsqueeze(0)
|
||||
|
||||
masked_rendered_img_batch.append(image_masked)
|
||||
masked_ref_img_batch.append(ref_imgs_masked[i])
|
||||
masks_batch.append(mask.unsqueeze(0))
|
||||
ref_masks_batch.append(ref_mask.unsqueeze(0))
|
||||
|
||||
masked_rendered_img_batch_torch = torch.cat(masked_rendered_img_batch, dim=0)
|
||||
masked_ref_img_batch_torch = torch.cat(masked_ref_img_batch, dim=0)
|
||||
masks_batch_torch = torch.cat(masks_batch, dim=0)
|
||||
ref_masks_batch_torch = torch.cat(ref_masks_batch, dim=0)
|
||||
|
||||
# rgb loss
|
||||
loss += (1 - self.gs_params.lambda_ssim) * self.gs_params.loss_scale * step_ratio * F.mse_loss(masked_rendered_img_batch_torch, masked_ref_img_batch_torch)
|
||||
|
||||
# alpha loss
|
||||
loss += self.gs_params.lambda_alpha * self.gs_params.loss_scale * step_ratio * F.mse_loss(masks_batch_torch, ref_masks_batch_torch)
|
||||
|
||||
# D-SSIM loss
|
||||
# [1, 3, H, W] in [0, 1]
|
||||
#loss += self.lambda_ssim * (1 - self.ssim_loss(X, Y))
|
||||
loss += self.gs_params.lambda_ssim * self.gs_params.loss_scale * step_ratio * (1 - self.ms_ssim_loss(masked_ref_img_batch_torch, masked_rendered_img_batch_torch))
|
||||
|
||||
if self.gs_params.lambda_offset > 0:
|
||||
# Reference offset loss
|
||||
offset_norm = self.renderer.gaussians.get_xyz_offset.norm(dim=-1, keepdim=True)
|
||||
loss += self.gs_params.lambda_offset * self.gs_params.loss_scale * step_ratio * torch.mean(offset_norm)
|
||||
|
||||
if self.gs_params.lambda_offset_opacity > 0:
|
||||
# Alpha penalty loss
|
||||
loss += self.gs_params.lambda_offset_opacity * self.gs_params.loss_scale * step_ratio * torch.mean(offset_norm.detach() * self.renderer.gaussians.get_opacity)
|
||||
|
||||
|
||||
# optimize step
|
||||
loss.backward()
|
||||
self.optimizer.step()
|
||||
self.optimizer.zero_grad()
|
||||
|
||||
# densify and prune
|
||||
if step >= self.gs_params.density_start_iter and step <= self.gs_params.density_end_iter:
|
||||
viewspace_point_tensor, visibility_filter, radii = out["viewspace_points"], out["visibility_filter"], out["radii"]
|
||||
self.renderer.gaussians.max_radii2D[visibility_filter] = torch.max(self.renderer.gaussians.max_radii2D[visibility_filter], radii[visibility_filter])
|
||||
self.renderer.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter)
|
||||
|
||||
if step % self.gs_params.densification_interval == 0:
|
||||
#self.renderer.gaussians.densify_and_prune(self.gs_params.densify_grad_threshold, min_opacity=0.01, extent=4, max_screen_size=1)
|
||||
#self.renderer.gaussians.densify_and_prune_by_compatness(self.gs_params.K, min_opacity=0.01, extent=4, max_screen_size=1)
|
||||
|
||||
#self.renderer.gaussians.densify_by_clone_and_split(self.gs_params.densify_grad_threshold, extent=4)
|
||||
self.renderer.gaussians.densify_by_compatness(self.gs_params.K)
|
||||
self.renderer.gaussians.prune(min_opacity=0.01, extent=4, max_screen_size=1, max_offset=0.1)
|
||||
|
||||
if step % self.gs_params.opacity_reset_interval == 0:
|
||||
self.renderer.gaussians.reset_opacity()
|
||||
|
||||
ender.record()
|
||||
torch.cuda.synchronize()
|
||||
#t = starter.elapsed_time(ender)
|
||||
|
||||
self.need_update = True
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,9 @@
|
||||
import torch
|
||||
|
||||
def sdf_reg_loss(sdf, all_edges):
|
||||
sdf_f1x6x2 = sdf[all_edges.reshape(-1)].reshape(-1,2)
|
||||
mask = torch.sign(sdf_f1x6x2[...,0]) != torch.sign(sdf_f1x6x2[...,1])
|
||||
sdf_f1x6x2 = sdf_f1x6x2[mask]
|
||||
sdf_diff = torch.nn.functional.binary_cross_entropy_with_logits(sdf_f1x6x2[...,0], (sdf_f1x6x2[...,1] > 0).float()) + \
|
||||
torch.nn.functional.binary_cross_entropy_with_logits(sdf_f1x6x2[...,1], (sdf_f1x6x2[...,0] > 0).float())
|
||||
return sdf_diff
|
||||
@@ -0,0 +1,165 @@
|
||||
# Copyright (c) 2021,22 NVIDIA CORPORATION & AFFILIATES.
|
||||
# All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
import torch
|
||||
|
||||
__all__ = ['marching_tetrahedra']
|
||||
|
||||
triangle_table = torch.tensor([
|
||||
[-1, -1, -1, -1, -1, -1],
|
||||
[1, 0, 2, -1, -1, -1],
|
||||
[4, 0, 3, -1, -1, -1],
|
||||
[1, 4, 2, 1, 3, 4],
|
||||
[3, 1, 5, -1, -1, -1],
|
||||
[2, 3, 0, 2, 5, 3],
|
||||
[1, 4, 0, 1, 5, 4],
|
||||
[4, 2, 5, -1, -1, -1],
|
||||
[4, 5, 2, -1, -1, -1],
|
||||
[4, 1, 0, 4, 5, 1],
|
||||
[3, 2, 0, 3, 5, 2],
|
||||
[1, 3, 5, -1, -1, -1],
|
||||
[4, 1, 2, 4, 3, 1],
|
||||
[3, 0, 4, -1, -1, -1],
|
||||
[2, 0, 1, -1, -1, -1],
|
||||
[-1, -1, -1, -1, -1, -1]
|
||||
], dtype=torch.long)
|
||||
|
||||
num_triangles_table = torch.tensor([0, 1, 1, 2, 1, 2, 2, 1, 1, 2, 2, 1, 2, 1, 1, 0], dtype=torch.long)
|
||||
base_tet_edges = torch.tensor([0, 1, 0, 2, 0, 3, 1, 2, 1, 3, 2, 3], dtype=torch.long)
|
||||
v_id = torch.pow(2, torch.arange(4, dtype=torch.long))
|
||||
|
||||
|
||||
def _sort_edges(edges):
|
||||
"""sort last dimension of edges of shape (E, 2)"""
|
||||
with torch.no_grad():
|
||||
order = (edges[:, 0] > edges[:, 1]).long()
|
||||
order = order.unsqueeze(dim=1)
|
||||
|
||||
a = torch.gather(input=edges, index=order, dim=1)
|
||||
b = torch.gather(input=edges, index=1 - order, dim=1)
|
||||
|
||||
return torch.stack([a, b], -1)
|
||||
|
||||
|
||||
def _unbatched_marching_tetrahedra(vertices, tets, sdf, return_tet_idx):
|
||||
"""unbatched marching tetrahedra.
|
||||
|
||||
Refer to :func:`marching_tetrahedra`.
|
||||
"""
|
||||
device = vertices.device
|
||||
with torch.no_grad():
|
||||
occ_n = sdf > 0
|
||||
occ_fx4 = occ_n[tets.reshape(-1)].reshape(-1, 4)
|
||||
occ_sum = torch.sum(occ_fx4, -1)
|
||||
valid_tets = (occ_sum > 0) & (occ_sum < 4)
|
||||
occ_sum = occ_sum[valid_tets]
|
||||
|
||||
# find all vertices
|
||||
all_edges = tets[valid_tets][:, base_tet_edges.to(device)].reshape(-1, 2)
|
||||
all_edges = _sort_edges(all_edges)
|
||||
unique_edges, idx_map = torch.unique(all_edges, dim=0, return_inverse=True)
|
||||
|
||||
unique_edges = unique_edges.long()
|
||||
mask_edges = occ_n[unique_edges.reshape(-1)].reshape(-1, 2).sum(-1) == 1
|
||||
mapping = torch.ones((unique_edges.shape[0]), dtype=torch.long, device=device) * -1
|
||||
mapping[mask_edges] = torch.arange(mask_edges.sum(), dtype=torch.long, device=device)
|
||||
idx_map = mapping[idx_map]
|
||||
|
||||
interp_v = unique_edges[mask_edges]
|
||||
edges_to_interp = vertices[interp_v.reshape(-1)].reshape(-1, 2, 3)
|
||||
edges_to_interp_sdf = sdf[interp_v.reshape(-1)].reshape(-1, 2, 1)
|
||||
edges_to_interp_sdf[:, -1] *= -1
|
||||
|
||||
denominator = edges_to_interp_sdf.sum(1, keepdim=True)
|
||||
|
||||
edges_to_interp_sdf = torch.flip(edges_to_interp_sdf, [1]) / denominator
|
||||
verts = (edges_to_interp * edges_to_interp_sdf).sum(1)
|
||||
|
||||
idx_map = idx_map.reshape(-1, 6)
|
||||
|
||||
tetindex = (occ_fx4[valid_tets] * v_id.to(device).unsqueeze(0)).sum(-1)
|
||||
num_triangles = num_triangles_table.to(device)[tetindex]
|
||||
triangle_table_device = triangle_table.to(device)
|
||||
|
||||
# Generate triangle indices
|
||||
faces = torch.cat((
|
||||
torch.gather(input=idx_map[num_triangles == 1], dim=1,
|
||||
index=triangle_table_device[tetindex[num_triangles == 1]][:, :3]).reshape(-1, 3),
|
||||
torch.gather(input=idx_map[num_triangles == 2], dim=1,
|
||||
index=triangle_table_device[tetindex[num_triangles == 2]][:, :6]).reshape(-1, 3),
|
||||
), dim=0)
|
||||
|
||||
if return_tet_idx:
|
||||
tet_idx = torch.arange(tets.shape[0], device=device)[valid_tets]
|
||||
tet_idx = torch.cat((tet_idx[num_triangles == 1], tet_idx[num_triangles ==
|
||||
2].unsqueeze(-1).expand(-1, 2).reshape(-1)), dim=0)
|
||||
return verts, faces, tet_idx
|
||||
return verts, faces
|
||||
|
||||
|
||||
def marching_tetrahedra(vertices, tets, sdf, return_tet_idx=False):
|
||||
r"""Convert discrete signed distance fields encoded on tetrahedral grids to triangle
|
||||
meshes using marching tetrahedra algorithm as described in `An efficient method of
|
||||
triangulating equi-valued surfaces by using tetrahedral cells`_. The output surface is differentiable with respect to
|
||||
input vertex positions and the SDF values. For more details and example usage in learning, see
|
||||
`Deep Marching Tetrahedra\: a Hybrid Representation for High-Resolution 3D Shape Synthesis`_ NeurIPS 2021.
|
||||
|
||||
|
||||
Args:
|
||||
vertices (torch.tensor): batched vertices of tetrahedral meshes, of shape
|
||||
:math:`(\text{batch_size}, \text{num_vertices}, 3)`.
|
||||
tets (torch.tensor): unbatched tetrahedral mesh topology, of shape
|
||||
:math:`(\text{num_tetrahedrons}, 4)`.
|
||||
sdf (torch.tensor): batched SDFs which specify the SDF value of each vertex, of shape
|
||||
:math:`(\text{batch_size}, \text{num_vertices})`.
|
||||
return_tet_idx (optional, bool): if True, return index of tetrahedron
|
||||
where each face is extracted. Default: False.
|
||||
|
||||
Returns:
|
||||
(list[torch.Tensor], list[torch.LongTensor], (optional) list[torch.LongTensor]):
|
||||
|
||||
- the list of vertices for mesh converted from each tetrahedral grid.
|
||||
- the list of faces for mesh converted from each tetrahedral grid.
|
||||
- the list of indices that correspond to tetrahedra where faces are extracted.
|
||||
|
||||
Example:
|
||||
>>> vertices = torch.tensor([[[0, 0, 0],
|
||||
... [1, 0, 0],
|
||||
... [0, 1, 0],
|
||||
... [0, 0, 1]]], dtype=torch.float)
|
||||
>>> tets = torch.tensor([[0, 1, 2, 3]], dtype=torch.long)
|
||||
>>> sdf = torch.tensor([[-1., -1., 0.5, 0.5]], dtype=torch.float)
|
||||
>>> verts_list, faces_list, tet_idx_list = marching_tetrahedra(vertices, tets, sdf, True)
|
||||
>>> verts_list[0]
|
||||
tensor([[0.0000, 0.6667, 0.0000],
|
||||
[0.0000, 0.0000, 0.6667],
|
||||
[0.3333, 0.6667, 0.0000],
|
||||
[0.3333, 0.0000, 0.6667]])
|
||||
>>> faces_list[0]
|
||||
tensor([[3, 0, 1],
|
||||
[3, 2, 0]])
|
||||
>>> tet_idx_list[0]
|
||||
tensor([0, 0])
|
||||
|
||||
.. _An efficient method of triangulating equi-valued surfaces by using tetrahedral cells:
|
||||
https://search.ieice.org/bin/summary.php?id=e74-d_1_214
|
||||
|
||||
.. _Deep Marching Tetrahedra\: a Hybrid Representation for High-Resolution 3D Shape Synthesis:
|
||||
https://arxiv.org/abs/2111.04276
|
||||
"""
|
||||
|
||||
list_of_outputs = [_unbatched_marching_tetrahedra(vertices[b], tets, sdf[b], return_tet_idx) for b in range(vertices.shape[0])]
|
||||
return list(zip(*list_of_outputs))
|
||||
Reference in New Issue
Block a user