def load_voxel(category, model_id):
"""Loads the voxel tensors given the model category and model ID.
Args:
category: Model category (e.g. '03001627')
model_id: Model ID (e.g. '587ee5822bb56bd07b11ae648ea92233')
Returns:
voxel_tensor: Voxel tensor of shape (height x width x depth x channels).
"""
voxel_fn = get_voxel_file(category, model_id)
voxel_tensor = read_nrrd(voxel_fn)
return voxel_tensor
def test_nrrd_rw():
"""Function for testing whether the read/write functions work correctly.
"""
from lib.config import cfg
from lib.data_io import get_voxel_file
# Read an NRRD file, then write it and manually make sure it looks the same
model_id = 'fcfb7012968416679c0b027ae5b223d6'
nrrd_filepath = get_voxel_file(None, model_id)
voxel_tensor = read_nrrd(nrrd_filepath)
write_nrrd(voxel_tensor, '/tmp/test_nrrd_rw.nrrd')
# Make a voxel tensor with red, green, blue, and black voxels (1 each) and manually verify it
sparse_voxel_tensor = np.zeros((32, 32, 32, 4))
sparse_voxel_tensor[0, 0, 0, :] = np.array([1, 0, 0, 1]) # Red
sparse_voxel_tensor[31, 0, 0, :] = np.array([0, 1, 0, 1]) # Green
sparse_voxel_tensor[0, 31, 0, :] = np.array([0, 0, 1, 1]) # Blue
sparse_voxel_tensor[0, 0, 31, :] = np.array([0, 0, 0, 1]) # Black
write_nrrd(sparse_voxel_tensor, '/tmp/test_nrrd_rw_2.nrrd')
def render_model_id(model_ids, out_dir, check=True):
"""Render models based on their model IDs.
Args:
model_ids: List of model ID strings.
nrrd_dir: Directory to write the NRRD files to.
out_dir: Directory to write the PNG files to.
check: Check if the output directory already exists and provide a warning if so.
"""
if cfg.CONST.DATASET == 'primitives':
categories = [model_id.split('_')[0] for model_id in model_ids]
elif cfg.CONST.DATASET == 'shapenet':
categories = [None] * len(model_ids)
else:
raise ValueError('Please use a valid dataset.')
if check is True:
if os.path.isdir(out_dir):
print('Output directory:', out_dir)
input('Output render directory exists! Continue?')
else:
os.makedirs(out_dir)
else:
os.makedirs(out_dir, exist_ok=True)
nrrd_files = []
for category, model_id in zip(categories, model_ids):
nrrd_files.append(get_voxel_file(category, model_id))
txt_filepath = os.path.join(out_dir, 'nrrd_filenames.txt')
with open(txt_filepath, 'w') as f:
for outfile in nrrd_files:
f.write(outfile + '\n')
print('Filenames written to {}'.format(txt_filepath))
render_nrrd(txt_filepath, out_dir, check=False)
def load_label(self, category, model_id):
voxel_fn = get_voxel_file(category, model_id)
with open(voxel_fn, 'rb') as f:
voxel = read_as_3d_array(f)
return voxel