277 lines
12 KiB
Python
277 lines
12 KiB
Python
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
|
|
|
|
import comfy.utils
|
|
|
|
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.detach(), f=faces.detach(), device=self.device)
|
|
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
|
|
|
|
comfy_pbar = comfy.utils.ProgressBar(self.iterations)
|
|
|
|
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()
|
|
|
|
comfy_pbar.update_absolute(step + 1) |