added support for textured mesh

README.md

@@ -82,6 +82,10 @@ torchrun --nproc_per_node=${GPUS} projects/neuralangelo/scripts/extract_mesh.py
     --resolution=${RESOLUTION} \
     --block_res=${BLOCK_RES}
 ```
+Some useful notes:
+- Add `--textured` to extract meshes with textures.
+- Lower `BLOCK_RES` to reduce GPU memory usage.
+- Lower `RESOLUTION` to reduce mesh size.
 
 --------------------------------------
 

projects/neuralangelo/scripts/extract_mesh.py

@@ -15,26 +15,27 @@
 import os
 import sys
 import numpy as np
+from functools import partial
 
 sys.path.append(os.getcwd())
 from imaginaire.config import Config, recursive_update_strict, parse_cmdline_arguments  # noqa: E402
 from imaginaire.utils.distributed import init_dist, get_world_size, is_master, master_only_print as print  # noqa: E402
 from imaginaire.utils.gpu_affinity import set_affinity  # noqa: E402
 from imaginaire.trainers.utils.logging import init_logging  # noqa: E402
 from imaginaire.trainers.utils.get_trainer import get_trainer  # noqa: E402
-from projects.neuralangelo.utils.mesh import extract_mesh  # noqa: E402
+from projects.neuralangelo.utils.mesh import extract_mesh, extract_texture  # noqa: E402
 
 
 def parse_args():
     parser = argparse.ArgumentParser(description="Training")
     parser.add_argument("--config", required=True, help="Path to the training config file.")
-    parser.add_argument("--logdir", help="Dir for saving logs and models.")
     parser.add_argument("--checkpoint", default="", help="Checkpoint path.")
     parser.add_argument('--local_rank', type=int, default=os.getenv('LOCAL_RANK', 0))
     parser.add_argument('--single_gpu', action='store_true')
     parser.add_argument("--resolution", default=512, type=int, help="Marching cubes resolution")
     parser.add_argument("--block_res", default=64, type=int, help="Block-wise resolution for marching cubes")
     parser.add_argument("--output_file", default="mesh.ply", type=str, help="Output file name")
+    parser.add_argument("--textured", action="store_true", help="Export mesh with texture")
     args, cfg_cmd = parser.parse_known_args()
     return args, cfg_cmd
 
@@ -56,7 +57,7 @@ def main():
         init_dist(cfg.local_rank, rank=-1, world_size=-1)
     print(f"Running mesh extraction with {get_world_size()} GPUs.")
 
-    cfg.logdir = init_logging(args.config, args.logdir, makedir=True)
+    cfg.logdir = ''
 
     # Initialize data loaders and models.
     trainer = get_trainer(cfg, is_inference=True, seed=0)
@@ -80,14 +81,25 @@ def main():
     else:
         bounds = np.array([[-1.0, 1.0], [-1.0, 1.0], [-1.0, 1.0]])
 
-    mesh = extract_mesh(sdf_func=lambda x: -trainer.model_module.neural_sdf.sdf(x),
-                        bounds=bounds, intv=(2.0 / args.resolution), block_res=args.block_res)
+    sdf_func = lambda x: -trainer.model_module.neural_sdf.sdf(x)
+    texture_func = partial(extract_texture, neural_sdf=trainer.model_module.neural_sdf, 
+                           neural_rgb=trainer.model_module.neural_rgb, 
+                           appear_embed=trainer.model_module.appear_embed) if args.textured else None
+
+    mesh = extract_mesh(sdf_func=sdf_func, bounds=bounds, 
+                        intv=(2.0 / args.resolution), block_res=args.block_res,
+                        texture_func=texture_func)
 
     if is_master():
         print(f"vertices: {len(mesh.vertices)}")
         print(f"faces: {len(mesh.faces)}")
+        if args.textured:
+            print(f"colors: {len(mesh.visual.vertex_colors)}")
+            print(mesh.vertices[0], mesh.vertices[100], mesh.vertices[1000])
+            print(mesh.visual.vertex_colors[0], mesh.visual.vertex_colors[100], mesh.visual.vertex_colors[1000])
         # center and scale
         mesh.vertices = mesh.vertices * meta["sphere_radius"] + np.array(meta["sphere_center"])
+        mesh.remove_degenerate_faces()
         mesh.export(args.output_file)
 
 

projects/neuralangelo/utils/mesh.py

@@ -15,13 +15,14 @@
 import mcubes
 import torch
 import torch.distributed as dist
+import torch.nn.functional as torch_F
 from tqdm import tqdm
 
 from imaginaire.utils.distributed import get_world_size, is_master
 
 
 @torch.no_grad()
-def extract_mesh(sdf_func, bounds, intv, block_res=64):
+def extract_mesh(sdf_func, bounds, intv, block_res=64, texture_func=None):
     lattice_grid = LatticeGrid(bounds, intv=intv, block_res=block_res)
     data_loader = get_lattice_grid_loader(lattice_grid)
     mesh_blocks = []
@@ -32,21 +33,36 @@ def extract_mesh(sdf_func, bounds, intv, block_res=64):
         xyz_cuda = xyz.cuda()
         sdf_cuda = sdf_func(xyz_cuda)[..., 0]
         sdf = sdf_cuda.cpu()
-        mesh = marching_cubes(sdf.numpy(), xyz.numpy(), intv)
+        mesh = marching_cubes(sdf.numpy(), xyz.numpy(), intv, texture_func)
         mesh_blocks.append(mesh)
     mesh_blocks_gather = [None] * get_world_size()
     if dist.is_initialized():
         dist.all_gather_object(mesh_blocks_gather, mesh_blocks)
     else:
         mesh_blocks_gather = [mesh_blocks]
     if is_master():
-        mesh_blocks_all = [mesh for mesh_blocks in mesh_blocks_gather for mesh in mesh_blocks]
+        mesh_blocks_all = [mesh for mesh_blocks in mesh_blocks_gather for mesh in mesh_blocks if mesh.vertices.shape[0] > 0]
         mesh = trimesh.util.concatenate(mesh_blocks_all)
         return mesh
     else:
         return None
 
 
+@torch.no_grad()
+def extract_texture(xyz, neural_rgb, neural_sdf, appear_embed):
+    num_samples, _ = xyz.shape
+    xyz_cuda = torch.from_numpy(xyz).float().cuda()[None, None]  # [N,3] -> [1,1,N,3]
+    sdfs, feats = neural_sdf(xyz_cuda)
+    gradients, _ = neural_sdf.compute_gradients(xyz_cuda, training=False, sdf=sdfs)
+    normals = torch_F.normalize(gradients, dim=-1)
+    if appear_embed is not None:
+        feat_dim = appear_embed.embedding_dim  # [1,1,N,C] 
+        app = torch.zeros([1, 1, num_samples, feat_dim], device=sdfs.device)  # TODO: hard-coded to zero. better way?
+    else:
+        app = None
+    rgbs = neural_rgb.forward(xyz_cuda, normals, -normals, feats, app=app)  # [1,1,N,3]
+    return (rgbs.squeeze().cpu().numpy() * 255).astype(np.uint8)
+
 class LatticeGrid(torch.utils.data.Dataset):
 
     def __init__(self, bounds, intv, block_res=64):
@@ -98,21 +114,32 @@ def get_lattice_grid_loader(dataset, num_workers=8):
     )
 
 
-def marching_cubes(sdf, xyz, intv):
+def marching_cubes(sdf, xyz, intv, texture_func):
     # marching cubes
     V, F = mcubes.marching_cubes(sdf, 0.)
-    V = V * intv + xyz[0, 0, 0]
-    mesh = trimesh.Trimesh(V, F)
-    mesh = filter_points_outside_bounding_sphere(mesh)
+    if V.shape[0] > 0:
+        V = V * intv + xyz[0, 0, 0]
+        if texture_func is not None:
+            C = texture_func(V)
+            mesh = trimesh.Trimesh(V, F, vertex_colors=C)
+        else:
+            mesh = trimesh.Trimesh(V, F)
+        mesh = filter_points_outside_bounding_sphere(mesh)
+    else:
+        mesh = trimesh.Trimesh()
     return mesh
 
 
 def filter_points_outside_bounding_sphere(old_mesh):
     mask = np.linalg.norm(old_mesh.vertices, axis=-1) < 1.0
-    indices = np.ones(len(old_mesh.vertices), dtype=int) * -1
-    indices[mask] = np.arange(mask.sum())
-    faces_mask = mask[old_mesh.faces[:, 0]] & mask[old_mesh.faces[:, 1]] & mask[old_mesh.faces[:, 2]]
-    new_faces = indices[old_mesh.faces[faces_mask]]
-    new_vertices = old_mesh.vertices[mask]
-    new_mesh = trimesh.Trimesh(new_vertices, new_faces)
+    if np.any(mask):
+        indices = np.ones(len(old_mesh.vertices), dtype=int) * -1
+        indices[mask] = np.arange(mask.sum())
+        faces_mask = mask[old_mesh.faces[:, 0]] & mask[old_mesh.faces[:, 1]] & mask[old_mesh.faces[:, 2]]
+        new_faces = indices[old_mesh.faces[faces_mask]]
+        new_vertices = old_mesh.vertices[mask]
+        new_colors = old_mesh.visual.vertex_colors[mask]
+        new_mesh = trimesh.Trimesh(new_vertices, new_faces, vertex_colors=new_colors)
+    else:
+        new_mesh = trimesh.Trimesh()
     return new_mesh