def test_features_at_coordinates(self):
coords = torch.IntTensor([[0, 0], [0, 1], [0, 2], [2, 0], [2, 2]])
feats = torch.FloatTensor([[1.1, 2.1, 3.1, 4.1, 5.1]]).t()
X = SparseTensor(features=feats, coordinates=coords)
feats = X.features_at_coordinates(
torch.FloatTensor([[0, 0], [0, 1], [0, 2], [2, 2], [0, 0],
[0, 0.5]])).flatten()
self.assertTrue(feats[0] == 1.1)
self.assertTrue(feats[3] == 5.1)
self.assertTrue(feats[4] == 1.1)
def get_neighbor_feature(neighbor: torch.Tensor, x: ME.SparseTensor):
"""
fetch neighbor voxel's feature tensor.
Input:
neighbor: torch.Tensor [B*npoint, k, xyz]
x: ME.SparseTensor
"""
B_npoint, k, _ = tuple(neighbor.size())
neighbor = neighbor.view(-1, 4).float() # [B*npoint*k, bxyz]
features = x.features_at_coordinates(neighbor)
_, dim = features.shape
features = features.view(-1, k, dim)
return features
