def backproject(voxel_dim, voxel_size, origin, projection, features):
""" Takes 2d features and fills them along rays in a 3d volume
This function implements eqs. 1,2 in https://arxiv.org/pdf/2003.10432.pdf
Each pixel in a feature image corresponds to a ray in 3d.
We fill all the voxels along the ray with that pixel's features.
Args:
voxel_dim: size of voxel volume to construct (nx,ny,nz)
voxel_size: metric size of each voxel (ex: .04m)
origin: origin of the voxel volume (xyz position of voxel (0,0,0))
projection: bx4x3 projection matrices (intrinsics@extrinsics)
features: bxcxhxw 2d feature tensor to be backprojected into 3d
Returns:
volume: b x c x nx x ny x nz 3d feature volume
valid: b x 1 x nx x ny x nz volume.
Each voxel contains a 1 if it projects to a pixel
and 0 otherwise (not in view frustrum of the camera)
"""
batch = features.size(0)
channels = features.size(1)
device = features.device
nx, ny, nz = voxel_dim
coords = coordinates(voxel_dim, device).unsqueeze(0).expand(batch, -1,
-1) # bx3xhwd
world = coords.type_as(projection) * voxel_size + origin.to(
device).unsqueeze(2)
world = torch.cat((world, torch.ones_like(world[:, :1])), dim=1)
torch.set_printoptions(sci_mode=False)
camera = torch.bmm(projection, world)
#print('camera shape:', camera.shape)
px = (camera[:, 0, :] / camera[:, 2, :]).round().type(torch.long)
py = (camera[:, 1, :] / camera[:, 2, :]).round().type(torch.long)
pz = camera[:, 2, :]
# voxels in view frustrum
height, width = features.size()[2:]
valid = (px >= 0) & (py >= 0) & (px < width) & (py < height) & (pz > 0
) # bxhwd
# put features in volume
volume = torch.zeros(batch,
channels,
nx * ny * nz,
dtype=features.dtype,
device=device)
for b in range(batch):
volume[b, :, valid[b]] = features[b, :, py[b, valid[b]], px[b,
valid[b]]]
volume = volume.view(batch, channels, nx, ny, nz)
valid = valid.view(batch, 1, nx, ny, nz)
return volume, valid
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'))