-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
54 lines (40 loc) · 1.81 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
import os
import numpy as np
import cv2
import theano
import lasagne
from constants import HOME_DIR
def chunks(l, n):
"""Yield successive n-sized chunks from l."""
for i in xrange(0, len(l), n):
yield l[i:i + n]
def load_weights(net, path, epochtoload):
"""
Load a pretrained model
:param epochtoload: epoch to load
:param net: model object
:param path: path of the weights to be set
"""
with np.load(HOME_DIR + path + "modelWeights{:04d}.npz".format(epochtoload)) as f:
param_values = [f['arr_%d' % i] for i in range(len(f.files))]
lasagne.layers.set_all_param_values(net, param_values)
def predict(model, image_stimuli, num_epoch=None, name=None, path_output_maps=None):
size = (image_stimuli.shape[1], image_stimuli.shape[0])
blur_size = 5
if image_stimuli.shape[:2] != (model.inputHeight, model.inputWidth):
image_stimuli = cv2.resize(image_stimuli, (model.inputWidth, model.inputHeight), interpolation=cv2.INTER_AREA)
blob = np.zeros((1, 3, model.inputHeight, model.inputWidth), theano.config.floatX)
blob[0, ...] = (image_stimuli.astype(theano.config.floatX).transpose(2, 0, 1))
result = np.squeeze(model.predictFunction(blob))
saliency_map = (result * 255).astype(np.uint8)
# resize back to original size
saliency_map = cv2.resize(saliency_map, size, interpolation=cv2.INTER_CUBIC)
# blur
saliency_map = cv2.GaussianBlur(saliency_map, (blur_size, blur_size), 0)
# clip again
saliency_map = np.clip(saliency_map, 0, 255)
if name is None:
# When we use for testing, there is no file name provided.
cv2.imwrite('./' + path_output_maps + '/validationRandomSaliencyPred_{:04d}.png'.format(num_epoch), saliency_map)
else:
cv2.imwrite(os.path.join(path_output_maps, name + '.jpg'), saliency_map)