1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
|
import os.path as osp
import torch_geometric.transforms as T
from torch_geometric.datasets import ModelNet
def get_dataset(num_points):
name = 'ModelNet10'
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', name)
pre_transform = T.NormalizeScale()
transform = T.SamplePoints(num_points)
train_dataset = ModelNet(path, name='10', train=True, transform=transform,
pre_transform=pre_transform)
test_dataset = ModelNet(path, name='10', train=False, transform=transform,
pre_transform=pre_transform)
return train_dataset, test_dataset
|
