import cv2
import numpy as np
from labels import LABEL_MAP # 1000 labels in imagenet dataset

# caffe model, googlenet aglo
weights = "bvlc_googlenet.caffemodel"
protxt = "bvlc_googlenet.prototxt"

# read caffe model from disk
net = cv2.dnn.readNetFromCaffe(protxt, weights)

# create input
image = cv2.imread("MY_TEST/ocean-liner.jpg")
blob = cv2.dnn.blobFromImage(image, 1.0, (224, 224), (104, 117, 123), False, crop=False)
result = np.copy(image)

# run!
net.setInput(blob)
out = net.forward()

# output probability, find the right index
out = out.flatten()
classId = np.argmax(out)
confidence = out[classId]

# time cost
t, _ = net.getPerfProfile()
label = 'cost time: %.2f ms' % (t * 1000.0 / cv2.getTickFrequency())
cv2.putText(result, label, (0, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 0), 2)

# render on image
label = '%s: %.4f' % (LABEL_MAP[classId] if LABEL_MAP else 'Class #%d' % classId, confidence)
cv2.putText(result, label, (0, 60), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

show_img = np.hstack((image, result))

# normal codes in opencv
cv2.imshow("Image", show_img)
cv2.waitKey(0)