-
Notifications
You must be signed in to change notification settings - Fork 2
/
utils.py
220 lines (164 loc) · 7.04 KB
/
utils.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
212
213
214
215
216
217
218
219
220
import os
import sys
import random
import numpy as np
import scipy.misc
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from PIL import Image
class ImagePool(object):
def __init__(self, pool_size=50):
self.pool_size = pool_size
self.imgs = []
def query(self, img):
if self.pool_size == 0:
return img
if len(self.imgs) < self.pool_size:
self.imgs.append(img)
return img
else:
if random.random() > 0.5:
# use old image
random_id = random.randrange(0, self.pool_size)
tmp_img = self.imgs[random_id].copy()
self.imgs[random_id] = img.copy()
return tmp_img
else:
return img
def all_files_under(path, extension=None, append_path=True, sort=True):
if append_path:
if extension is None:
filenames = [os.path.join(path, fname) for fname in os.listdir(path)]
else:
filenames = [os.path.join(path, fname)
for fname in os.listdir(path) if fname.endswith(extension)]
else:
if extension is None:
filenames = [os.path.basename(fname) for fname in os.listdir(path)]
else:
filenames = [os.path.basename(fname)
for fname in os.listdir(path) if fname.endswith(extension)]
if sort:
filenames = sorted(filenames)
return filenames
def imagefiles2arrs(filenames):
img_shape = image_shape(filenames[0])
images_arr = None
if len(img_shape) == 3: # color image
images_arr = np.zeros((len(filenames), img_shape[0], img_shape[1], img_shape[2]), dtype=np.float32)
elif len(img_shape) == 2: # gray scale image
images_arr = np.zeros((len(filenames), img_shape[0], img_shape[1]), dtype=np.float32)
for file_index in range(len(filenames)):
img = Image.open(filenames[file_index])
images_arr[file_index] = np.asarray(img).astype(np.float32)
return images_arr
def image_shape(filename):
img = Image.open(filename, mode="r")
img_arr = np.asarray(img)
img_shape = img_arr.shape
return img_shape
def print_metrics(itr, kargs):
print("*** Iteration {} ====> ".format(itr))
for name, value in kargs.items():
print("{} : {}, ".format(name, value))
print("")
sys.stdout.flush()
def transform(img):
return img / 127.5 - 1.0
def inverse_transform(img):
return (img + 1.) / 2.
def preprocess_pair(img_a, img_b, load_size=286, fine_size=256, flip=True, is_test=False):
if is_test:
img_a = scipy.misc.imresize(img_a, [fine_size, fine_size])
img_b = scipy.misc.imresize(img_b, [fine_size, fine_size])
else:
img_a = scipy.misc.imresize(img_a, [load_size, load_size])
img_b = scipy.misc.imresize(img_b, [load_size, load_size])
h1 = int(np.ceil(np.random.uniform(1e-2, load_size - fine_size)))
w1 = int(np.ceil(np.random.uniform(1e-2, load_size - fine_size)))
img_a = img_a[h1:h1 + fine_size, w1:w1 + fine_size]
img_b = img_b[h1:h1 + fine_size, w1:w1 + fine_size]
if flip and np.random.random() > 0.5:
img_a = np.fliplr(img_a)
img_b = np.fliplr(img_b)
return img_a, img_b
def imread(path, is_gray_scale=False, img_size=None):
if is_gray_scale:
img = scipy.misc.imread(path, flatten=True).astype(np.float)
else:
img = scipy.misc.imread(path, mode='RGB').astype(np.float)
if not (img.ndim == 3 and img.shape[2] == 3):
img = np.dstack((img, img, img))
if img_size is not None:
img = scipy.misc.imresize(img, img_size)
return img
def load_image(image_path, which_direction=0, is_gray_scale=True, img_size=(256, 256, 1)):
input_img = imread(image_path, is_gray_scale=is_gray_scale, img_size=img_size)
w_pair = int(input_img.shape[1])
w_single = int(w_pair / 2)
if which_direction == 0: # A to B
img_a = input_img[:, 0:w_single]
img_b = input_img[:, w_single:w_pair]
else: # B to A
img_a = input_img[:, w_single:w_pair]
img_b = input_img[:, 0:w_single]
return img_a, img_b
def load_data(image_path, is_gray_scale=False):
img = imread(path=image_path, is_gray_scale=is_gray_scale)
img_trans = transform(img) # from [0, 255] to [-1., 1.]
if is_gray_scale and (img_trans.ndim == 2):
img_trans = np.expand_dims(img_trans, axis=2)
return img_trans
def plots(imgs, iter_time, save_file, grid_cols, grid_rows, sample_batch, name=None):
# parameters for plot size
scale, margin = 0.02, 0.02
# save more bigger image
img_h, img_w, img_c = imgs.shape[1:]
fig = plt.figure(figsize=(img_w * grid_cols * scale, img_h * grid_rows * scale)) # (column, row)
gs = gridspec.GridSpec(grid_rows, grid_cols) # (row, column)
gs.update(wspace=margin, hspace=margin)
for img_idx in range(sample_batch):
ax = plt.subplot(gs[img_idx])
plt.axis('off')
ax.set_xticklabels([])
ax.set_yticklabels([])
ax.set_aspect('equal')
if imgs[img_idx].shape[2] == 1: # gray scale
plt.imshow((imgs[img_idx]).reshape(img_h, img_w), cmap='Greys_r')
else:
plt.imshow((imgs[img_idx]).reshape(img_h, img_w, img_c), cmap='Greys_r')
plt.savefig(save_file + '/{}_{}.png'.format(str(iter_time), name), bbox_inches='tight')
plt.close(fig)
def _merge(images, size, resize_ratio=1.):
h, w = images.shape[1], images.shape[2]
h_ = int(h * resize_ratio)
w_ = int(w * resize_ratio)
img_canvas = np.zeros((h_ * size[0], w_ * size[1]))
for idx, image in enumerate(images):
i = int(idx % size[1])
j = int(idx / size[1])
image_resize = scipy.misc.imresize(image, size=(h_, w_), interp='bicubic')
img_canvas[j * h_:j * h_ + h_, i * w_:i * w_ + w_] = image_resize
return img_canvas
# borrowed from https://github.com/ykwon0407/variational_autoencoder/blob/master/variational_bayes.ipynb
def save_scattered_image(z, id_, iter_time, save_file, z_range=2):
num_labels = 10
plt.figure(figsize=(8, 6))
plt.scatter(z[:, 0], z[:, 1], c=np.argmax(id_, 1), marker='o', edgecolor='none',
cmap=discrete_cmap(num_labels, 'jet'))
plt.colorbar(ticks=range(num_labels))
axes = plt.gca()
axes.set_xlim([-z_range-2, z_range+2])
axes.set_ylim([-z_range-2, z_range+2])
plt.grid(True)
plt.savefig(os.path.join(save_file, 'embedding_{}.png'.format(str(iter_time))))
# borrowed from https://gist.github.com/jakevdp/91077b0cae40f8f8244a
def discrete_cmap(N, base_cmap=None):
"""Create an N-bin discrete colormap from the specified input map"""
# Note that if base_cmap is a string or None, you can simply do
# return plt.cm.get_cmap(base_cmap, N)
# The following works for string, None, or a colormap instance:
base = plt.cm.get_cmap(base_cmap)
color_list = base(np.linspace(0, 1, N))
cmap_name = base.name + str(N)
return base.from_list(cmap_name, color_list, N)