def predict(dataset_name, input_path, output_path):
dataset = Dataset(dataset_name)
label_margin = 186
# Create theano graph
input_var = T.tensor4('input')
net = build_model(input_var)
outputs = lasagne.layers.get_output(net['prob'], deterministic=True)
fn = theano.function([input_var], outputs)
# Set the parameters from lasagne
f = open(dataset.pretrained_path, 'rb')
params = pickle.load(f)
[p.set_value(pval) for (p, pval) in zip(lasagne.layers.get_all_params(net['prob']), params)]
# Image processing
input_dims = dataset.shape
batch_size, num_channels, input_height, input_width = input_dims
image = cv2.imread(input_path, 1).astype(np.float32) - dataset.mean_pixel
image_size = image.shape
output_height = input_height - 2 * label_margin
output_width = input_width - 2 * label_margin
image = cv2.copyMakeBorder(image, label_margin, label_margin,
label_margin, label_margin,
cv2.BORDER_REFLECT_101)
num_tiles_h = image_size[0] // output_height + \
(1 if image_size[0] % output_height else 0)
num_tiles_w = image_size[1] // output_width + \
(1 if image_size[1] % output_width else 0)
prediction = []
for h in range(num_tiles_h):
col_prediction = []
for w in range(num_tiles_w):
offset = [output_height * h,
output_width * w]
tile = image[offset[0]:offset[0] + input_height,
offset[1]:offset[1] + input_width, :]
margin = [0, input_height - tile.shape[0],
0, input_width - tile.shape[1]]
tile = cv2.copyMakeBorder(tile, margin[0], margin[1],
margin[2], margin[3],
cv2.BORDER_REFLECT_101)
lasagne_in = tile.transpose([2, 0, 1])
# Get theano graph prediction
prob = fn(np.asarray([lasagne_in]))
col_prediction.append(prob)
col_prediction = np.concatenate(col_prediction, axis=1)
prediction.append(col_prediction)
prob = np.concatenate(prediction, axis=1).transpose().reshape((21, 66, 66))
if dataset.zoom > 1:
prob = interp_map(prob, dataset.zoom, image_size[1], image_size[0])
prediction = np.argmax(prob.transpose([1, 2, 0]), axis=2)
# Save the segmentation prediction
color_image = dataset.palette[prediction.ravel()].reshape(image_size)
color_image = cv2.cvtColor(color_image, cv2.COLOR_RGB2BGR)
print('Writing', output_path)
cv2.imwrite(output_path, color_image)
def convert(dataset_name):
dataset = Dataset(dataset_name)
# Create theano graph
input_var = T.tensor4('input')
net = build_model(input_var)
# Load caffe model
net_caffe = caffe.Net(dataset.model_path, dataset.pretrained_path, caffe.TEST)
# Set the parameters from caffe into lasagne
load_caffe_model(net, net_caffe)
# Save the parameters
p = join(dirname(__file__), 'pretrained', dataset.model_name + '.pkl')
output = open(p, 'wb')
params = get_all_param_values(net['prob'])
pickle.dump(params, output)
output.close()
