def eval_tsdf(file_pred, file_trgt):
""" Compute TSDF metrics between prediction and target.
Opens the TSDFs, aligns the voxels and runs the metrics
Args:
file_pred: file path of prediction
file_trgt: file path of target
Returns:
Dict of TSDF metrics
"""
tsdf_pred = TSDF.load(file_pred)
tsdf_trgt = TSDF.load(file_trgt)
# align prediction voxels to target voxels
# we use align_corners=True here so that when shifting by integer number
# of voxels we do not interpolate.
# TODO: verify align corners when we need to do interpolation (non integer
# voxel shifts)
shift = (tsdf_trgt.origin - tsdf_pred.origin) / tsdf_trgt.voxel_size
assert torch.allclose(shift, shift.round())
tsdf_pred = tsdf_pred.transform(voxel_dim=list(tsdf_trgt.tsdf_vol.shape),
origin=tsdf_trgt.origin,
align_corners=True)
metrics = {'l1': l1(tsdf_pred, tsdf_trgt)}
return metrics
def postprocess(self, batch):
""" Wraps the network output into a TSDF data structure
Args:
batch: dict containg network outputs
Returns:
list of TSDFs (one TSDF per scene in the batch)
"""
key = 'vol_%05d' % self.voxel_sizes[
0] # only get vol of final resolution
out = []
batch_size = len(batch[key + '_tsdf'])
for i in range(batch_size):
tsdf = TSDF(self.voxel_size, self.origin,
batch[key + '_tsdf'][i].squeeze(0))
# add semseg vol
if ('semseg' in self.voxel_types) and (key + '_semseg' in batch):
semseg = batch[key + '_semseg'][i]
if semseg.ndim == 4:
semseg = semseg.argmax(0)
tsdf.attribute_vols['semseg'] = semseg
# add color vol
if 'color' in self.voxel_types:
color = batch[key + '_color'][i]
tsdf.attribute_vols['color'] = color
out.append(tsdf)
return out
def map_tsdf(info, data, voxel_types, voxel_sizes):
""" Load TSDFs from paths in info.
Args:
info: dict with paths to TSDF files (see datasets/README)
data: dict to add TSDF data to
voxel_types: list of voxel attributes to load with the TSDF
voxel_sizes: list of voxel sizes to load
Returns:
dict with TSDFs included
"""
if len(voxel_types)>0:
for scale in voxel_sizes:
data['vol_%02d'%scale] = TSDF.load(info['file_name_vol_%02d'%scale],
voxel_types)
return data
def process(info_file, save_path, total_scenes_index, total_scenes_count):
# gt depth data loader
width, height = 640, 480
transform = transforms.Compose([
transforms.ResizeImage((width,height)),
transforms.ToTensor(),
])
dataset = SceneDataset(info_file, transform, frame_types=['depth'])
dataloader = torch.utils.data.DataLoader(dataset, batch_size=None,
batch_sampler=None, num_workers=2)
scene = dataset.info['scene']
# get info about tsdf
file_tsdf_pred = os.path.join(save_path, '%s.npz'%scene)
temp = TSDF.load(file_tsdf_pred)
voxel_size = int(temp.voxel_size*100)
# re-fuse to remove hole filling since filled holes are penalized in
# mesh metrics
vol_dim = list(temp.tsdf_vol.shape)
origin = temp.origin
tsdf_fusion = TSDFFusion(vol_dim, float(voxel_size)/100, origin, color=False)
device = tsdf_fusion.device
# mesh renderer
renderer = Renderer()
mesh_file = os.path.join(save_path, '%s.ply'%scene)
mesh = trimesh.load(mesh_file, process=False)
mesh_opengl = renderer.mesh_opengl(mesh)
for i, d in enumerate(dataloader):
if i%25==0:
print(total_scenes_index, total_scenes_count,scene, i, len(dataloader))
depth_trgt = d['depth'].numpy()
_, depth_pred = renderer(height, width, d['intrinsics'], d['pose'], mesh_opengl)
temp = eval_depth(depth_pred, depth_trgt)
if i==0:
metrics_depth = temp
else:
metrics_depth = {key:value+temp[key]
for key, value in metrics_depth.items()}
# # play video visualizations of depth
# viz1 = (np.clip((depth_trgt-.5)/5,0,1)*255).astype(np.uint8)
# viz2 = (np.clip((depth_pred-.5)/5,0,1)*255).astype(np.uint8)
# viz1 = cv2.applyColorMap(viz1, cv2.COLORMAP_JET)
# viz2 = cv2.applyColorMap(viz2, cv2.COLORMAP_JET)
# viz1[depth_trgt==0]=0
# viz2[depth_pred==0]=0
# viz = np.hstack((viz1,viz2))
# cv2.imshow('test', viz)
# cv2.waitKey(1)
tsdf_fusion.integrate((d['intrinsics'] @ d['pose'].inverse()[:3,:]).to(device),
torch.as_tensor(depth_pred).to(device))
metrics_depth = {key:value/len(dataloader)
for key, value in metrics_depth.items()}
# save trimed mesh
file_mesh_trim = os.path.join(save_path, '%s_trim.ply'%scene)
tsdf_fusion.get_tsdf().get_mesh().export(file_mesh_trim)
# eval tsdf
file_tsdf_trgt = dataset.info['file_name_vol_%02d'%voxel_size]
metrics_tsdf = eval_tsdf(file_tsdf_pred, file_tsdf_trgt)
# eval trimed mesh
file_mesh_trgt = dataset.info['file_name_mesh_gt']
metrics_mesh = eval_mesh(file_mesh_trim, file_mesh_trgt)
# transfer labels from pred mesh to gt mesh using nearest neighbors
file_attributes = os.path.join(save_path, '%s_attributes.npz'%scene)
if os.path.exists(file_attributes):
mesh.vertex_attributes = np.load(file_attributes)
print(mesh.vertex_attributes)
mesh_trgt = trimesh.load(file_mesh_trgt, process=False)
mesh_transfer = project_to_mesh(mesh, mesh_trgt, 'semseg')
semseg = mesh_transfer.vertex_attributes['semseg']
# save as txt for benchmark evaluation
np.savetxt(os.path.join(save_path, '%s.txt'%scene), semseg, fmt='%d')
mesh_transfer.export(os.path.join(save_path, '%s_transfer.ply'%scene))
# TODO: semseg val evaluation
metrics = {**metrics_depth, **metrics_mesh, **metrics_tsdf}
print(metrics)
rslt_file = os.path.join(save_path, '%s_metrics.json'%scene)
json.dump(metrics, open(rslt_file, 'w'))
return scene, metrics
def label_scene(path_meta, scene, voxel_size, dist_thresh=.05, verbose=2):
""" Transfer instance labels from ground truth mesh to TSDF
For each voxel find the nearest vertex and transfer the label if
it is close enough to the voxel.
Args:
path_meta: path to save the TSDFs
(we recommend creating a parallel directory structure to save
derived data so that we don't modify the original dataset)
scene: name of scene to process
voxel_size: voxel size of TSDF to process
dist_thresh: beyond this distance labels are not transferd
verbose: how much logging to print
Returns:
Updates the TSDF (.npz) file with the instance volume
"""
# dist_thresh: beyond this distance to nearest gt mesh vertex,
# voxels are not labeled
if verbose > 0:
print('labeling', scene)
info_file = os.path.join(path_meta, scene, 'info.json')
data = load_info_json(info_file)
# each vertex in gt mesh indexs a seg group
segIndices = json.load(open(data['file_name_seg_indices'],
'r'))['segIndices']
# maps seg groups to instances
segGroups = json.load(open(data['file_name_seg_groups'], 'r'))['segGroups']
mapping = {
ind: group['id'] + 1
for group in segGroups for ind in group['segments']
}
# get per vertex instance ids (0 is unknown, [1,...] are objects)
n = len(segIndices)
instance_verts = torch.zeros(n, dtype=torch.long)
for i in range(n):
if segIndices[i] in mapping:
instance_verts[i] = mapping[segIndices[i]]
# load vertex locations
mesh = trimesh.load(data['file_name_mesh_gt'], process=False)
verts = mesh.vertices
# construct kdtree of vertices for fast nn lookup
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(verts)
kdtree = o3d.geometry.KDTreeFlann(pcd)
# load tsdf volume
tsdf = TSDF.load(data['file_name_vol_%02d' % voxel_size])
coords = coordinates(tsdf.tsdf_vol.size(), device=torch.device('cpu'))
coords = coords.type(torch.float) * tsdf.voxel_size + tsdf.origin.T
mask = tsdf.tsdf_vol.abs().view(-1) < 1
# transfer vertex instance ids to voxels near surface
instance_vol = torch.zeros(len(mask), dtype=torch.long)
for i in mask.nonzero():
_, inds, dist = kdtree.search_knn_vector_3d(coords[:, i], 1)
if dist[0] < dist_thresh:
instance_vol[i] = instance_verts[inds[0]]
tsdf.attribute_vols['instance'] = instance_vol.view(
list(tsdf.tsdf_vol.size()))
tsdf.save(data['file_name_vol_%02d' % voxel_size])
key = 'vol_%02d' % voxel_size
temp_data = {
key: tsdf,
'instances': data['instances'],
'dataset': data['dataset']
}
tsdf = transforms.InstanceToSemseg('nyu40')(temp_data)[key]
mesh = tsdf.get_mesh('semseg')
fname = data['file_name_vol_%02d' % voxel_size]
mesh.export(fname.replace('tsdf', 'mesh').replace('.npz', '_semseg.ply'))