File size: 7,050 Bytes
535964c
 
 
 
 
55d2f59
535964c
 
 
 
55d2f59
 
535964c
 
 
 
 
 
 
 
 
 
 
 
55d2f59
 
535964c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
55d2f59
 
535964c
 
 
55d2f59
 
535964c
 
55d2f59
 
535964c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
55d2f59
535964c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
55d2f59
 
 
535964c
 
 
 
 
 
 
55d2f59
535964c
55d2f59
535964c
 
 
 
 
 
 
 
 
 
 
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
import os, argparse, importlib
import torch
import time
import trimesh
import numpy as np
from MeshAnything.models.meshanything_v2 import MeshAnythingV2
import datetime
from accelerate import Accelerator
from accelerate.utils import set_seed
from accelerate.utils import DistributedDataParallelKwargs
from safetensors.torch import load_model

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] >= 8192, "input pc_normal should have at least 4096 points"
                idx = np.random.choice(cur_data.shape[0], 8192, 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('--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(
        '--input_type',
        choices=['mesh','pc_normal'],
        default='pc',
        help="Type of the asset to process (default: pc)"
    )

    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_v2():
    model = MeshAnythingV2()
    print("load model over!!!")

    ckpt_path = hf_hub_download(
        repo_id="Yiwen-ntu/MeshAnythingV2",
        filename="350m.pth",
    )

    load_model(model, ckpt_path)

    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_v2()
    # 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]
                valid_mask = torch.all(~torch.isnan(recon_mesh.reshape((-1, 9))), dim=1)
                recon_mesh = recon_mesh[valid_mask]  # 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.nondegenerate_faces())
                scene_mesh.update_faces(scene_mesh.unique_faces())
                scene_mesh.remove_unreferenced_vertices()
                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}")