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