-
Notifications
You must be signed in to change notification settings - Fork 1
/
datasetsHelper.py
91 lines (74 loc) · 3.63 KB
/
datasetsHelper.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
import os
from torch.utils.data.dataset import Subset
from torchvision import datasets, transforms
root_path = os.path.dirname(__file__)
def get_subclass_dataset(dataset, classes):
if not isinstance(classes, list):
classes = [classes]
indices = []
for idx, data in enumerate(dataset):
if data[1] in classes:
indices.append(idx)
dataset = Subset(dataset, indices)
return dataset
def get_dataset(dataset, train_transform, test_transform, download=False, seen=None):
"""Get datasets for setting 1 (OOD Detection on the Same Dataset)."""
if dataset == 'cifar10':
DATA_PATH = os.path.join(root_path, 'cifar10')
class_idx = [int(num) for num in seen]
for i in range(10):
if i in class_idx:
continue
class_idx.append(i)
train_set = datasets.CIFAR10(DATA_PATH, train=True, download=download, transform=train_transform,
target_transform=lambda x: class_idx.index(x))
test_set = datasets.CIFAR10(DATA_PATH, train=False, download=download, transform=test_transform,
target_transform=lambda x: class_idx.index(x))
seen_class_idx = [0, 1, 2, 3, 4, 5]
unseen_class_idx = [6, 7, 8, 9]
train_set = get_subclass_dataset(train_set, seen_class_idx)
test_set_seen = get_subclass_dataset(test_set, seen_class_idx)
test_set_unseen = get_subclass_dataset(test_set, unseen_class_idx)
elif dataset == 'mnist':
DATA_PATH = os.path.join(root_path, 'mnist')
class_idx = [int(num) for num in seen]
for i in range(10):
if i in class_idx:
continue
class_idx.append(i)
train_set = datasets.MNIST(DATA_PATH, train=True, download=download, transform=train_transform,
target_transform=lambda x: class_idx.index(x))
test_set = datasets.MNIST(DATA_PATH, train=False, download=download, transform=test_transform,
target_transform=lambda x: class_idx.index(x))
seen_class_idx = [0, 1, 2, 3, 4, 5]
unseen_class_idx = [6, 7, 8, 9]
train_set = get_subclass_dataset(train_set, seen_class_idx)
test_set_seen = get_subclass_dataset(test_set, seen_class_idx)
test_set_unseen = get_subclass_dataset(test_set, unseen_class_idx)
else:
raise NotImplementedError
return train_set, test_set_seen, test_set_unseen
def get_ood_dataset(dataset):
"""Get datasets for setting 2 (OOD Detection on Different Datasets)."""
if dataset == 'SVHN':
dir = os.path.join(root_path, 'svhn')
data = datasets.SVHN(root=dir, split='test', download=True, transform=transforms.ToTensor())
elif dataset == 'LSUN':
dir = os.path.join(root_path, 'LSUN_resize')
data = datasets.ImageFolder(dir, transform=transforms.ToTensor())
elif dataset == 'tinyImageNet':
dir = os.path.join(root_path, 'Imagenet_resize')
data = datasets.ImageFolder(dir, transform=transforms.ToTensor())
elif dataset == 'LSUN-FIX':
dir = os.path.join(root_path, 'LSUN_fix')
data = datasets.ImageFolder(dir, transform=transforms.ToTensor())
elif dataset == 'ImageNet-FIX':
dir = os.path.join(root_path, 'Imagenet_fix')
data = datasets.ImageFolder(dir, transform=transforms.ToTensor())
elif dataset == 'CIFAR100':
dir = os.path.join(root_path, 'cifar100')
data = datasets.CIFAR100(
root=dir, train=False, transform=transforms.ToTensor(), download=True)
else:
raise NotImplementedError
return data