forked from BR-IDL/PaddleViT
-
Notifications
You must be signed in to change notification settings - Fork 0
/
datasets.py
211 lines (179 loc) · 7.96 KB
/
datasets.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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
# Copyright (c) 2021 PPViT Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Dataset related classes and methods for Styleformer training and validation
Cifar10, STL10Dataset, LSUNchurch, Celeba and ImageNet2012 are supported
"""
import os
import math
from paddle.io import Dataset
from paddle.io import DataLoader
from paddle.io import DistributedBatchSampler
from paddle.vision import transforms
from paddle.vision import datasets
from paddle.vision import image_load
from stl10_dataset import STL10Dataset
from lsun_church_dataset import LSUNchurchDataset
from celeba_dataset import CelebADataset
class ImageNet2012Dataset(Dataset):
"""Build ImageNet2012 dataset
This class gets train/val imagenet datasets, which loads transfomed data and labels.
Attributes:
file_folder: path where imagenet images are stored
transform: preprocessing ops to apply on image
img_path_list: list of full path of images in whole dataset
label_list: list of labels of whole dataset
"""
def __init__(self, file_folder, mode="train", transform=None):
"""Init ImageNet2012 Dataset with dataset file path, mode(train/val), and transform"""
super(ImageNet2012Dataset, self).__init__()
assert mode in ["train", "val"]
self.file_folder = file_folder
self.transform = transform
self.img_path_list = []
self.label_list = []
if mode == "train":
self.list_file = os.path.join(self.file_folder, "train_list.txt")
else:
self.list_file = os.path.join(self.file_folder, "val_list.txt")
with open(self.list_file, 'r') as infile:
for line in infile:
img_path = line.strip().split()[0]
img_label = int(line.strip().split()[1])
self.img_path_list.append(os.path.join(self.file_folder, img_path))
self.label_list.append(img_label)
print(f'----- Imagenet2012 image {mode} list len = {len(self.label_list)}')
def __len__(self):
return len(self.label_list)
def __getitem__(self, index):
data = image_load(self.img_path_list[index]).convert('RGB')
data = self.transform(data)
label = self.label_list[index]
return data, label
def get_train_transforms(config):
""" Get training transforms
For training, a RandomResizedCrop is applied, then normalization is applied with
[0.5, 0.5, 0.5] mean and std. The input pixel values must be rescaled to [0, 1.]
Outputs is converted to tensor
Args:
config: configs contains IMAGE_SIZE, see config.py for details
Returns:
transforms_train: training transforms
"""
scale_size = int(math.floor(config.DATA.IMAGE_SIZE / config.DATA.CROP_PCT))
transforms_train = transforms.Compose([
transforms.Resize(scale_size, interpolation='lanczos'),
transforms.ToTensor(),
])
return transforms_train
def get_val_transforms(config):
""" Get training transforms
For validation, image is first Resize then CenterCrop to image_size.
Then normalization is applied with [0.5, 0.5, 0.5] mean and std.
The input pixel values must be rescaled to [0, 1.]
Outputs is converted to tensor
Args:
config: configs contains IMAGE_SIZE, see config.py for details
Returns:
transforms_train: training transforms
"""
scale_size = int(math.floor(config.DATA.IMAGE_SIZE / config.DATA.CROP_PCT))
transforms_val = transforms.Compose([
transforms.Resize([scale_size, scale_size], interpolation='lanczos'),
transforms.ToTensor(),
])
return transforms_val
def get_dataset(config, mode='train'):
""" Get dataset from config and mode (train/val)
Returns the related dataset object according to configs and mode(train/val)
Args:
config: configs contains dataset related settings. see config.py for details
Returns:
dataset: dataset object
"""
assert mode in ['train', 'val']
if config.DATA.DATASET == "cifar10":
if mode == 'train':
dataset = datasets.Cifar10(mode=mode, transform=get_train_transforms(config))
else:
mode = 'train'
dataset = datasets.Cifar10(mode=mode, transform=get_val_transforms(config))
elif config.DATA.DATASET == "stl10":
if mode == 'train':
dataset = STL10Dataset(file_folder=config.DATA.DATA_PATH,
mode=mode,
transform=get_train_transforms(config))
else:
#mode = 'test'
mode = 'unlabeled'
dataset = STL10Dataset(file_folder=config.DATA.DATA_PATH,
mode=mode,
transform=get_val_transforms(config))
elif config.DATA.DATASET == "lsun":
if mode == 'train':
dataset = LSUNchurchDataset(file_folder=config.DATA.DATA_PATH,
mode=mode,
transform=get_train_transforms(config))
else:
dataset = LSUNchurchDataset(file_folder=config.DATA.DATA_PATH,
mode=mode,
transform=get_val_transforms(config))
elif config.DATA.DATASET == "celeba":
if mode == 'train':
dataset = CelebADataset(file_folder=config.DATA.DATA_PATH,
transform=get_train_transforms(config))
else:
dataset = CelebADataset(file_folder=config.DATA.DATA_PATH,
transform=get_val_transforms(config))
elif config.DATA.DATASET == "imagenet2012":
if mode == 'train':
dataset = ImageNet2012Dataset(config.DATA.DATA_PATH,
mode=mode,
transform=get_train_transforms(config))
else:
dataset = ImageNet2012Dataset(config.DATA.DATA_PATH,
mode=mode,
transform=get_val_transforms(config))
else:
raise NotImplementedError(
"Only support cifar10, cifar100, imagenet2012, celeba, stl10, lsun")
return dataset
def get_dataloader(config, dataset, mode='train', multi_process=False):
"""Get dataloader with config, dataset, mode as input, allows multiGPU settings.
Multi-GPU loader is implements as distributedBatchSampler.
Args:
config: see config.py for details
dataset: paddle.io.dataset object
mode: train/val
multi_process: if True, use DistributedBatchSampler to support multi-processing
Returns:
dataloader: paddle.io.DataLoader object.
"""
if mode == 'train':
batch_size = config.DATA.BATCH_SIZE
else:
batch_size = config.DATA.BATCH_SIZE_EVAL
if multi_process is True:
sampler = DistributedBatchSampler(dataset,
batch_size=batch_size,
shuffle=(mode == 'train'))
dataloader = DataLoader(dataset,
batch_sampler=sampler,
num_workers=config.DATA.NUM_WORKERS)
else:
dataloader = DataLoader(dataset,
batch_size=batch_size,
num_workers=config.DATA.NUM_WORKERS,
shuffle=(mode == 'train'))
return dataloader