import sys
#sys.path.insert(0, '/path/to/caffe/python')
import caffe
import cv2
import numpy as np
from caffe.proto import caffe_pb2

#Read mean image
mean_blob = caffe_pb2.BlobProto()
with open('examples/mnist/mnist_mean.binaryproto') as f:
    mean_blob.ParseFromString(f.read())
mean_array = np.asarray(mean_blob.data, dtype=np.float32).reshape(
    (mean_blob.channels, mean_blob.height, mean_blob.width))

#Create net. Load model
net = caffe.Net('examples/mnist/alexnet_deploy.prototxt','examples/mnist/alexnet_iter_10000.caffemodel', caffe.TEST)
print("net loaded")

##Preprocessing the image
transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
#The function of the transformer is to preprocess the input image and transform it into something that Caffe can understand.
#The mean image needs to be subtracted from each input image.
transformer.set_mean('data', mean_array)

#A couple of other params:
print("\n Before transpose net.blobs[data]shape",net.blobs['data'].data.shape)
transformer.set_transpose('data', (2,0,1))
transformer.set_raw_scale('data', 255.0)

#We need to reshape the blobs so that they match the image shape.
net.blobs['data'].reshape(1,1,28,28)
print("\nAfter reshape net.blobs[data]shape",net.blobs['data'].data.shape)

#Examining model parameters
print("------------Examining Blobs dimensions---------------\n")
# for each layer, show the output shape
for layer_name, blob in net.blobs.items():
    print(layer_name + '\t' + str(blob.data.shape))
print("------------contain the weights and biases of the layers------------------\n")
for layer_name, param in net.params.items():
    print(layer_name + '\t' + str(param[0].data.shape), str(param[1].data.shape))

#load the input image:
img = caffe.io.load_image('two(4).png',False)
net.blobs['data'].data[...] = transformer.preprocess('data', img)

### perform classification
output = net.forward()
output_prob = output['prob'][0]  # the output probability vector for the first image in the batch
# Text representation of the digit
digits_label = ["Zero","One","Two","Three","Four","Five","Six","Seven","Eight","Nine"]

#Find the index with the highest probablility
highest_index = -1
highest_probability = 0.1 #If nothing's more than 10% sure then don't print out anything

for i in range(0,10):
    if output_prob[i] > highest_probability:
        highest_index = i
        highest_probability = output_prob[i]

#Print our result
if highest_index < 0:
    print("\nDid not detect a number!")
else:
    print("\nDigit "+ str(digits_label[highest_index]) + " detected with " + str(highest_probability*100.0)+"%  probability.\n")
