| import os |
| import time |
| import argparse |
| import numpy as np |
| import torch |
| from src.utils import load_config |
| from src.mind3d import MinD3D |
| from src.utils import torch_init_model, set_random_seed |
|
|
|
|
| if __name__ == '__main__': |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--config', type=str, default="./configs/mind3d.yaml") |
| parser.add_argument('--check_point_path', type=str, default="ckpt/mind3d_sub01.pt") |
| parser.add_argument('--fmri_dir', type=str, default="/home/hubin/data/fMRI-Shape/sub-0001/npy_data") |
| parser.add_argument('--test_list', type=str, default="/home/hubin/data/fMRI-Shape/annotations/core_test_list.txt") |
| parser.add_argument('--out_dir', type=str, default="/home/hubin/trellis_work/outputs/mind3d_sub01") |
| parser.add_argument('--topk', type=int, default=250) |
| parser.add_argument('--seed', type=int, default=100) |
| parser.add_argument('--shard', type=int, default=0) |
| parser.add_argument('--num_shards', type=int, default=1) |
| args = parser.parse_args() |
|
|
| cfg = load_config(args.config, 'configs/default.yaml') |
| device = torch.device("cuda") |
| os.makedirs(os.path.join(args.out_dir, "meshes"), exist_ok=True) |
|
|
| model_full = MinD3D(cfg, device=device) |
| state_dict = torch.load(args.check_point_path, map_location='cpu', weights_only=False) |
| torch_init_model(model_full, state_dict) |
| del state_dict |
| model_full.eval() |
|
|
| ids = [l.strip() for l in open(args.test_list) if l.strip()][args.shard::args.num_shards] |
| for i, obj in enumerate(ids): |
| out_path = os.path.join(args.out_dir, "meshes", obj.replace("/", "_") + ".obj") |
| if os.path.exists(out_path): |
| continue |
| set_random_seed(args.seed) |
| t0 = time.time() |
| fmri_data = np.load(os.path.join(args.fmri_dir, obj + ".npy")) |
| fmri_data = torch.from_numpy(fmri_data[2:8]).unsqueeze(0).cuda() |
| with torch.no_grad(): |
| fmri_embedding = model_full.get_fmri_embedding_batch(fmri_data, diffusion_bs=128, decoder_bs=1) |
| mesh_list = model_full.generate_mesh_from_embedding(fmri_embedding, bs=1, top_k=args.topk) |
| if mesh_list and mesh_list[0].vertices.shape[0]: |
| mesh_list[0].export(out_path) |
| print(f"[{i + 1}/{len(ids)}] {obj} {time.time() - t0:.1f}s", flush=True) |
|
|
