#! /usr/bin/env python3

# Copyright(c) 2017 Intel Corporation.
# License: MIT See LICENSE file in root directory.


from mvnc import mvncapi as mvnc
import numpy as np
import cv2
import sys
import time
import matplotlib.pyplot as plt

dim=(300,300)
EXAMPLES_BASE_DIR='../../'
IMAGES_DIR = EXAMPLES_BASE_DIR + 'data/images/'
IMAGE_FULL_PATH = IMAGES_DIR + '1.jpg'

# ***************************************************************
# Labels for the classifications for the network.
# ***************************************************************
LABELS = ('background',
          'aeroplane', 'bicycle', 'bird', 'boat',
          'bottle', 'bus', 'car', 'cat', 'chair',
          'cow', 'diningtable', 'dog', 'horse',
          'motorbike', 'person', 'pottedplant',
          'sheep', 'sofa', 'train', 'tvmonitor')


def run_inference(image_to_classify, ssd_mobilenet_graph):

    # get a resized version of the image that is the dimensions
    # SSD Mobile net expects
    resized_image = preprocess_image(image_to_classify)
    start=time.time()
    # ***************************************************************
    # Send the image to the NCS
    # ***************************************************************
    ssd_mobilenet_graph.LoadTensor(resized_image.astype(np.float16), None)

    # ***************************************************************
    # Get the result from the NCS
    # ***************************************************************
    detections, userobj = ssd_mobilenet_graph.GetResult()
    print (detections.shape)

    detections = np.reshape(detections,(1,1,201,7))


    # Parse the outputs.
    det_label = detections[0,0,:,1]
    det_conf = detections[0,0,:,2]
    det_xmin = detections[0,0,:,3]
    det_ymin = detections[0,0,:,4]
    det_xmax = detections[0,0,:,5]
    det_ymax = detections[0,0,:,6]

    # Get detections with confidence higher than 0.6.
    top_indices = [i for i, conf in enumerate(det_conf) if conf >= 0.05]

    top_conf = det_conf[top_indices]
    top_label_indices = det_label[top_indices].tolist()
    #top_labels = get_labelname(labelmap, top_label_indices)
    top_xmin = det_xmin[top_indices]
    top_ymin = det_ymin[top_indices]
    top_xmax = det_xmax[top_indices]
    top_ymax = det_ymax[top_indices]
    #print('total num boxes: ' + str(num_valid_boxes))
    colors = plt.cm.hsv(np.linspace(0, 1, 21)).tolist()
    infer_image = cv2.imread(IMAGE_FULL_PATH)

    plt.imshow(infer_image)
    currentAxis = plt.gca()

    for i in range(top_conf.shape[0]):
        xmin = int(round(top_xmin[i] * infer_image.shape[1]))
        ymin = int(round(top_ymin[i] * infer_image.shape[0]))
        xmax = int(round(top_xmax[i] * infer_image.shape[1]))
        ymax = int(round(top_ymax[i] * infer_image.shape[0]))
        score = top_conf[i]
        label = int(top_label_indices[i])
        label_name = "top_labels[i]"
        display_txt = '%s: %.2f'%(label_name, score)
        coords = (xmin, ymin), xmax-xmin+1, ymax-ymin+1
        color = colors[label]
        currentAxis.add_patch(plt.Rectangle(*coords, fill=False, edgecolor=color, linewidth=2))
        currentAxis.text(xmin, ymin, display_txt, bbox={'facecolor':color, 'alpha':0.5})
    plt.show()

def preprocess_image(src):

    # scale the image
    NETWORK_WIDTH = 300
    NETWORK_HEIGHT = 300
    img = cv2.resize(src, (NETWORK_WIDTH, NETWORK_HEIGHT))

    # adjust values to range between -1.0 and + 1.0
    img = img - 127.5
    img = img * 0.007843
    return img


# This function is called from the entry point to do
# all the work of the program
def main():
    # name of the opencv window
    cv_window_name = "SSD MobileNet - hit any key to exit"

    # Get a list of ALL the sticks that are plugged in
    # we need at least one
    devices = mvnc.EnumerateDevices()
    if len(devices) == 0:
        print('No devices found')
        quit()

    # Pick the first stick to run the network
    device = mvnc.Device(devices[0])

    # Open the NCS
    device.OpenDevice()

    # The graph file that was created with the ncsdk compiler
    graph_file_name = 'graph'

    # read in the graph file to memory buffer
    with open(graph_file_name, mode='rb') as f:
        graph_in_memory = f.read()

    # create the NCAPI graph instance from the memory buffer containing the graph file.
    graph = device.AllocateGraph(graph_in_memory)

    # read the image to run an inference on from the disk
    infer_image = cv2.imread(IMAGE_FULL_PATH)

    # run a single inference on the image and overwrite the
    # boxes and labels
    start=time.time()
    run_inference(infer_image, graph)
    stop=time.time()-start
    #print("Time taken to process an image is "+str(stop))

    # display the results and wait for user to hit a key
    #cv2.imshow(cv_window_name, infer_image)
    cv2.waitKey(0)

    # Clean up the graph and the device
    graph.DeallocateGraph()
    device.CloseDevice()


# main entry point for program. we'll call main() to do what needs to be done.
if __name__ == "__main__":
    sys.exit(main())