1

I have a point cloud that I want to use a graph neural network on. Each point in the point cloud is characterised by its positional coordinates as well as it's color. So a single node is (X, Y, Z, C).

Now I want to apply an Edge Convolution on this (as described in the DGL Edge-Conv example, and to do it I should build a Nearest Neighbors graph on (X, Y, Z) (And not on C), then use all the 4 properties as features for my neural network.

What would be a clean and efficient way to do this? (I have a lot of data so I want to batch and collate well)

desertnaut
  • 57,590
  • 26
  • 140
  • 166
Animesh Sinha
  • 817
  • 2
  • 11
  • 24

1 Answers1

1

Supposing you have a tensor pc of shape (NUM_POINTS, 4) where each row is (X, Y, Z, C), then you could use sklearn as follows:

from sklearn.neighbors import NearestNeighbors
import dgl

k = 3 # number of neighbours you want

neigh = NearestNeighbors(n_neighbors=k)
neigh.fit(pc[:, :3].numpy()) # selects only (X, Y, Z)

knn = neigh.kneighbors_graph()
graph = dgl.from_scipy(knn)
graph.ndata['x'] = pc

I would recommend saving these graphs to disk, so they are not computed each time you train etc.

Pete
  • 156
  • 6