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'))