Ejemplo n.º 1
0
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)
Ejemplo n.º 2
0
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()