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