-
Notifications
You must be signed in to change notification settings - Fork 32
/
Copy pathdata_loading_functions.py
161 lines (120 loc) · 6.01 KB
/
data_loading_functions.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
import numpy as np
import os
from skimage.color import gray2rgb
def train_val_shuffle_inds_per_epoch(valid_inds, train_inds, train_iter, batch_size, log_path, save_log=True):
"""shuffle image indices for each training epoch and save to log"""
np.random.seed(0)
num_train_images = len(train_inds)
num_epochs = int(np.ceil((1. * train_iter) / (1. * num_train_images / batch_size)))+1
epoch_inds_shuffle = np.zeros((num_epochs, num_train_images)).astype(int)
img_inds = np.arange(num_train_images)
for i in range(num_epochs):
np.random.shuffle(img_inds)
epoch_inds_shuffle[i, :] = img_inds
if save_log:
with open(os.path.join(log_path, "train_val_shuffle_inds.csv"), "wb") as f:
if valid_inds is not None:
f.write(b'valid inds\n')
np.savetxt(f, valid_inds.reshape(1, -1), fmt='%i', delimiter=",")
f.write(b'train inds\n')
np.savetxt(f, train_inds.reshape(1, -1), fmt='%i', delimiter=",")
f.write(b'shuffle inds\n')
np.savetxt(f, epoch_inds_shuffle, fmt='%i', delimiter=",")
return epoch_inds_shuffle
def gaussian(x, y, x0, y0, sigma=6):
return 1./(np.sqrt(2*np.pi)*sigma) * np.exp(-0.5 * ((x-x0)**2 + (y-y0)**2) / sigma**2)
def create_gaussian_filter(sigma=6, win_mult=3.5):
win_size = int(win_mult * sigma)
x, y = np.mgrid[0:2*win_size+1, 0:2*win_size+1]
gauss_filt = (8./3)*sigma*gaussian(x, y, win_size, win_size, sigma=sigma) # same as in ECT
return gauss_filt
def load_images(img_list, batch_inds, image_size=256, c_dim=3, scale=255):
""" load images as a numpy array from menpo image list """
num_inputs = len(batch_inds)
batch_menpo_images = img_list[batch_inds]
images = np.zeros([num_inputs, image_size, image_size, c_dim]).astype('float32')
for ind, img in enumerate(batch_menpo_images):
if img.n_channels < 3 and c_dim == 3:
images[ind, :, :, :] = gray2rgb(img.pixels_with_channels_at_back())
else:
images[ind, :, :, :] = img.pixels_with_channels_at_back()
if scale is 255:
images *= 255
elif scale is 0:
images = 2 * images - 1
return images
# loading functions with pre-allocation and approx heat-map generation
def create_approx_heat_maps_alloc_once(landmarks, maps, gauss_filt=None, win_mult=3.5, num_landmarks=68, image_size=256,
sigma=6):
""" create heatmaps from input landmarks"""
maps.fill(0.)
win_size = int(win_mult * sigma)
filt_size = 2 * win_size + 1
landmarks = landmarks.astype(int)
if gauss_filt is None:
x_small, y_small = np.mgrid[0:2 * win_size + 1, 0:2 * win_size + 1]
gauss_filt = (8. / 3) * sigma * gaussian(x_small, y_small, win_size, win_size, sigma=sigma) # same as in ECT
for i in range(num_landmarks):
min_row = landmarks[i, 0] - win_size
max_row = landmarks[i, 0] + win_size + 1
min_col = landmarks[i, 1] - win_size
max_col = landmarks[i, 1] + win_size + 1
if min_row < 0:
min_row_gap = -1 * min_row
min_row = 0
else:
min_row_gap = 0
if min_col < 0:
min_col_gap = -1 * min_col
min_col = 0
else:
min_col_gap = 0
if max_row > image_size:
max_row_gap = max_row - image_size
max_row = image_size
else:
max_row_gap = 0
if max_col > image_size:
max_col_gap = max_col - image_size
max_col = image_size
else:
max_col_gap = 0
maps[min_row:max_row, min_col:max_col, i] =\
gauss_filt[min_row_gap:filt_size - 1 * max_row_gap, min_col_gap:filt_size - 1 * max_col_gap]
def load_images_landmarks_approx_maps_alloc_once(
img_list, batch_inds, images, maps_small, maps, landmarks, image_size=256, num_landmarks=68,
scale=255, gauss_filt_large=None, gauss_filt_small=None, win_mult=3.5, sigma=6, save_landmarks=False):
""" load images and gt landmarks from menpo image list, and create matching heatmaps """
batch_menpo_images = img_list[batch_inds]
c_dim = images.shape[-1]
grp_name = batch_menpo_images[0].landmarks.group_labels[0]
win_size_large = int(win_mult * sigma)
win_size_small = int(win_mult * (1.*sigma/4))
if gauss_filt_small is None:
x_small, y_small = np.mgrid[0:2 * win_size_small + 1, 0:2 * win_size_small + 1]
gauss_filt_small = (8. / 3) * (1.*sigma/4) * gaussian(
x_small, y_small, win_size_small, win_size_small, sigma=1.*sigma/4) # same as in ECT
if gauss_filt_large is None:
x_large, y_large = np.mgrid[0:2 * win_size_large + 1, 0:2 * win_size_large + 1]
gauss_filt_large = (8. / 3) * sigma * gaussian(x_large, y_large, win_size_large, win_size_large, sigma=sigma) # same as in ECT
for ind, img in enumerate(batch_menpo_images):
if img.n_channels < 3 and c_dim == 3:
images[ind, :, :, :] = gray2rgb(img.pixels_with_channels_at_back())
else:
images[ind, :, :, :] = img.pixels_with_channels_at_back()
lms = img.landmarks[grp_name].points
lms = np.minimum(lms, image_size - 1)
create_approx_heat_maps_alloc_once(
landmarks=lms, maps=maps[ind, :, :, :], gauss_filt=gauss_filt_large, win_mult=win_mult,
num_landmarks=num_landmarks, image_size=image_size, sigma=sigma)
lms_small = img.resize([image_size / 4, image_size / 4]).landmarks[grp_name].points
lms_small = np.minimum(lms_small, image_size / 4 - 1)
create_approx_heat_maps_alloc_once(
landmarks=lms_small, maps=maps_small[ind, :, :, :], gauss_filt=gauss_filt_small, win_mult=win_mult,
num_landmarks=num_landmarks, image_size=image_size / 4, sigma=1. * sigma / 4)
if save_landmarks:
landmarks[ind, :, :] = lms
if scale is 255:
images *= 255
elif scale is 0:
images = 2 * images - 1