Upload 2 files

main.py

@@ -0,0 +1,177 @@
+import os, argparse, importlib
+import torch
+import time
+import trimesh
+import numpy as np
+from MeshAnything.models.meshanything import MeshAnything
+import datetime
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+from accelerate.utils import DistributedDataParallelKwargs
+from safetensors import safe_open
+from mesh_to_pc import process_mesh_to_pc
+from huggingface_hub import hf_hub_download
+
+class Dataset:
+    def __init__(self, input_type, input_list, mc=False):
+        super().__init__()
+        self.data = []
+        if input_type == 'pc_normal':
+            for input_path in input_list:
+                # load npy
+                cur_data = np.load(input_path)
+                # sample 4096
+                assert cur_data.shape[0] >= 4096, "input pc_normal should have at least 4096 points"
+                idx = np.random.choice(cur_data.shape[0], 4096, replace=False)
+                cur_data = cur_data[idx]
+                self.data.append({'pc_normal': cur_data, 'uid': input_path.split('/')[-1].split('.')[0]})
+
+        elif input_type == 'mesh':
+            mesh_list = []
+            for input_path in input_list:
+                # load ply
+                cur_data = trimesh.load(input_path)
+                mesh_list.append(cur_data)
+            if mc:
+                print("First Marching Cubes and then sample point cloud, need several minutes...")
+            pc_list, _ = process_mesh_to_pc(mesh_list, marching_cubes=mc)
+            for input_path, cur_data in zip(input_list, pc_list):
+                self.data.append({'pc_normal': cur_data, 'uid': input_path.split('/')[-1].split('.')[0]})
+        print(f"dataset total data samples: {len(self.data)}")
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        data_dict = {}
+        data_dict['pc_normal'] = self.data[idx]['pc_normal']
+        # normalize pc coor
+        pc_coor = data_dict['pc_normal'][:, :3]
+        normals = data_dict['pc_normal'][:, 3:]
+        bounds = np.array([pc_coor.min(axis=0), pc_coor.max(axis=0)])
+        pc_coor = pc_coor - (bounds[0] + bounds[1])[None, :] / 2
+        pc_coor = pc_coor / np.abs(pc_coor).max() * 0.9995
+        assert (np.linalg.norm(normals, axis=-1) > 0.99).all(), "normals should be unit vectors, something wrong"
+        data_dict['pc_normal'] = np.concatenate([pc_coor, normals], axis=-1, dtype=np.float16)
+        data_dict['uid'] = self.data[idx]['uid']
+
+        return data_dict
+
+def get_args():
+    parser = argparse.ArgumentParser("MeshAnything", add_help=False)
+
+    parser.add_argument('--llm', default="facebook/opt-350m", type=str)
+    parser.add_argument('--input_dir', default=None, type=str)
+    parser.add_argument('--input_path', default=None, type=str)
+
+    parser.add_argument('--out_dir', default="inference_out", type=str)
+    parser.add_argument('--pretrained_weights', default="MeshAnything_350m.pth", type=str)
+
+    parser.add_argument(
+        '--input_type',
+        choices=['mesh','pc_normal'],
+        default='pc',
+        help="Type of the asset to process (default: pc)"
+    )
+
+    parser.add_argument("--codebook_size", default=8192, type=int)
+    parser.add_argument("--codebook_dim", default=1024, type=int)
+
+    parser.add_argument("--n_max_triangles", default=800, type=int)
+
+    parser.add_argument("--batchsize_per_gpu", default=1, type=int)
+    parser.add_argument("--seed", default=0, type=int)
+
+    parser.add_argument("--mc", default=False, action="store_true")
+    parser.add_argument("--sampling", default=False, action="store_true")
+
+    args = parser.parse_args()
+    return args
+
+def load_model(args):
+    model = MeshAnything(args)
+    print("load model over!!!")
+
+    ckpt_path = hf_hub_download(
+        repo_id="Yiwen-ntu/MeshAnything",
+        filename="MeshAnything_350m.pth",
+    )
+    tensors = {}
+    with safe_open(ckpt_path, framework="pt", device=0) as f:
+        for k in f.keys():
+            tensors[k] = f.get_tensor(k)
+
+    model.load_state_dict(tensors, strict=True)
+    print("load weights over!!!")
+    return model
+if __name__ == "__main__":
+    args = get_args()
+
+    cur_time = datetime.datetime.now().strftime("%d_%H-%M-%S")
+    checkpoint_dir = os.path.join(args.out_dir, cur_time)
+    os.makedirs(checkpoint_dir, exist_ok=True)
+    kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
+    accelerator = Accelerator(
+        mixed_precision="fp16",
+        project_dir=checkpoint_dir,
+        kwargs_handlers=[kwargs]
+    )
+
+    model = load_model(args)
+    # create dataset
+    if args.input_dir is not None:
+        input_list = sorted(os.listdir(args.input_dir))
+        # only ply, obj or npy
+        if args.input_type == 'pc_normal':
+            input_list = [os.path.join(args.input_dir, x) for x in input_list if x.endswith('.npy')]
+        else:
+            input_list = [os.path.join(args.input_dir, x) for x in input_list if x.endswith('.ply') or x.endswith('.obj') or x.endswith('.npy')]
+        set_seed(args.seed)
+        dataset = Dataset(args.input_type, input_list, args.mc)
+    elif args.input_path is not None:
+        set_seed(args.seed)
+        dataset = Dataset(args.input_type, [args.input_path], args.mc)
+    else:
+        raise ValueError("input_dir or input_path must be provided.")
+
+    dataloader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=args.batchsize_per_gpu,
+        drop_last = False,
+        shuffle = False,
+    )
+
+    if accelerator.state.num_processes > 1:
+        model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
+    dataloader, model = accelerator.prepare(dataloader, model)
+    begin_time = time.time()
+    print("Generation Start!!!")
+    with accelerator.autocast():
+        for curr_iter, batch_data_label in enumerate(dataloader):
+            curr_time = time.time()
+            outputs = model(batch_data_label['pc_normal'], sampling=args.sampling)
+            batch_size = outputs.shape[0]
+            device = outputs.device
+
+            for batch_id in range(batch_size):
+                recon_mesh = outputs[batch_id]
+                recon_mesh = recon_mesh[~torch.isnan(recon_mesh[:, 0, 0])]  # nvalid_face x 3 x 3
+                vertices = recon_mesh.reshape(-1, 3).cpu()
+                vertices_index = np.arange(len(vertices))  # 0, 1, ..., 3 x face
+                triangles = vertices_index.reshape(-1, 3)
+
+                scene_mesh = trimesh.Trimesh(vertices=vertices, faces=triangles, force="mesh",
+                                             merge_primitives=True)
+                scene_mesh.merge_vertices()
+                scene_mesh.update_faces(scene_mesh.unique_faces())
+                scene_mesh.fix_normals()
+                save_path = os.path.join(checkpoint_dir, f'{batch_data_label["uid"][batch_id]}_gen.obj')
+                num_faces = len(scene_mesh.faces)
+                brown_color = np.array([255, 165, 0, 255], dtype=np.uint8)
+                face_colors = np.tile(brown_color, (num_faces, 1))
+
+                scene_mesh.visual.face_colors = face_colors
+                scene_mesh.export(save_path)
+                print(f"{save_path} Over!!")
+    end_time = time.time()
+    print(f"Total time: {end_time - begin_time}")

