I have a dataframe representing a point cloud plus node data: I have hence X,Y,Z positions as well as "node features". I would like to represent it as a graph to use it in a GNN.
I would hence like to use the GNN transform in PyGeometric to extract edge connectivity info. An example in the docs can be found here.
I am not managing to get it to work, I think I need to represent my dataset as a torch_geometric.datasets.GeometricShapes object first?
This is what I tried
import pandas as pd
import numpy as np
import torch
from torch_geometric.data import Data, Dataset,InMemoryDataset
from torch_geometric.transforms import SamplePoints, KNNGraph
import torch_geometric.transforms as T
from torch_geometric.datasets import GeometricShapes
class CustomDataset(InMemoryDataset):
def __init__(self, listOfDataObjects):
super().__init__()
self.data, self.slices = self.collate(listOfDataObjects)
def __len__(self):
return len(self.slices)
def __getitem__(self, idx):
sample = self.get(idx)
return sample
## My data
data_fake = pd.DataFrame(data=np.random.rand(20,5), columns =['X','Y', 'Z','Node_feature_1', 'Node_feature_2'])
# Get a data object and then Dataset
my_fake_data = Data()
my_fake_data.pos = torch.from_numpy(data_fake[['X','Y', 'Z']].values)
dataset_cloud = CustomDataset([my_fake_data])
# Transform
dataset_cloud.transform = T.Compose([SamplePoints(num=20), KNNGraph(k=5)])
alas, dataset_cloud[0] returns
>>>Data(pos=[20, 3])
so no edge connectivity info for the nearest neighbours is added, as in the docs.
Any hints please? Thanks
EDIT
I think a workaround is not to use dataset.transforms and apply KNNGraph to the point cloud directly, as in
knn =KNNGraph(k=3)
data_knn = knn(dataset_cloud[0])
which return ad edge_index attribute as expected.
Still would like to understand what goes wrong with the approach detailed above