main.py

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}")