CAFFE2_ROOT = "/caffe2"
CAFFE_MODELS = "/usr/local/caffe2/python/models"

import numpy as np
import skimage.io
import skimage.transform
import os

from caffe2.proto import caffe2_pb2
from caffe2.python import core, workspace

print("Required modules imported.")

# modify /usr/local/caffe2/python/models/download.py
# DOWNLOAD_BASE_URL = "https://github.com/caffe2/models/raw/master/"
# for f in ['predict_net.pb', 'exec_net.pb']:

# https://github.com/caffe2/models/tree/master/squeezenet
# $ python -m caffe2.python.models.download -i squeezenet
MODEL = 'squeezenet', 'exec_net.pb', 'predict_net.pb'

IMAGE = 'cat.jpg' # from http://cdn1-www.cattime.com/assets/uploads/gallery/persian-cats-and-kittens/persian-cats-and-kittens-8.jpg

INIT_NET = os.path.join(CAFFE_MODELS, MODEL[0], MODEL[1])
print 'INIT_NET =', INIT_NET
PREDICT_NET = os.path.join(CAFFE_MODELS, MODEL[0], MODEL[2])
print 'PREDICT_NET =' , PREDICT_NET

with open(INIT_NET) as f:
        init_net = f.read()
with open(PREDICT_NET) as f:
        predict_net = f.read()


mean = 128
print "mean was set to: ", mean

INPUT_IMAGE_SIZE = 227

def crop_center(img,cropx,cropy):
    y,x,c = img.shape
    startx = x//2-(cropx//2)
    starty = y//2-(cropy//2)
    return img[starty:starty+cropy,startx:startx+cropx]

def rescale(img, input_height, input_width):
    print("Original image shape:" + str(img.shape) + " and remember it should be in H, W, C!")
    print("Model's input shape is %dx%d") % (input_height, input_width)
    aspect = img.shape[1]/float(img.shape[0])
    print("Orginal aspect ratio: " + str(aspect))
    if(aspect>1):
        # landscape orientation - wide image
        res = int(aspect * input_height)
        imgScaled = skimage.transform.resize(img, (input_width, res))
    if(aspect<1):
        # portrait orientation - tall image
        res = int(input_width/aspect)
        imgScaled = skimage.transform.resize(img, (res, input_height))
    if(aspect == 1):
        imgScaled = skimage.transform.resize(img, (input_width, input_height))
    #pyplot.figure()
    #pyplot.imshow(imgScaled)
    #pyplot.axis('on')
    #pyplot.title('Rescaled image')
    print("New image shape:" + str(imgScaled.shape) + " in HWC")
    return imgScaled

img = skimage.img_as_float(skimage.io.imread(IMAGE)).astype(np.float32)
img = skimage.img_as_float(skimage.io.imread(IMAGE)).astype(np.float32)
img = rescale(img, INPUT_IMAGE_SIZE, INPUT_IMAGE_SIZE)
img = crop_center(img, INPUT_IMAGE_SIZE, INPUT_IMAGE_SIZE)
print("New image shape:" + str(img.shape) + " in HWC")

# switch to CHW
img = img.swapaxes(1, 2).swapaxes(0, 1)

# switch to BGR
img = img[(2, 1, 0), :, :]

# remove mean for better results
img = img * 255 - mean

# add batch size
img = img[np.newaxis, :, :, :].astype(np.float32)

#workspace.FeedBlob("data", [img])
#workspace.FeedBlob("data", np.random.rand(16, 100).astype(np.float32))

# net_def = caffe2_pb2.NetDef()
# net_def.ParseFromString(init_net)
# init_net = net_def.SerializeToString()

# net_def2 = caffe2_pb2.NetDef()
# net_def2.ParseFromString(predict_net)
# predict_net = net_def2.SerializeToString()

# workspace.FeedBlob('data', np.random.rand(1, 3, 227, 227).astype(np.float32))

p = workspace.Predictor(init_net, predict_net)

results = p.run([img])
results = np.asarray(results)
print "results shape: ", results.shape


codes = dict()
with open('alexnet_codes') as f:
    for line in f.readlines():
        code, result = line.partition(":")[::2]
        #print code, result
        codes[code.strip()] = result.strip()[1:-2]

results = np.delete(results, 1)
print "results shape: ", results.shape

index = 0
highest = 0
arr = np.empty((0, 2), dtype=object)
arr[:,0] = int(10)
arr[:,1:] = float(10)

for i, r in enumerate(results):
    i = i + 1
    print i, '{:.6f}'.format(r), codes[str(i)]
    arr = np.append(arr, np.array([[i, r]]), axis = 0)
    if r > highest:
        highest = r
        index = i
        
print index, " :: ", highest, codes[str(index)]

# $ python example.py
# WARNING:root:This caffe2 python run does not have GPU support. Will run in CPU only mode.
# WARNING:root:Debug message: No module named caffe2_pybind11_state_gpu
# Required modules imported.
# INIT_NET = /usr/local/caffe2/python/models/squeezenet/exec_net.pb
# PREDICT_NET = /usr/local/caffe2/python/models/squeezenet/predict_net.pb
# mean was set to:  128
# Original image shape:(680, 680, 3) and remember it should be in H, W, C!
# Model's input shape is 227x227
# Orginal aspect ratio: 1.0
# /usr/local/lib/python2.7/dist-packages/skimage/transform/_warps.py:84: UserWarning: The default mode, 'constant', will be changed to 'reflect' in skimage 0.15.
#  warn("The default mode, 'constant', will be changed to 'reflect' in "
# New image shape:(227, 227, 3) in HWC
# New image shape:(227, 227, 3) in HWC
# WARNING: Logging before InitGoogleLogging() is written to STDERR
# W1013 07:50:50.003103  1254 conv_pool_op_base.h:554] You are hitting a case where Caffe's legacy padding calculation is hit. This leads to inefficient and sometimes incorrect results. We are keeping this behavior for backward compatibility, but you are strongly recommended to move away from it.
# W1013 07:50:50.003170  1254 conv_pool_op_base.h:554] You are hitting a case where Caffe's legacy padding calculation is hit. This leads to inefficient and sometimes incorrect results. We are keeping this behavior for backward compatibility, but you are strongly recommended to move away from it.
# W1013 07:50:50.064188  1254 conv_pool_op_base.h:554] You are hitting a case where Caffe's legacy padding calculation is hit. This leads to inefficient and sometimes incorrect results. We are keeping this behavior for backward compatibility, but you are strongly recommended to move away from it.
# W1013 07:50:50.064229  1254 conv_pool_op_base.h:554] You are hitting a case where Caffe's legacy padding calculation is hit. This leads to inefficient and sometimes incorrect results. We are keeping this behavior for backward compatibility, but you are strongly recommended to move away from it.
# results shape:  (1, 1, 1000, 1, 1)
# results shape:  (999,)
# 283  ::  0.999376 Persian cat