File size: 2,300 Bytes
87e9895 | 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 | 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)
|