-
Notifications
You must be signed in to change notification settings - Fork 91
/
main.py
177 lines (153 loc) · 7.05 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
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}")