# -*- coding: utf-8 -*-

import cv2
import numpy as np
import caffe
import sys

def main(args):

    caffe.set_mode_gpu()
    model_def = './deploy_finetuned.prototxt'
    model_weights = './deep_river_finetuned.caffemodel'

    net = caffe.Net(model_def, model_weights, caffe.TEST)
    
    I = cv2.imread('./input2.png', 0)
    I = I.astype('float')/np.max(I)

    h, w = I.shape
    net.blobs['data'].reshape(1, 1, h, w)
    
    # Forward propagate the image in the network
    net.blobs['data'].data[...] = I
    net.forward()
    
    cv2.imwrite("centerline.jpg", net.blobs['prob'].data[0][1,:,:]*255)
    
if __name__ == '__main__':
    main(sys.argv)