-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathdataset.py
101 lines (75 loc) · 3.39 KB
/
dataset.py
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
91
92
93
94
95
96
97
98
99
100
101
import os
import numpy as np
import torch
# from torch.utils.data import Dataset, DataLoader
from torchvision import datasets
from PIL import Image
class CIFAR10(datasets.CIFAR10):
def __init__(self, root, train=True, transform=None, target_transform=None, download=False):
super(CIFAR10, self).__init__(root, train=train, transform=transform,
target_transform=target_transform, download=download)
# unify the interface
if not hasattr(self, 'data'): # torch <= 0.4.1
if self.train:
self.data, self.targets = self.train_data, self.train_labels
else:
self.data, self.targets = self.test_data, self.test_labels
def __getitem__(self, index):
img, target = self.data[index], self.targets[index]
# doing this so that it is consistent with all other datasets
# to return a PIL Image
img = Image.fromarray(img)
if self.transform is not None:
img = self.transform(img)
if self.target_transform is not None:
target = self.target_transform(target)
return img, target, index
@property
def num_classes(self):
return 10
class CIFAR100(datasets.CIFAR100):
def __init__(self, root, train=True, transform=None, target_transform=None, download=False):
super(CIFAR100, self).__init__(root, train=train, transform=transform,
target_transform=target_transform, download=download)
# unify the interface
if not hasattr(self, 'data'): # torch <= 0.4.1
if self.train:
self.data, self.targets = self.train_data, self.train_labels
else:
self.data, self.targets = self.test_data, self.test_labels
def __getitem__(self, index):
img, target = self.data[index], self.targets[index]
# doing this so that it is consistent with all other datasets
# to return a PIL Image
img = Image.fromarray(img)
if self.transform is not None:
img = self.transform(img)
if self.target_transform is not None:
target = self.target_transform(target)
return img, target, index
@property
def num_classes(self):
return 100
class SVHN(datasets.SVHN):
def __init__(self, root, split='train', transform=None, target_transform=None, download=False):
super(SVHN, self).__init__(root, split=split, transform=transform,
target_transform=target_transform, download=download)
# unify the interface
if not hasattr(self, 'data'): # torch <= 0.4.1
if self.train:
self.data, self.targets = self.train_data, self.train_labels
else:
self.data, self.targets = self.test_data, self.test_labels
def __getitem__(self, index):
img, target = self.data[index], int(self.labels[index])
# doing this so that it is consistent with all other datasets
# to return a PIL Image
img = Image.fromarray(np.transpose(img, (1, 2, 0)))
if self.transform is not None:
img = self.transform(img)
if self.target_transform is not None:
target = self.target_transform(target)
return img, target, index
@property
def num_classes(self):
return 10