### Copyright (C) 2017 NVIDIA Corporation. All rights reserved.
### Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
from options.test_options import TestOptions
from data.data_loader import CreateDataLoader
from models.models import create_PRSNet
import torch
import scipy.io as sio
import os
import numpy as np
opt = TestOptions().parse(save=False)
opt.nThreads = 0 # test code only supports nThreads = 1
opt.batchSize = 1 # test code only supports batchSize = 1
opt.noshuffle = True # no shuffle
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
save_dir = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch))
if not os.path.exists(save_dir):
os.makedirs(save_dir)
# test
PRSNet = create_PRSNet(opt)
if opt.data_type == 16:
PRSNet.half()
elif opt.data_type == 8:
PRSNet.type(torch.uint8)
for i, data in enumerate(dataset):
plane, quat = PRSNet.inference(data['voxel'])
data_path = data['path'][0]
print('[%s] process mat ... %s' % (str(i),data_path))
matdata = sio.loadmat(data_path,verify_compressed_data_integrity=False)
import ntpath
short_path = ntpath.basename(data_path)
name = os.path.splitext(short_path)[0]
model = {'name':name, 'voxel':matdata['Volume'], 'vertices':matdata['vertices'], 'faces':matdata['faces'], 'sample':np.transpose(matdata['surfaceSamples']), 'axisangle': matdata['axisangle']}
for j in range(opt.num_plane):
model['plane'+str(j)] = plane[j].cpu().numpy()
for j in range(opt.num_quat):
model['quat'+str(j)] = quat[j].cpu().numpy()
sio.savemat(save_dir+"/"+name+".mat",model)