mesh_to_pc.py

@@ -0,0 +1,58 @@
+import mesh2sdf.core
+import numpy as np
+import skimage.measure
+import trimesh
+
+def normalize_vertices(vertices, scale=0.9):
+    bbmin, bbmax = vertices.min(0), vertices.max(0)
+    center = (bbmin + bbmax) * 0.5
+    scale = 2.0 * scale / (bbmax - bbmin).max()
+    vertices = (vertices - center) * scale
+    return vertices, center, scale
+
+def export_to_watertight(normalized_mesh, octree_depth: int = 7):
+    """
+        Convert the non-watertight mesh to watertight.
+
+        Args:
+            input_path (str): normalized path
+            octree_depth (int):
+
+        Returns:
+            mesh(trimesh.Trimesh): watertight mesh
+
+        """
+    size = 2 ** octree_depth
+    level = 2 / size
+
+    scaled_vertices, to_orig_center, to_orig_scale = normalize_vertices(normalized_mesh.vertices)
+
+    sdf = mesh2sdf.core.compute(scaled_vertices, normalized_mesh.faces, size=size)
+
+    vertices, faces, normals, _ = skimage.measure.marching_cubes(np.abs(sdf), level)
+
+    # watertight mesh
+    vertices = vertices / size * 2 - 1 # -1 to 1
+    vertices = vertices / to_orig_scale + to_orig_center
+    # vertices = vertices / to_orig_scale + to_orig_center
+    mesh = trimesh.Trimesh(vertices, faces, normals=normals)
+
+    return mesh
+
+def process_mesh_to_pc(mesh_list, marching_cubes = False, sample_num = 4096):
+    # mesh_list : list of trimesh
+    pc_normal_list = []
+    return_mesh_list = []
+    for mesh in mesh_list:
+        if marching_cubes:
+            mesh = export_to_watertight(mesh)
+            print("MC over!")
+        return_mesh_list.append(mesh)
+        points, face_idx = mesh.sample(sample_num, return_index=True)
+        normals = mesh.face_normals[face_idx]
+
+        pc_normal = np.concatenate([points, normals], axis=-1, dtype=np.float16)
+        pc_normal_list.append(pc_normal)
+        print("process mesh success")
+    return pc_normal_list, return_mesh_list
+