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()