forked from hhb072/MOSS
-
Notifications
You must be signed in to change notification settings - Fork 0
/
dataset.py
205 lines (171 loc) · 7.93 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
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
import torch
import torch.utils.data as data
from os import listdir
from os.path import join
from PIL import Image, ImageOps
import random
import torchvision.transforms as transforms
import torchvision.transforms.functional as TF
def load_double_image(file_path, input_height=None, input_width=None, output_height=None, output_width=None,
crop_height=None, crop_width=None, is_random_crop=True, is_mirror=True, is_gray=False, random_scale=None):
if input_width is None:
input_width = input_height
if output_width is None:
output_width = output_height
if crop_width is None:
crop_width = crop_height
img = Image.open(file_path)
if is_gray is False and img.mode != 'RGB':
img = img.convert('RGB')
if is_gray and img.mode != 'L':
img = img.convert('L')
if random_scale is not None:
if random.random() < 0.2:
[w, h] = img.size
ww = random.randint(int(w*random_scale), w)
hh = random.randint(int(h*random_scale), h)
img = img.resize((ww, hh), Image.BICUBIC)
[w, h] = img.size
# print(img.size, flush=True)
imgR = ImageOps.crop(img, (0, 0, w//2, 0))
imgL = ImageOps.crop(img, (w//2, 0, 0, 0))
if is_mirror and random.randint(0,1) == 0:
imgR = ImageOps.mirror(imgR)
imgL = ImageOps.mirror(imgL)
if input_height is not None:
imgR = imgR.resize((input_width, input_height),Image.BICUBIC)
imgL = imgL.resize((input_width, input_height),Image.BICUBIC)
[w, h] = imgR.size
if crop_height is not None:
if is_random_crop:
#print([w,cropSize])
cx1 = random.randint(0, w-crop_width) if crop_width < w else 0
cx2 = w - crop_width - cx1
cy1 = random.randint(0, h-crop_height) if crop_height < h else 0
cy2 = h - crop_height - cy1
else:
cx2 = cx1 = int(round((w-crop_width)/2.))
cy2 = cy1 = int(round((h-crop_height)/2.))
imgR = ImageOps.crop(imgR, (cx1, cy1, cx2, cy2))
imgL = ImageOps.crop(imgL, (cx1, cy1, cx2, cy2))
if output_height is not None:
imgR = imgR.resize((output_width, output_height),Image.BICUBIC)
imgL = imgL.resize((output_width, output_height),Image.BICUBIC)
return imgR, imgL
class DoubleImageDataset(data.Dataset):
def __init__(self, image_list, root_path,
input_height=None, input_width=None, output_height=None, output_width=None,
crop_height=None, crop_width=None, is_random_crop=False,
is_mirror=True, is_gray=False, normalize=None, augment=1, random_scale=None):
super().__init__()
self.image_filenames = image_list
self.is_random_crop = is_random_crop
self.is_mirror = is_mirror
self.input_height = input_height
self.input_width = input_width
self.output_height = output_height
self.output_width = output_width
self.root_path = root_path
self.crop_height = crop_height
self.crop_width = crop_width
self.is_gray = is_gray
self.random_scale = random_scale
if normalize is None:
self.transform = transforms.Compose([transforms.ToTensor()])
else:
self.transform = transforms.Compose([transforms.ToTensor(), normalize])
self.length = len(self.image_filenames)
self.augment = augment
def __getitem__(self, index):
idx = index % self.length
imgR, imgL = load_double_image(join(self.root_path, self.image_filenames[idx]),
self.input_height, self.input_width, self.output_height, self.output_width,
self.crop_height, self.crop_width, self.is_random_crop, self.is_mirror, self.is_gray, random_scale=self.random_scale)
imgR = self.transform(imgR)
imgL = self.transform(imgL)
return imgR, imgL
def __len__(self):
return len(self.image_filenames) * self.augment
def load_single_image(file_path, input_height=128, input_width=None, output_height=128, output_width=None,
crop_height=None, crop_width=None, is_random_crop=True, is_mirror=True, is_gray=False, random_scale=None):
if input_width is None:
input_width = input_height
if output_width is None:
output_width = output_height
if crop_width is None:
crop_width = crop_height
img = Image.open(file_path)
if is_gray is False and img.mode is not 'RGB':
img = img.convert('RGB')
if is_gray and img.mode is not 'L':
img = img.convert('L')
if random_scale is not None:
if random.random() < 0.5:
[w, h] = img.size
ww = random.randint(int(w*random_scale), w)
hh = random.randint(int(h*random_scale), h)
img = img.resize((ww, hh), Image.BICUBIC)
if is_mirror and random.randint(0,1) is 0:
img = ImageOps.mirror(img)
if input_height is not None:
img = img.resize((input_width, input_height),Image.BICUBIC)
[w, h] = img.size
if crop_height is not None:
if is_random_crop:
#print([w,cropSize])
if crop_width < w:
cx1 = random.randint(0, w-crop_width)
cx2 = w - crop_width - cx1
else:
cx1, cx2 = 0, 0
if crop_height < h:
cy1 = random.randint(0, h-crop_height)
cy2 = h - crop_height - cy1
else:
cy1, cy2 = 0, 0
else:
if crop_width < w:
cx2 = cx1 = int(round((w-crop_width)/2.))
else:
cx2 = cx1 = 0
if crop_height < h:
cy2 = cy1 = int(round((h-crop_height)/2.))
else:
cy2 = cy1 = 0
img = ImageOps.crop(img, (cx1, cy1, cx2, cy2))
if output_height is not None:
img = img.resize((output_width, output_height),Image.BICUBIC)
return img
class SingleImageDataset(data.Dataset):
def __init__(self, image_list, root_path,
input_height=None, input_width=None, output_height=None, output_width=None,
crop_height=None, crop_width=None, is_random_crop=False,
is_mirror=True, is_gray=False, normalize=None, augment=1, random_scale=None):
super().__init__()
self.image_filenames = image_list
self.is_random_crop = is_random_crop
self.is_mirror = is_mirror
self.input_height = input_height
self.input_width = input_width
self.output_height = output_height
self.output_width = output_width
self.root_path = root_path
self.crop_height = crop_height
self.crop_width = crop_width
self.is_gray = is_gray
self.random_scale = random_scale
if normalize is None:
self.transform = transforms.Compose([transforms.ToTensor()])
else:
self.transform = transforms.Compose([transforms.ToTensor(), normalize])
self.length = len(self.image_filenames)
self.augment = augment
def __getitem__(self, index):
idx = index % self.length
img = load_single_image(join(self.root_path, self.image_filenames[idx]),
self.input_height, self.input_width, self.output_height, self.output_width,
self.crop_height, self.crop_width, self.is_random_crop, self.is_mirror, self.is_gray, random_scale=self.random_scale)
img = self.transform(img)
return img
def __len__(self):
return len(self.image_filenames) * self.augment