-
Notifications
You must be signed in to change notification settings - Fork 0
/
data.py
129 lines (100 loc) · 3.85 KB
/
data.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
"""
Copyright (C) 2018 NVIDIA Corporation. All rights reserved.
Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
"""
import torch.utils.data as data
import os.path
def default_loader(path):
return Image.open(path).convert('RGB')
def default_flist_reader(flist):
"""
flist format: impath label\nimpath label\n ...(same to caffe's filelist)
"""
imlist = []
with open(flist, 'r') as rf:
for line in rf.readlines():
impath = line.strip()
imlist.append(impath)
return imlist
class ImageFilelist(data.Dataset):
def __init__(self, root, flist, transform=None,
flist_reader=default_flist_reader, loader=default_loader):
self.root = root
self.imlist = flist_reader(flist)
self.transform = transform
self.loader = loader
def __getitem__(self, index):
impath = self.imlist[index]
img = self.loader(os.path.join(self.root, impath))
if self.transform is not None:
img = self.transform(img)
return img
def __len__(self):
return len(self.imlist)
class ImageLabelFilelist(data.Dataset):
def __init__(self, root, flist, transform=None,
flist_reader=default_flist_reader, loader=default_loader):
self.root = root
self.imlist = flist_reader(os.path.join(self.root, flist))
self.transform = transform
self.loader = loader
self.classes = sorted(list(set([path.split('/')[0] for path in self.imlist])))
self.class_to_idx = {self.classes[i]: i for i in range(len(self.classes))}
self.imgs = [(impath, self.class_to_idx[impath.split('/')[0]]) for impath in self.imlist]
def __getitem__(self, index):
impath, label = self.imgs[index]
img = self.loader(os.path.join(self.root, impath))
if self.transform is not None:
img = self.transform(img)
return img, label
def __len__(self):
return len(self.imgs)
###############################################################################
# Code from
# https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py
# Modified the original code so that it also loads images from the current
# directory as well as the subdirectories
###############################################################################
import torch.utils.data as data
from PIL import Image
import os
import os.path
IMG_EXTENSIONS = [
'.jpg', '.JPG', '.jpeg', '.JPEG',
'.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
]
def is_image_file(filename):
return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
def make_dataset(dir):
images = []
assert os.path.isdir(dir), '%s is not a valid directory' % dir
for root, _, fnames in sorted(os.walk(dir)):
for fname in fnames:
if is_image_file(fname):
path = os.path.join(root, fname)
images.append(path)
return images
class ImageFolder(data.Dataset):
def __init__(self, root, transform=None, return_paths=False,
loader=default_loader):
imgs = sorted(make_dataset(root))
if len(imgs) == 0:
raise(RuntimeError("Found 0 images in: " + root + "\n"
"Supported image extensions are: " +
",".join(IMG_EXTENSIONS)))
self.root = root
self.imgs = imgs
self.transform = transform
self.return_paths = return_paths
self.loader = loader
def __getitem__(self, index):
path = self.imgs[index]
img = self.loader(path)
if self.transform is not None:
img = self.transform(img)
if self.return_paths:
return img, path
else:
return img
def __len__(self):
return len(self.imgs)