1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
|
import os
import os.path as osp
import torch
from torch_geometric.data import download_url, extract_gz
from torch_geometric.loader import DataLoader
from torch_geometric.testing import onlyOnline, withPackage
from torch_geometric.transforms import ToSLIC
resources = [
'https://ossci-datasets.s3.amazonaws.com/mnist/t10k-images-idx3-ubyte.gz',
'https://ossci-datasets.s3.amazonaws.com/mnist/t10k-labels-idx1-ubyte.gz',
]
@onlyOnline
@withPackage('torchvision', 'skimage')
def test_to_superpixels(tmp_path):
import torchvision.transforms as T
from torchvision.datasets.mnist import (
MNIST,
read_image_file,
read_label_file,
)
raw_folder = osp.join(tmp_path, 'MNIST', 'raw')
processed_folder = osp.join(tmp_path, 'MNIST', 'processed')
os.makedirs(raw_folder, exist_ok=True)
os.makedirs(processed_folder, exist_ok=True)
for resource in resources:
path = download_url(resource, raw_folder)
extract_gz(path, osp.join(tmp_path, raw_folder))
test_set = (
read_image_file(osp.join(raw_folder, 't10k-images-idx3-ubyte')),
read_label_file(osp.join(raw_folder, 't10k-labels-idx1-ubyte')),
)
torch.save(test_set, osp.join(processed_folder, 'training.pt'))
torch.save(test_set, osp.join(processed_folder, 'test.pt'))
dataset = MNIST(tmp_path, download=False)
dataset.transform = T.Compose([T.ToTensor(), ToSLIC()])
data, y = dataset[0]
assert len(data) == 2
assert data.pos.dim() == 2 and data.pos.size(1) == 2
assert data.x.dim() == 2 and data.x.size(1) == 1
assert data.pos.size(0) == data.x.size(0)
assert y == 7
loader = DataLoader(dataset, batch_size=2, shuffle=False)
for batch, y in loader:
assert batch.num_graphs == len(batch) == 2
assert batch.pos.dim() == 2 and batch.pos.size(1) == 2
assert batch.x.dim() == 2 and batch.x.size(1) == 1
assert batch.batch.dim() == 1
assert batch.ptr.dim() == 1
assert batch.pos.size(0) == batch.x.size(0) == batch.batch.size(0)
assert y.tolist() == [7, 2]
break
dataset.transform = T.Compose(
[T.ToTensor(), ToSLIC(add_seg=True, add_img=True)])
data, y = dataset[0]
assert len(data) == 4
assert data.pos.dim() == 2 and data.pos.size(1) == 2
assert data.x.dim() == 2 and data.x.size(1) == 1
assert data.pos.size(0) == data.x.size(0)
assert data.seg.size() == (1, 28, 28)
assert data.img.size() == (1, 1, 28, 28)
assert data.seg.max().item() + 1 == data.x.size(0)
assert y == 7
loader = DataLoader(dataset, batch_size=2, shuffle=False)
for batch, y in loader:
assert batch.num_graphs == len(batch) == 2
assert batch.pos.dim() == 2 and batch.pos.size(1) == 2
assert batch.x.dim() == 2 and batch.x.size(1) == 1
assert batch.batch.dim() == 1
assert batch.ptr.dim() == 1
assert batch.pos.size(0) == batch.x.size(0) == batch.batch.size(0)
assert batch.seg.size() == (2, 28, 28)
assert batch.img.size() == (2, 1, 28, 28)
assert y.tolist() == [7, 2]
break
|
