diff --git a/README.md b/README.md index 392a4ef..6e4c31e 100644 --- a/README.md +++ b/README.md @@ -72,6 +72,7 @@ Copy the files inside folder [__New_ComfyUI_Bats](./_New_ComfyUI_Bats/) to your - Load 3D file (.obj, .ply, .glb) + - Also have node that support save and load 3DGS .ply file - 3D Gaussian Splatting, with: - [Improved Differential Gaussian Rasterization](https://github.com/ashawkey/diff-gaussian-rasterization) - Better Compactness-based Densification method from [Gsgen](https://gsgen3d.github.io/), @@ -83,9 +84,14 @@ Copy the files inside folder [__New_ComfyUI_Bats](./_New_ComfyUI_Bats/) to your - Bake Multi-View images into UVTexture of given 3D mesh using [Nvdiffrast](https://github.com/NVlabs/nvdiffrast), supports: - Export to .obj, .ply, .glb +- Deep Marching Tetrahedrons + - Allow convert 3DGS .ply file to 3D mesh
+ *Note: I didn't spent time to turn the hyperprameters yet, the result will be improved in the future!* + ### To-Do Next: - Add interactive 3D UI inside ComfuUI to visulaize training and generated results for 3D representations -- Add DMTet algorithm to allow convertion from points cloud(Gaussian/.ply) to mesh (.obj, .ply, .glb) +- ~~Add DMTet algorithm to allow convertion from points cloud(Gaussian/.ply) to mesh (.obj, .ply, .glb)~~ +- Improve DMTet result & add support to training it with images(RGB, Alpha, Normal Map) - Add a general SDS/VSD Optimization algorithm to allow training 3D representations with diffusion model, *The real fun begins here* ;) - Add a few best Nerf algorithms (No idea yet, [instant-ngp](https://github.com/NVlabs/instant-ngp) maybe?) diff --git a/_Example_Workflows/Images_to_3DGS_to_DMTetMesh.json b/_Example_Workflows/Images_to_3DGS_to_DMTetMesh.json new file mode 100644 index 0000000..09df809 --- /dev/null +++ b/_Example_Workflows/Images_to_3DGS_to_DMTetMesh.json @@ -0,0 +1,602 @@ +{ + "last_node_id": 22, + "last_link_id": 26, + "nodes": [ + { + "id": 10, + "type": "PreviewImage", + "pos": [ + 1410, + -320 + ], + "size": { + "0": 210, + "1": 246 + }, + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 4 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 9, + "type": "MaskToImage", + "pos": [ + 1110, + -320 + ], + "size": { + "0": 210, + "1": 26 + }, + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "mask", + "type": "MASK", + "link": 8 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 4 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "MaskToImage" + } + }, + { + "id": 8, + "type": "InvertMask", + "pos": [ + 746, + -14 + ], + "size": { + "0": 210, + "1": 26 + }, + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "mask", + "type": "MASK", + "link": 2 + } + ], + "outputs": [ + { + "name": "MASK", + "type": "MASK", + "links": [ + 8, + 18 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "InvertMask" + } + }, + { + "id": 4, + "type": "[ComfyUI-3D] Generate Orbit Camera Poses", + "pos": [ + 830, + 132 + ], + "size": { + "0": 400, + "1": 200 + }, + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "reference_images", + "type": "IMAGE", + "link": 1 + } + ], + "outputs": [ + { + "name": "orbit_camposes", + "type": "ORBIT_CAMPOSES", + "links": [ + 19 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[ComfyUI-3D] Generate Orbit Camera Poses" + }, + "widgets_values": [ + "([0:120], 1.75, 0, 0, 360)" + ] + }, + { + "id": 7, + "type": "VHS_LoadImages", + "pos": [ + 113, + 54 + ], + "size": [ + 243.60000610351562, + 194 + ], + "flags": {}, + "order": 0, + "mode": 0, + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 1, + 17, + 22 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": [ + 2 + ], + "shape": 3, + "slot_index": 1 + }, + { + "name": "INT", + "type": "INT", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "VHS_LoadImages" + }, + "widgets_values": { + "directory": "Images_enhanced", + "image_load_cap": 0, + "skip_first_images": 0, + "select_every_nth": 1, + "choose folder to upload": "image", + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "frame_load_cap": 0, + "skip_first_images": 0, + "filename": "Images_enhanced", + "type": "input", + "format": "folder", + "select_every_nth": 1 + } + } + } + }, + { + "id": 19, + "type": "PreviewImage", + "pos": [ + 420, + -320 + ], + "size": { + "0": 210, + "1": 246 + }, + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 22 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 18, + "type": "[ComfyUI-3D] Gaussian Splatting", + "pos": [ + 1486, + 53 + ], + "size": { + "0": 468.5999755859375, + "1": 718 + }, + "flags": { + "collapsed": false + }, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "reference_images", + "type": "IMAGE", + "link": 17 + }, + { + "name": "reference_masks", + "type": "MASK", + "link": 18 + }, + { + "name": "reference_orbit_camera_poses", + "type": "ORBIT_CAMPOSES", + "link": 19 + }, + { + "name": "mesh_to_initialize_gaussian", + "type": "MESH", + "link": null + } + ], + "outputs": [ + { + "name": "raw_3DGS", + "type": "GS_RAW", + "links": [ + 21, + 24 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[ComfyUI-3D] Gaussian Splatting" + }, + "widgets_values": [ + 49.1, + 1000, + 1, + 10000, + 0.2, + 3, + 0, + 0, + 0.5, + 0.01, + 0.05, + 0.005, + 0.005, + 0.001, + 0.00002, + 0.02, + 500, + 5000, + 3, + 0.01, + 100, + 100000, + 100, + 700, + 0.01, + 0 + ] + }, + { + "id": 17, + "type": "[ComfyUI-3D] Save 3DGS", + "pos": [ + 2243, + 52 + ], + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "raw_3DGS", + "type": "GS_RAW", + "link": 21 + } + ], + "properties": { + "Node name for S&R": "[ComfyUI-3D] Save 3DGS" + }, + "widgets_values": [ + "MeshTest\\3DGS_01.ply" + ] + }, + { + "id": 20, + "type": "[ComfyUI-3D] Convert 3DGS To Pointcloud", + "pos": [ + 2242, + 251 + ], + "size": { + "0": 267, + "1": 26 + }, + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "raw_3DGS", + "type": "GS_RAW", + "link": 24 + } + ], + "outputs": [ + { + "name": "points_cloud", + "type": "POINTCLOUD", + "links": [ + 25 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[ComfyUI-3D] Convert 3DGS To Pointcloud" + } + }, + { + "id": 13, + "type": "[ComfyUI-3D] Load 3D Mesh", + "pos": [ + 1053, + 415 + ], + "size": { + "0": 315, + "1": 154 + }, + "flags": {}, + "order": 1, + "mode": 0, + "outputs": [ + { + "name": "mesh", + "type": "MESH", + "links": [], + "shape": 3, + "slot_index": 0 + } + ], + "title": "Load 3D Mesh [Optional]", + "properties": { + "Node name for S&R": "[ComfyUI-3D] Load 3D Mesh" + }, + "widgets_values": [ + "Meshes\\Cammy_Uniform.obj", + false, + true, + false, + false + ] + }, + { + "id": 22, + "type": "[ComfyUI-3D] Save 3D Mesh", + "pos": [ + 3261, + 55 + ], + "size": { + "0": 353.38916015625, + "1": 61.879417419433594 + }, + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "mesh", + "type": "MESH", + "link": 26 + } + ], + "properties": { + "Node name for S&R": "[ComfyUI-3D] Save 3D Mesh" + }, + "widgets_values": [ + "MeshTest\\Cammy_Uniform_DMTet_01.obj" + ] + }, + { + "id": 21, + "type": "[ComfyUI-3D] Deep Marching Tetrahedrons", + "pos": [ + 2710, + 56 + ], + "size": { + "0": 443.4000244140625, + "1": 314 + }, + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "reference_points_cloud", + "type": "POINTCLOUD", + "link": 25 + }, + { + "name": "reference_images", + "type": "IMAGE", + "link": null + }, + { + "name": "reference_masks", + "type": "MASK", + "link": null + } + ], + "outputs": [ + { + "name": "trained_mesh", + "type": "MESH", + "links": [ + 26 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[ComfyUI-3D] Deep Marching Tetrahedrons" + }, + "widgets_values": [ + 5000, + 1, + 0.1, + 1, + 1, + 128, + 0.0001, + 2, + 128, + 5 + ] + } + ], + "links": [ + [ + 1, + 7, + 0, + 4, + 0, + "IMAGE" + ], + [ + 2, + 7, + 1, + 8, + 0, + "MASK" + ], + [ + 4, + 9, + 0, + 10, + 0, + "IMAGE" + ], + [ + 8, + 8, + 0, + 9, + 0, + "MASK" + ], + [ + 17, + 7, + 0, + 18, + 0, + "IMAGE" + ], + [ + 18, + 8, + 0, + 18, + 1, + "MASK" + ], + [ + 19, + 4, + 0, + 18, + 2, + "ORBIT_CAMPOSES" + ], + [ + 21, + 18, + 0, + 17, + 0, + "GS_RAW" + ], + [ + 22, + 7, + 0, + 19, + 0, + "IMAGE" + ], + [ + 24, + 18, + 0, + 20, + 0, + "GS_RAW" + ], + [ + 25, + 20, + 0, + 21, + 0, + "POINTCLOUD" + ], + [ + 26, + 21, + 0, + 22, + 0, + "MESH" + ] + ], + "groups": [], + "config": {}, + "extra": {}, + "version": 0.4 +} \ No newline at end of file diff --git a/_Example_Workflows/Load_3DGS_to_DMTetMesh.json b/_Example_Workflows/Load_3DGS_to_DMTetMesh.json new file mode 100644 index 0000000..06e2fab --- /dev/null +++ b/_Example_Workflows/Load_3DGS_to_DMTetMesh.json @@ -0,0 +1,190 @@ +{ + "last_node_id": 4, + "last_link_id": 3, + "nodes": [ + { + "id": 2, + "type": "[ComfyUI-3D] Convert 3DGS To Pointcloud", + "pos": [ + 561, + 484 + ], + "size": { + "0": 267, + "1": 26 + }, + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "raw_3DGS", + "type": "GS_RAW", + "link": 1 + } + ], + "outputs": [ + { + "name": "points_cloud", + "type": "POINTCLOUD", + "links": [ + 2 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[ComfyUI-3D] Convert 3DGS To Pointcloud" + } + }, + { + "id": 1, + "type": "[ComfyUI-3D] Load 3DGS", + "pos": [ + 166, + 484 + ], + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 0, + "mode": 0, + "outputs": [ + { + "name": "raw_3DGS", + "type": "GS_RAW", + "links": [ + 1 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[ComfyUI-3D] Load 3DGS" + }, + "widgets_values": [ + "Meshes\\Generated\\Cammy_Uniform_3DGS_01.ply" + ] + }, + { + "id": 4, + "type": "[ComfyUI-3D] Save 3D Mesh", + "pos": [ + 1446, + 485 + ], + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "mesh", + "type": "MESH", + "link": 3 + } + ], + "properties": { + "Node name for S&R": "[ComfyUI-3D] Save 3D Mesh" + }, + "widgets_values": [ + "MeshTest\\Cammy_Uniform_DMTet_01.obj" + ] + }, + { + "id": 3, + "type": "[ComfyUI-3D] Deep Marching Tetrahedrons", + "pos": [ + 913, + 484 + ], + "size": { + "0": 443.4000244140625, + "1": 314 + }, + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "reference_points_cloud", + "type": "POINTCLOUD", + "link": 2 + }, + { + "name": "reference_images", + "type": "IMAGE", + "link": null + }, + { + "name": "reference_masks", + "type": "MASK", + "link": null + } + ], + "outputs": [ + { + "name": "trained_mesh", + "type": "MESH", + "links": [ + 3 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[ComfyUI-3D] Deep Marching Tetrahedrons" + }, + "widgets_values": [ + 5000, + 1, + 0.1, + 1, + 1, + 128, + 0.0001, + 2, + 128, + 5 + ] + } + ], + "links": [ + [ + 1, + 1, + 0, + 2, + 0, + "GS_RAW" + ], + [ + 2, + 2, + 0, + 3, + 0, + "POINTCLOUD" + ], + [ + 3, + 3, + 0, + 4, + 0, + "MESH" + ] + ], + "groups": [], + "config": {}, + "extra": {}, + "version": 0.4 +} \ No newline at end of file diff --git a/algorithms/_Test.py b/algorithms/_Test.py new file mode 100644 index 0000000..3d5a616 --- /dev/null +++ b/algorithms/_Test.py @@ -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) + diff --git a/diff_rast/diff_mesh_renderer.py b/algorithms/diff_mesh_renderer.py similarity index 100% rename from diff_rast/diff_mesh_renderer.py rename to algorithms/diff_mesh_renderer.py diff --git a/diff_rast/diff_texturing.py b/algorithms/diff_texturing.py similarity index 93% rename from diff_rast/diff_texturing.py rename to algorithms/diff_texturing.py index 2dff5e3..0f559c2 100644 --- a/diff_rast/diff_texturing.py +++ b/algorithms/diff_texturing.py @@ -6,25 +6,14 @@ import torch.nn.functional as F import numpy as np from pytorch_msssim import SSIM, MS_SSIM -from ..diff_rast.diff_mesh_renderer import Renderer +from .diff_mesh_renderer import Renderer from ..shared_utils.camera_utils import orbit_camera, OrbitCamera class DiffTextureBaker: - def __init__(self, reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy, mesh, - training_iterations, batch_size, texture_learning_rate, train_mesh_geometry, geometry_learning_rate, ms_ssim_loss_weight, force_cuda_rasterize): + 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") - 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 main components for optimization self.renderer = Renderer(mesh, force_cuda_rasterize).to(self.device) @@ -36,6 +25,22 @@ class DiffTextureBaker: 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 = [] @@ -47,11 +52,6 @@ class DiffTextureBaker: 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 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 training(self): starter = torch.cuda.Event(enable_timing=True) ender = torch.cuda.Event(enable_timing=True) diff --git a/algorithms/dmtet.py b/algorithms/dmtet.py new file mode 100644 index 0000000..6f68258 --- /dev/null +++ b/algorithms/dmtet.py @@ -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() \ No newline at end of file diff --git a/algorithms/dmtet_network.py b/algorithms/dmtet_network.py new file mode 100644 index 0000000..715c950 --- /dev/null +++ b/algorithms/dmtet_network.py @@ -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 diff --git a/gaussian_splatting/main_3DGS.py b/algorithms/main_3DGS.py similarity index 97% rename from gaussian_splatting/main_3DGS.py rename to algorithms/main_3DGS.py index f3ee6d0..4564c55 100644 --- a/gaussian_splatting/main_3DGS.py +++ b/algorithms/main_3DGS.py @@ -73,20 +73,8 @@ class GSParams: class GaussianSplatting: - def __init__(self, reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy, gs_params=None, init_input=None): - self.device = torch.device("cuda") - - 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() + 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: @@ -106,6 +94,24 @@ class GaussianSplatting: 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 = [] @@ -115,11 +121,6 @@ class GaussianSplatting: 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 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 training(self): starter = torch.cuda.Event(enable_timing=True) diff --git a/gaussian_splatting/main_3DGS_renderer.py b/algorithms/main_3DGS_renderer.py similarity index 96% rename from gaussian_splatting/main_3DGS_renderer.py rename to algorithms/main_3DGS_renderer.py index a637fad..65999c0 100644 --- a/gaussian_splatting/main_3DGS_renderer.py +++ b/algorithms/main_3DGS_renderer.py @@ -2,7 +2,7 @@ import os import math import random import numpy as np -from typing import NamedTuple, Optional +from typing import Optional from torchtyping import TensorType from plyfile import PlyData, PlyElement @@ -10,6 +10,8 @@ import torch from torch import nn import torch.nn.functional as F +from ..mesh_processer.mesh import PointCloud + pytorch3d_capable = True try: import pytorch3d @@ -255,11 +257,6 @@ def distance_to_gaussian_surface(points, svec, rotmat, query): def qvec2rotmat_batched(qvec: TensorType["N", 4]): return quaternion_to_rotation_matrix(qvec) -class BasicPointCloud(NamedTuple): - points: np.array - colors: np.array - normals: np.array - class GaussianModel: def setup_functions(self): @@ -440,34 +437,6 @@ class GaussianModel: return occ - def extract_mesh(self, path, density_thresh=1, resolution=128, decimate_target=1e5): - - os.makedirs(os.path.dirname(path), exist_ok=True) - - occ = self.extract_fields(resolution).detach().cpu().numpy() - - import mcubes - vertices, triangles = mcubes.marching_cubes(occ, density_thresh) - vertices = vertices / (resolution - 1.0) * 2 - 1 - - # transform back to the original space - vertices = vertices / self.scale + self.center.detach().cpu().numpy() - - vertices, triangles = clean_mesh(vertices, triangles, remesh=True, remesh_size=0.015) - if decimate_target > 0 and triangles.shape[0] > decimate_target: - vertices, triangles = decimate_mesh(vertices, triangles, decimate_target) - - v = torch.from_numpy(vertices.astype(np.float32)).contiguous().cuda() - f = torch.from_numpy(triangles.astype(np.int32)).contiguous().cuda() - - print( - f"[INFO] marching cubes result: {v.shape} ({v.min().item()}-{v.max().item()}), {f.shape}" - ) - - mesh = Mesh(v=v, f=f, device='cuda') - - return mesh - def get_covariance(self, scaling_modifier = 1, gaussain_idx = None): if gaussain_idx is None: return self.covariance_activation(self.get_scaling, scaling_modifier, self._rotation) @@ -478,7 +447,7 @@ class GaussianModel: if self.active_sh_degree < self.max_sh_degree: self.active_sh_degree += 1 - def create_from_pcd(self, pcd : BasicPointCloud, spatial_lr_scale : float = 1): + def create_from_pcd(self, pcd : PointCloud, spatial_lr_scale : float = 1): self.spatial_lr_scale = spatial_lr_scale fused_point_cloud = torch.tensor(np.asarray(pcd.points)).float().cuda() fused_color = RGB2SH(torch.tensor(np.asarray(pcd.colors)).float().cuda()) @@ -545,6 +514,16 @@ class GaussianModel: for i in range(self._rotation.shape[1]): l.append('rot_{}'.format(i)) return l + + def get_points_cloud(self): + xyz = self._xyz.detach().cpu().numpy() + normals = np.zeros_like(xyz) + shs = self.get_features.detach().cpu().numpy() + pcd = PointCloud(points=xyz, colors=SH2RGB(shs), normals=normals) + return pcd + + def get_mesh(self): + xyz = self._xyz.detach().cpu().numpy() def save_ply(self, path): os.makedirs(os.path.dirname(path), exist_ok=True) @@ -633,7 +612,7 @@ class GaussianModel: xyz = np.array(all_points) shs = np.random.random((num_pts, 3)) / 255.0 - pcd = BasicPointCloud( + pcd = PointCloud( points=xyz, colors=SH2RGB(shs), normals=np.zeros((num_pts, 3)) ) self.create_from_pcd(pcd, 10) @@ -650,7 +629,7 @@ class GaussianModel: num_pts = len(uv_coords_3d) xyz = np.array(uv_coords_3d) shs = np.random.random((num_pts, 3)) / 255.0 - pcd = BasicPointCloud( + pcd = PointCloud( points=xyz, colors=SH2RGB(shs), normals=np.zeros((num_pts, 3)) ) self.create_from_pcd(pcd, 10) @@ -930,11 +909,11 @@ class Renderer: # xyz = np.random.random((num_pts, 3)) * 2.6 - 1.3 shs = np.random.random((num_pts, 3)) / 255.0 - pcd = BasicPointCloud( + pcd = PointCloud( points=xyz, colors=SH2RGB(shs), normals=np.zeros((num_pts, 3)) ) self.gaussians.create_from_pcd(pcd, 10) - elif isinstance(input, BasicPointCloud): + elif isinstance(input, PointCloud): # load from a provided pcd self.gaussians.create_from_pcd(input, 1) else: diff --git a/algorithms/metrics_utils.py b/algorithms/metrics_utils.py new file mode 100644 index 0000000..57a5074 --- /dev/null +++ b/algorithms/metrics_utils.py @@ -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 \ No newline at end of file diff --git a/algorithms/tetmesh.py b/algorithms/tetmesh.py new file mode 100644 index 0000000..5283306 --- /dev/null +++ b/algorithms/tetmesh.py @@ -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)) diff --git a/data/samples_tets/128_tets_0.npz b/data/samples_tets/128_tets_0.npz new file mode 100644 index 0000000..8684f02 Binary files /dev/null and b/data/samples_tets/128_tets_0.npz differ diff --git a/data/samples_tets/128_tets_1.npz b/data/samples_tets/128_tets_1.npz new file mode 100644 index 0000000..e7f1c97 Binary files /dev/null and b/data/samples_tets/128_tets_1.npz differ diff --git a/data/samples_tets/128_tets_2.npz b/data/samples_tets/128_tets_2.npz new file mode 100644 index 0000000..07e5202 Binary files /dev/null and b/data/samples_tets/128_tets_2.npz differ diff --git a/data/samples_tets/128_tets_3.npz b/data/samples_tets/128_tets_3.npz new file mode 100644 index 0000000..3e0aa5e Binary files /dev/null and b/data/samples_tets/128_tets_3.npz differ diff --git a/data/samples_tets/128_verts.npz b/data/samples_tets/128_verts.npz new file mode 100644 index 0000000..ab071cc Binary files /dev/null and b/data/samples_tets/128_verts.npz differ diff --git a/data/tets/128_tets.npz b/data/tets/128_tets.npz new file mode 100644 index 0000000..156b7fd Binary files /dev/null and b/data/tets/128_tets.npz differ diff --git a/data/tets/32_tets.npz b/data/tets/32_tets.npz new file mode 100644 index 0000000..cffeeba Binary files /dev/null and b/data/tets/32_tets.npz differ diff --git a/data/tets/64_tets.npz b/data/tets/64_tets.npz new file mode 100644 index 0000000..8515e35 Binary files /dev/null and b/data/tets/64_tets.npz differ diff --git a/diff_rast/__init__.py b/diff_rast/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/mesh_processer/mesh.py b/mesh_processer/mesh.py index d19ec9f..43d6935 100644 --- a/mesh_processer/mesh.py +++ b/mesh_processer/mesh.py @@ -1,9 +1,13 @@ import os import cv2 +from typing import NamedTuple import torch import trimesh import numpy as np + +from ..shared_utils.sh_utils import SH2RGB + def dot(x, y): return torch.sum(x * y, -1, keepdim=True) @@ -42,6 +46,8 @@ class Mesh: self.ori_center = 0 self.ori_scale = 1 + + @classmethod def load(cls, path=None, resize=True, renormal=True, retex=False, front_dir='+z', **kwargs): @@ -615,8 +621,24 @@ class Mesh: fp.write(f"Tr 1 \n") fp.write(f"illum 1 \n") fp.write(f"Ns 0 \n") - fp.write(f"map_Kd {os.path.basename(albedo_path)} \n") + if self.albedo is not None: + fp.write(f"map_Kd {os.path.basename(albedo_path)} \n") - albedo = self.albedo.detach().cpu().numpy() - albedo = (albedo * 255).astype(np.uint8) - cv2.imwrite(albedo_path, cv2.cvtColor(albedo, cv2.COLOR_RGB2BGR)) \ No newline at end of file + if self.albedo is not None: + albedo = self.albedo.detach().cpu().numpy() + albedo = (albedo * 255).astype(np.uint8) + cv2.imwrite(albedo_path, cv2.cvtColor(albedo, cv2.COLOR_RGB2BGR)) + + def convert_to_pointcloud(self): + xyz = self.v.detach().cpu().numpy() + num_pts = self.v.shape[0] + shs = np.random.random((num_pts, 3)) / 255.0 + normals = np.zeros((num_pts, 3)) + pcd = PointCloud(points=xyz, colors=SH2RGB(shs), normals=normals) + return pcd + + +class PointCloud(NamedTuple): + points: np.array + colors: np.array + normals: np.array \ No newline at end of file diff --git a/mesh_processer/mesh_utils.py b/mesh_processer/mesh_utils.py index ca9fce9..f454dad 100644 --- a/mesh_processer/mesh_utils.py +++ b/mesh_processer/mesh_utils.py @@ -1,6 +1,158 @@ +import torch import numpy as np import pymeshlab as pml +def _base_face_areas(face_vertices_0, face_vertices_1, face_vertices_2): + """Base function to compute the face areas.""" + x1, x2, x3 = torch.split(face_vertices_0 - face_vertices_1, 1, dim=-1) + y1, y2, y3 = torch.split(face_vertices_1 - face_vertices_2, 1, dim=-1) + + a = (x2 * y3 - x3 * y2) ** 2 + b = (x3 * y1 - x1 * y3) ** 2 + c = (x1 * y2 - x2 * y1) ** 2 + areas = torch.sqrt(a + b + c) * 0.5 + + return areas + + +def _base_sample_points_selected_faces(face_vertices, face_features=None): + """Base function to sample points over selected faces, sample one point per face. + The coordinates of the face vertices are interpolated to generate new samples. + + Args: + face_vertices (tuple of torch.Tensor): + Coordinates of vertices, corresponding to selected faces to sample from. + A tuple of 3 entries corresponding to each of the face vertices. + Each entry is a torch.Tensor of shape :math:`(\\text{batch_size}, \\text{num_samples}, 3)`. + face_features (tuple of torch.Tensor, Optional): + Features of face vertices, corresponding to selected faces to sample from. + A tuple of 3 entries corresponding to each of the face vertices. + Each entry is a torch.Tensor of shape + :math:`(\\text{batch_size}, \\text{num_samples}, \\text{feature_dim})`. + + Returns: + (torch.Tensor, torch.Tensor): + Sampled point coordinates of shape :math:`(\\text{batch_size}, \\text{num_samples}, 3)`. + Sampled points interpolated features of shape + :math:`(\\text{batch_size}, \\text{num_samples}, \\text{feature_dim})`. + If `face_vertices_features` arg is not specified, the returned interpolated features are None. + """ + + face_vertices0, face_vertices1, face_vertices2 = face_vertices + + sampling_shape = tuple(int(d) for d in face_vertices0.shape[:-1]) + (1,) + # u is proximity to middle point between v1 and v2 against v0. + # v is proximity to v2 against v1. + # + # The probability density for u should be f_U(u) = 2u. + # However, torch.rand use a uniform (f_X(x) = x) distribution, + # so using torch.sqrt we make a change of variable to have the desired density + # f_Y(y) = f_X(y ^ 2) * |d(y ^ 2) / dy| = 2y + u = torch.sqrt(torch.rand(sampling_shape, + device=face_vertices0.device, + dtype=face_vertices0.dtype)) + + v = torch.rand(sampling_shape, + device=face_vertices0.device, + dtype=face_vertices0.dtype) + w0 = 1 - u + w1 = u * (1 - v) + w2 = u * v + + points = w0 * face_vertices0 + w1 * face_vertices1 + w2 * face_vertices2 + + features = None + if face_features is not None: + face_features0, face_features1, face_features2 = face_features + features = w0 * face_features0 + w1 * face_features1 + \ + w2 * face_features2 + + return points, features + +# Modified from https://github.com/NVIDIAGameWorks/kaolin/blob/master/kaolin/ops/mesh/trianglemesh.py#L159 +def sample_points(vertices, faces, num_samples, areas=None, face_features=None): + r"""Uniformly sample points over the surface of triangle meshes. + + First, face on which the point is sampled is randomly selected, + with the probability of selection being proportional to the area of the face. + then the coordinate on the face is uniformly sampled. + + If ``face_features`` is defined for the mesh faces, + the sampled points will be returned with interpolated features as well, + otherwise, no feature interpolation will occur. + + Args: + vertices (torch.Tensor): + The vertices of the meshes, of shape + :math:`(\text{batch_size}, \text{num_vertices}, 3)`. + faces (torch.LongTensor): + The faces of the mesh, of shape :math:`(\text{num_faces}, 3)`. + num_samples (int): + The number of point sampled per mesh. + Also the number of faces sampled per mesh, and then sample a single point per face. + areas (torch.Tensor, optional): + The areas of each face, of shape :math:`(\text{batch_size}, \text{num_faces})`, + can be preprocessed, for fast on-the-fly sampling, + will be computed if None (default). + face_features (torch.Tensor, optional): + Per-vertex-per-face features, matching ``faces`` order, + of shape :math:`(\text{batch_size}, \text{num_faces}, 3, \text{feature_dim})`. + For example: + + 1. Texture uv coordinates would be of shape + :math:`(\text{batch_size}, \text{num_faces}, 3, 2)`. + 2. RGB color values would be of shape + :math:`(\text{batch_size}, \text{num_faces}, 3, 3)`. + + When specified, it is used to interpolate the features for new sampled points. + + See also: + :func:`~kaolin.ops.mesh.index_vertices_by_faces` for conversion of features defined per vertex + and need to be converted to per-vertex-per-face shape of :math:`(\text{num_faces}, 3)`. + + Returns: + (torch.Tensor, torch.LongTensor, (optional) torch.Tensor): + the pointclouds of shape :math:`(\text{batch_size}, \text{num_samples}, 3)`, + and the indexes of the faces selected, + of shape :math:`(\text{batch_size}, \text{num_samples})`. + + If ``face_features`` arg is specified, then the interpolated features of sampled points of shape + :math:`(\text{batch_size}, \text{num_samples}, \text{feature_dim})` are also returned. + """ + if faces.shape[-1] != 3: + raise NotImplementedError("sample_points is only implemented for triangle meshes") + faces_0, faces_1, faces_2 = torch.split(faces, 1, dim=1) # (num_faces, 3) -> tuple of (num_faces,) + face_v_0 = torch.index_select(vertices, 1, faces_0.reshape(-1)) # (batch_size, num_faces, 3) + face_v_1 = torch.index_select(vertices, 1, faces_1.reshape(-1)) # (batch_size, num_faces, 3) + face_v_2 = torch.index_select(vertices, 1, faces_2.reshape(-1)) # (batch_size, num_faces, 3) + + if areas is None: + areas = _base_face_areas(face_v_0, face_v_1, face_v_2).squeeze(-1) + face_dist = torch.distributions.Categorical(areas) + face_choices = face_dist.sample([num_samples]).transpose(0, 1) + _face_choices = face_choices.unsqueeze(-1).repeat(1, 1, 3) + v0 = torch.gather(face_v_0, 1, _face_choices) # (batch_size, num_samples, 3) + v1 = torch.gather(face_v_1, 1, _face_choices) # (batch_size, num_samples, 3) + v2 = torch.gather(face_v_2, 1, _face_choices) # (batch_size, num_samples, 3) + face_vertices_choices = (v0, v1, v2) + + # UV coordinates are available, make sure to calculate them for sampled points as well + face_features_choices = None + if face_features is not None: + feat_dim = face_features.shape[-1] + # (num_faces, 3) -> tuple of (num_faces,) + _face_choices = face_choices[..., None, None].repeat(1, 1, 3, feat_dim) + face_features_choices = torch.gather(face_features, 1, _face_choices) + face_features_choices = tuple( + tmp_feat.squeeze(2) for tmp_feat in torch.split(face_features_choices, 1, dim=2)) + + points, point_features = _base_sample_points_selected_faces( + face_vertices_choices, face_features_choices) + + if point_features is not None: + return points, face_choices, point_features + else: + return points, face_choices def poisson_mesh_reconstruction(points, normals=None): # points/normals: [N, 3] np.ndarray diff --git a/nodes.py b/nodes.py index 4a08b4c..86d0822 100644 --- a/nodes.py +++ b/nodes.py @@ -8,10 +8,11 @@ import folder_paths as comfy_paths import torch import numpy as np -from .diff_rast.diff_texturing import DiffTextureBaker from .shared_utils.common_utils import cstr from .mesh_processer.mesh import Mesh -from .gaussian_splatting.main_3DGS import GaussianSplatting, GSParams +from .algorithms.main_3DGS import GaussianSplatting, GSParams +from .algorithms.diff_texturing import DiffTextureBaker +from .algorithms.dmtet import DMTetMesh MANIFEST = { "name": "ComfyUI-3D-Pack", @@ -27,6 +28,10 @@ SUPPORTED_3D_EXTENSIONS = ( '.glb', ) +SUPPORTED_3DGS_EXTENSIONS = ( + '.ply', +) + ELEVATION_MIN = -90 ELEVATION_MAX = 90.0 AZIMUTH_MIN = -180.0 @@ -72,6 +77,45 @@ class Load_3D_Mesh: cstr(f"[{self.__class__.__name__}] File {mesh_file_path} does not exist").error.print() return (mesh, ) +class Load_3DGS: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "raw_3DGS_file_path": ("STRING", {"default": '', "multiline": False}), + }, + } + + RETURN_TYPES = ( + "GS_RAW", + + ) + RETURN_NAMES = ( + "raw_3DGS", + ) + FUNCTION = "load_mesh" + CATEGORY = "ComfyUI3D/Import|Export" + + def load_mesh(self, raw_3DGS_file_path): + raw_3DGS = None + + if not os.path.isabs(raw_3DGS_file_path): + raw_3DGS_file_path = os.path.join(comfy_paths.input_directory, raw_3DGS_file_path) + + if os.path.exists(raw_3DGS_file_path): + folder, filename = os.path.split(raw_3DGS_file_path) + if filename.lower().endswith(SUPPORTED_3DGS_EXTENSIONS): + with torch.inference_mode(False): + raw_3DGS = GaussianSplatting() + raw_3DGS.renderer.gaussians.load_ply(raw_3DGS_file_path) + # TODO: better way of loading 3DGS, since this method omit render & training parameters + else: + cstr(f"[{self.__class__.__name__}] File name {filename} does not end with supported 3DGS file extensions: {SUPPORTED_3DGS_EXTENSIONS}").error.print() + else: + cstr(f"[{self.__class__.__name__}] File {raw_3DGS_file_path} does not exist").error.print() + return (raw_3DGS, ) + class Save_3D_Mesh: @classmethod @@ -126,16 +170,70 @@ class Save_3DGS: folder_path, filename = os.path.split(save_path) - if not os.path.isabs(save_path): - folder_path = os.path.join(comfy_paths.output_directory, folder_path) - save_path = os.path.join(folder_path, filename) - - os.makedirs(folder_path, exist_ok=True) - - raw_3DGS.renderer.gaussians.save_ply(save_path) + if filename.lower().endswith(SUPPORTED_3DGS_EXTENSIONS): + if not os.path.isabs(save_path): + folder_path = os.path.join(comfy_paths.output_directory, folder_path) + save_path = os.path.join(folder_path, filename) + + os.makedirs(folder_path, exist_ok=True) + + raw_3DGS.renderer.gaussians.save_ply(save_path) + # TODO: better way of saving 3DGS, since this method omit render & training parameters + else: + cstr(f"[{self.__class__.__name__}] File name {filename} does not end with supported 3DGS file extensions: {SUPPORTED_3DGS_EXTENSIONS}").error.print() return () +class Convert_3DGS_To_Pointcloud: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "raw_3DGS": ("GS_RAW",), + }, + } + + RETURN_TYPES = ( + "POINTCLOUD", + ) + RETURN_NAMES = ( + "points_cloud", + ) + FUNCTION = "convert_3DGS" + CATEGORY = "ComfyUI3D/Preprocessor" + + def convert_3DGS(self, raw_3DGS): + + points_cloud = raw_3DGS.renderer.gaussians.get_points_cloud() + + return (points_cloud, ) + +class Convert_Mesh_To_Pointcloud: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mesh": ("MESH",), + }, + } + + RETURN_TYPES = ( + "POINTCLOUD", + ) + RETURN_NAMES = ( + "points_cloud", + ) + FUNCTION = "convert_mesh" + CATEGORY = "ComfyUI3D/Preprocessor" + + def convert_mesh(self, mesh): + + points_cloud = mesh.convert_to_pointcloud() + + return (points_cloud, ) + class Stack_Orbit_Camera_Poses: @classmethod @@ -413,29 +511,29 @@ class Gaussian_Splatting: "reference_orbit_camera_fovy": ("FLOAT", {"default": 49.1, "min": 0.0, "max": 180.0, "step": 0.1}), "training_iterations": ("INT", {"default": 1000, "min": 1, "max": 100000}), "batch_size": ("INT", {"default": 5, "min": 1, "max": 0xffffffffffffffff}), - "loss_value_scale": ("FLOAT", {"default": 10000.0, }), - "ms_ssim_loss_weight": ("FLOAT", {"default": 0.2, }), - "alpha_loss_weight": ("FLOAT", {"default": 3, }), - "offset_loss_weight": ("FLOAT", {"default": 0.0, }), - "offset_opacity_loss_weight": ("FLOAT", {"default": 0.0, }), + "loss_value_scale": ("FLOAT", {"default": 10000.0, "min": 1.0}), + "ms_ssim_loss_weight": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, }), + "alpha_loss_weight": ("FLOAT", {"default": 3, "min": 0.0, }), + "offset_loss_weight": ("FLOAT", {"default": 0.0, "min": 0.0, }), + "offset_opacity_loss_weight": ("FLOAT", {"default": 0.0, "min": 0.0, }), "invert_background_probability": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.1}), - "feature_learning_rate": ("FLOAT", {"default": 0.01, }), - "opacity_learning_rate": ("FLOAT", {"default": 0.05, }), - "scaling_learning_rate": ("FLOAT", {"default": 0.005, }), - "rotation_learning_rate": ("FLOAT", {"default": 0.005, }), - "position_learning_rate_init": ("FLOAT", {"default": 0.001, }), - "position_learning_rate_final": ("FLOAT", {"default": 0.00002, }), - "position_learning_rate_delay_mult": ("FLOAT", {"default": 0.02, }), - "position_learning_rate_max_steps": ("INT", {"default": 500, }), - "initial_gaussians_num": ("INT", {"default": 5000, }), - "K_nearest_neighbors": ("INT", {"default": 3, }), - "percent_dense": ("FLOAT", {"default": 0.01, }), - "density_start_iterations": ("INT", {"default": 100, }), - "density_end_iterations": ("INT", {"default": 100000, }), - "densification_interval": ("INT", {"default": 100, }), - "opacity_reset_interval": ("INT", {"default": 700, }), - "densify_grad_threshold": ("FLOAT", {"default": 0.01 }), - "gaussian_sh_degree": ("INT", {"default": 0, }), + "feature_learning_rate": ("FLOAT", {"default": 0.01, "min": 0.00001, "step": 0.00001}), + "opacity_learning_rate": ("FLOAT", {"default": 0.05, "min": 0.00001, "step": 0.00001}), + "scaling_learning_rate": ("FLOAT", {"default": 0.005, "min": 0.00001, "step": 0.00001}), + "rotation_learning_rate": ("FLOAT", {"default": 0.005, "min": 0.00001, "step": 0.00001}), + "position_learning_rate_init": ("FLOAT", {"default": 0.001, "min": 0.00001, "step": 0.00001}), + "position_learning_rate_final": ("FLOAT", {"default": 0.00002, "min": 0.00001, "step": 0.00001}), + "position_learning_rate_delay_mult": ("FLOAT", {"default": 0.02, "min": 0.00001, "step": 0.00001}), + "position_learning_rate_max_steps": ("INT", {"default": 500, "min": 1}), + "initial_gaussians_num": ("INT", {"default": 5000, "min": 1}), + "K_nearest_neighbors": ("INT", {"default": 3, "min": 1}), + "percent_dense": ("FLOAT", {"default": 0.01, "min": 0.00001, "step": 0.00001}), + "density_start_iterations": ("INT", {"default": 100, "min": 0}), + "density_end_iterations": ("INT", {"default": 100000, "min": 0}), + "densification_interval": ("INT", {"default": 100, "min": 1}), + "opacity_reset_interval": ("INT", {"default": 700, "min": 1}), + "densify_grad_threshold": ("FLOAT", {"default": 0.01, "min": 0.00001, "step": 0.00001}), + "gaussian_sh_degree": ("INT", {"default": 0, "min": 0}), }, "optional": { @@ -526,8 +624,8 @@ class Gaussian_Splatting: densify_grad_threshold, gaussian_sh_degree) - gs = GaussianSplatting(reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy, gs_params, mesh_to_initialize_gaussian) - + gs = GaussianSplatting(gs_params, mesh_to_initialize_gaussian) + gs.prepare_training(reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy) gs.training() raw_3DGS = gs @@ -555,9 +653,9 @@ class Bake_Texture_To_Mesh: "mesh": ("MESH",), "training_iterations": ("INT", {"default": 1000, "min": 1, "max": 100000}), "batch_size": ("INT", {"default": 5, "min": 1, "max": 0xffffffffffffffff}), - "texture_learning_rate": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01}), + "texture_learning_rate": ("FLOAT", {"default": 0.1, "min": 0.00001, "step": 0.00001}), "train_mesh_geometry": ("BOOLEAN", {"default": False},), - "geometry_learning_rate": ("FLOAT", {"default": 0.01, "min": 0.0, "max": 0.1, "step": 0.001}), + "geometry_learning_rate": ("FLOAT", {"default": 0.01, "min": 0.00001, "step": 0.00001}), "ms_ssim_loss_weight": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), "force_cuda_rasterize": ("BOOLEAN", {"default": False},), }, @@ -593,9 +691,9 @@ class Bake_Texture_To_Mesh: with torch.inference_mode(False): - texture_baker = DiffTextureBaker(reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy, mesh, - training_iterations, batch_size, texture_learning_rate, train_mesh_geometry, geometry_learning_rate, ms_ssim_loss_weight, force_cuda_rasterize) + texture_baker = DiffTextureBaker(mesh, training_iterations, batch_size, texture_learning_rate, train_mesh_geometry, geometry_learning_rate, ms_ssim_loss_weight, force_cuda_rasterize) + texture_baker.prepare_training(reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy) texture_baker.training() trained_mesh, baked_texture = texture_baker.get_mesh_and_texture() @@ -605,4 +703,62 @@ class Bake_Texture_To_Mesh: else: cstr(f"[{self.__class__.__name__}] Number of reference images {ref_imgs_num} does not equal to number of masks {ref_masks_num}").error.print() - return (trained_mesh, baked_texture, ) \ No newline at end of file + return (trained_mesh, baked_texture, ) + +class Deep_Marching_Tetrahedrons: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "training_iterations": ("INT", {"default": 5000, "min": 1, "max": 100000}), + "points_cloud_fitting_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.01}), + "mesh_smoothing_weight": ("FLOAT", {"default": 0.1, "min": 0.0, "step": 0.01}), + "chamfer_faces_sample_scale": ("FLOAT", {"default": 1.0, "min": 0.01, "step": 0.01}), + "mesh_scale": ("FLOAT", {"default": 1.0, "min": 0.01, "step": 0.01}), + "grid_resolution": ([128, 64, 32], ), + "geometry_learning_rate": ("FLOAT", {"default": 0.0001, "min": 0.00001, "step": 0.00001}), + "positional_encoding_multires": ("INT", {"default": 2, "min": 2}), + "mlp_internal_dims": ("INT", {"default": 128, "min": 8}), + "mlp_hidden_layer_num": ("INT", {"default": 5, "min": 1}), + }, + + "optional": { + "reference_points_cloud": ("POINTCLOUD",), + "reference_images": ("IMAGE",), + "reference_masks": ("MASK",), + } + } + + RETURN_TYPES = ( + "MESH", + ) + RETURN_NAMES = ( + "trained_mesh", + ) + FUNCTION = "run_dmtet" + CATEGORY = "ComfyUI3D/Algorithm" + + def run_dmtet(self, training_iterations, points_cloud_fitting_weight, mesh_smoothing_weight, chamfer_faces_sample_scale, mesh_scale, grid_resolution, geometry_learning_rate, + positional_encoding_multires, mlp_internal_dims, mlp_hidden_layer_num, + reference_points_cloud=None, reference_images=None, reference_masks=None): + + with torch.inference_mode(False): + dmtet_mesh = DMTetMesh(training_iterations, + points_cloud_fitting_weight, + mesh_smoothing_weight, + chamfer_faces_sample_scale, + mesh_scale, + grid_resolution, + geometry_learning_rate, + positional_encoding_multires, + mlp_internal_dims, + mlp_hidden_layer_num) + + if reference_points_cloud is None and reference_images is None: + cstr(f"[{self.__class__.__name__}] reference_points_cloud and reference_images cannot both be None, you need at least provide one of them!").error.print() + raise Exception("User didn't provide necessary inputs, stop executing the code!") + + dmtet_mesh.training(reference_points_cloud, reference_images, reference_masks) + mesh = dmtet_mesh.get_mesh() + return (mesh, ) \ No newline at end of file