nodes/MrForExample-ComfyUI-3D-Pack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add Deep Marching Tetrahedra (DMTet)

Browse Source
This commit is contained in:
MrForExample
2024-01-12 21:43:44 +01:00
parent 01b0a25138
commit 49fed463a2
24 changed files with 1773 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+7 -1
View File
@@ -72,6 +72,7 @@ Copy the files inside folder [__New_ComfyUI_Bats](./_New_ComfyUI_Bats/) to your
<img src="_Example_Workflows/_Example_Outputs/Cammy_Cam_Rotate_Clockwise_Camposes.png" width="256"/> <img src="_Example_Workflows/_Example_Outputs/Cammy_Cam_Rotate_Counter_Clockwise_Camposes.png" width="256"/>
<img src="_Example_Workflows/_Example_Outputs/Cammy_Cam_Rotate_Clockwise.gif" width="256"/> <img src="_Example_Workflows/_Example_Outputs/Cammy_Cam_Rotate_Counter_Clockwise.gif" width="256"/>
- 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 <br>
*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?)
_Example_Workflows/Images_to_3DGS_to_DMTetMesh.json
+602
View File
@@ -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
}
_Example_Workflows/Load_3DGS_to_DMTetMesh.json
+190
View File
@@ -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
}
algorithms/_Test.py
+8
View File
@@ -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_rast/diff_mesh_renderer.py → algorithms/diff_mesh_renderer.py
View File
diff_rast/diff_texturing.py → algorithms/diff_texturing.py
+18 -18
View File
@@ -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)
algorithms/dmtet.py
+271
View File
@@ -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()
algorithms/dmtet_network.py
+93
View File
@@ -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
gaussian_splatting/main_3DGS.py → algorithms/main_3DGS.py
+20 -19
View File
@@ -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)
gaussian_splatting/main_3DGS_renderer.py → algorithms/main_3DGS_renderer.py
+18 -39
View File
@@ -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:
algorithms/metrics_utils.py
+9
View File
@@ -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
algorithms/tetmesh.py
+165
View File
@@ -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))
data/samples_tets/128_tets_0.npz
BIN
View File
Binary file not shown.
data/samples_tets/128_tets_1.npz
BIN
View File
Binary file not shown.
data/samples_tets/128_tets_2.npz
BIN
View File
Binary file not shown.
data/samples_tets/128_tets_3.npz
BIN
View File
Binary file not shown.
data/samples_tets/128_verts.npz
BIN
View File
Binary file not shown.
data/tets/128_tets.npz
BIN
View File
Binary file not shown.
data/tets/32_tets.npz
BIN
View File
Binary file not shown.
data/tets/64_tets.npz
BIN
View File
Binary file not shown.
diff_rast/__init__.py
View File
mesh_processer/mesh.py
+26 -4
View File
@@ -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))
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
mesh_processer/mesh_utils.py
+152
View File
@@ -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
nodes.py
+194 -38
View File
@@ -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, )
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, )