# Upload the required wheel files, models and images in a google drive folder
# Uncomment and run the below command to copy them in your current workspace
#!cp /content/drive/MyDrive/TF-OV/working_dir_files/* . 

import os

# Enable these variables for runtime inference optimizations
os.environ["OPENVINO_TF_CONVERT_VARIABLES_TO_CONSTANTS"] = "1"
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "1"

!python3 -m pip -q install --upgrade pip
!python3 -m pip -q install pillow
!python3 -m pip -q install keras_applications

# Install TensorFlow (v2.8.0) and OpenVINO-TensorFlow (v2.0.0) only if they aren't found
!if python3 -c "import tensorflow"; then echo "Found TensorFlow. Skipping."; else echo "TensorFlow Not Found. Installing."; python3 -m pip -q install tensorflow==2.8.0; fi
!if python3 -c "import openvino_tensorflow"; then echo "Found OpenVINO-TensorFlow. Skipping."; else echo "OpenVINO-TensorFlow Not Found. Installing."; 

OVTF_DIR = "/"
RAW_GITHUB_COMMON = "https://raw.githubusercontent.com/openvinotoolkit/openvino_tensorflow/master/examples/common/"
GITHUB_EXAMPLES = "https://github.com/openvinotoolkit/openvino_tensorflow/raw/master/examples/data/"
RAW_GITHUB_EXAMPLES = "https://raw.githubusercontent.com/openvinotoolkit/openvino_tensorflow/master/examples/"

import os

files = os.listdir('.')
if ('common' not in files or 'examples' not in files) and 'openvino_tensorflow' not in os.listdir(OVTF_DIR):
    !mkdir ./common
    !wget {RAW_GITHUB_COMMON}/post_process.py -O ./common/post_process.py
    !wget {RAW_GITHUB_COMMON}/pre_process.py -O ./common/pre_process.py
    !wget {RAW_GITHUB_COMMON}/utils.py -O ./common/utils.py
    !mkdir -p ./examples/data
    !wget {GITHUB_EXAMPLES}/grace_hopper.jpg -O ./examples/data/grace_hopper.jpg
    !wget {GITHUB_EXAMPLES}/yolov4_anchors.txt -O ./examples/data/yolov4_anchors.txt
    !wget {RAW_GITHUB_EXAMPLES}/convert_yolov4.sh -O ./examples/convert_yolov4.sh
    !wget {RAW_GITHUB_EXAMPLES}/keras_to_tensorflow.patch -O ./examples/keras_to_tensorflow.patch

import sys

if 'openvino_tensorflow' in os.listdir(OVTF_DIR):
    sys_append = os.path.abspath(OVTF_DIR + "/openvino_tensorflow/examples/")
    sys.path.append(sys_append)

from __future__ import absolute_import, division, print_function, unicode_literals

import numpy as np
import tensorflow as tf
import openvino_tensorflow as ovtf
from PIL import Image
import cv2
import matplotlib.pyplot as plt

from common.utils import get_input_mode, get_colors, draw_boxes, get_anchors, rename_file
from common.pre_process import preprocess_image_yolov3 as preprocess_image
from common.post_process import yolo3_postprocess_np

# Download and Convert the YoloV4 model
files = os.listdir('.')
if 'examples' in files:
    path = "examples"
else:
    path = "{0}/openvino_tensorflow/examples/".format(OVTF_DIR)
%cd {path}
!chmod +x convert_yolov4.sh && bash convert_yolov4.sh
# Once the model conversion is completed; move back to outside of examples directory
%cd ../
  
def load_coco_names(file_name):
    """Parses the label file with only class names,
      and returns a dictionary mapping the class IDs to class names.
    """
    names = {}
    with open(file_name) as f:
        for id_, name in enumerate(f):
            names[id_] = name
    return names

def load_labels(label_file):
    """Parses the label file, assuming that labels are separated with a newline
       in the file and returns the list of labels.
    """  
    label = []
    proto_as_ascii_lines = tf.io.gfile.GFile(label_file).readlines()
    for l in proto_as_ascii_lines:
        label.append(l.rstrip())
    return label

def infer_openvino_tensorflow(model_file, image_file , input_height, input_width, label_file, anchor_file, conf_threshold, iou_threshold):
    """Takes the tensorflow model and all other input parameters as arguments. 
       Runs inference with the object detection model and prints the predictions.
    """
    print("CREATE MODEL - BEGIN")

    # Load model and process input image
    model =     model = tf.saved_model.load(model_file)
    print("CREATE MODEL - END")

    if label_file:
        classes = load_coco_names(label_file)
        labels = load_labels(label_file)
        colors = get_colors(labels)

    if anchor_file:
        anchors = get_anchors(anchor_file)

    print("PREDICTION - BEGIN")
    
    #Preprocess Image
    image = Image.open(image_file)
    img = np.asarray(image)
    image_width, image_height = image.size
    img_resized = tf.convert_to_tensor(preprocess_image(image, (input_height, input_width)))

    # Warmup
    detected_boxes = model(img_resized)
    # Run
    import time
    start = time.time()
    detected_boxes = model(img_resized)
    elapsed = time.time() - start
    print('Inference time in ms: %f' % (elapsed * 1000))
    print("PREDICTION - END")  
    
    image_shape = tuple((image_height, image_width))  
    # apply non max suppresion, draw boxes and save updated image
    out_boxes, out_classes, out_scores = yolo3_postprocess_np(
            detected_boxes,
            image_shape,
            anchors,
            len(labels), (input_height, input_width),
            max_boxes=10,
            confidence=conf_threshold,
            iou_threshold=iou_threshold,
            elim_grid_sense=True)
    img_bbox = draw_boxes(img, out_boxes, out_classes, out_scores,
                        labels, colors)
    if output_dir:
        cv2.imwrite(os.path.join(output_dir, "detections.jpg"), img_bbox)
    else:
        cv2.imwrite("detections.jpg", img_bbox)
    plt.imshow(img)

input_file = "examples/data/grace_hopper.jpg"
model_file = "examples/data/yolo_v4"
label_file = "examples/data/coco.names"
anchor_file = "examples/data/yolov4_anchors.txt"
input_height = 416
input_width = 416
backend_name = "CPU"
output_dir = "."
conf_threshold = 0.6
iou_threshold = 0.5

#Print list of available backends
print('Available Backends:')
backends_list = ovtf.list_backends()
for backend in backends_list:
    print(backend)
ovtf.set_backend(backend_name)

print("OpenVINO TensorFlow is enabled")
infer_openvino_tensorflow(model_file, input_file, input_height, input_width, label_file, anchor_file, conf_threshold, iou_threshold )

ovtf.disable() ## Disabling OVTF
print("OpenVINO TensorFlow is disabled")
infer_openvino_tensorflow(model_file, input_file, input_height, input_width, label_file, anchor_file, conf_threshold, iou_threshold )
ovtf.enable()