from tflite_runtime.interpreter import Interpreter
import tensorflow as tf
import tflite_runtime.interpreter as tflite
import os
import argparse
import cv2
import numpy as np
import sys
import glob
import matplotlib.pyplot as plt

# Define and parse input arguments
parser = argparse.ArgumentParser()
parser.add_argument(
    "--modeldir", help="Folder the .tflite file is located in", required=True
)
parser.add_argument(
    "--graph",
    help="Name of the .tflite file, if different than detect.tflite",
    default="detect.tflite",
)
parser.add_argument(
    "--labels",
    help="Name of the labelmap file, if different than labelmap.txt",
    default="labelmap.txt",
)
parser.add_argument(
    "--threshold",
    help="Minimum confidence threshold for displaying detected objects",
    default=0.4,
)
parser.add_argument(
    "--image",
    help="Name of the single image to perform detection on. To run detection on multiple images, use --imagedir",
    default=None,
)
parser.add_argument(
    "--imagedir",
    help="Name of the folder containing images to perform detection on. Folder must contain only images.",
    default=None,
)
parser.add_argument(
    "--savedir",
    help="Name of the folder to save results",
    default=None,
)
parser.add_argument(
    "--edgetpu",
    help="Use Coral Edge TPU Accelerator to speed up detection",
    action="store_true",
)

args = parser.parse_args()

MODEL_NAME = args.modeldir
GRAPH_NAME = args.graph
LABELMAP_NAME = args.labels
min_conf_threshold = float(args.threshold)
use_TPU = args.edgetpu
savedir = args.savedir

# Parse input image name and directory.
IM_NAME = args.image
IM_DIR = args.imagedir

# If both an image AND a folder are specified, throw an error
if IM_NAME and IM_DIR:
    print(
        'Error! Please only use the --image argument or the --imagedir argument, not both. Issue "python TFLite_detection_image.py -h" for help.'
    )
    sys.exit()

# If neither an image or a folder are specified, default to using 'test1.jpg' for image name
if not IM_NAME and not IM_DIR:
    IM_NAME = "test1.jpg"

if IM_DIR and not savedir:
    print(
        'Error! Please specify the folder to save detection images in. Issue "python TFLite_detection_image.py -h" for help.'
    )
    sys.exit()

if savedir:
    if not os.path.isdir(savedir):
        os.makedirs(savedir)

CWD_PATH = os.getcwd()

# Define path to images and grab all image filenames
if IM_DIR:
    PATH_TO_IMAGES = os.path.join(CWD_PATH, IM_DIR)
    images = glob.glob(PATH_TO_IMAGES + "/*.jpg")

elif IM_NAME:
    PATH_TO_IMAGES = os.path.join(CWD_PATH, IM_NAME)
    images = glob.glob(PATH_TO_IMAGES)

# Path to .tflite file, which contains the model that is used for object detection
PATH_TO_CKPT = os.path.join(CWD_PATH, MODEL_NAME)

# Path to label map file
PATH_TO_LABELS = os.path.join(CWD_PATH, LABELMAP_NAME)

# Load the label map
with open(PATH_TO_LABELS, "r") as f:
    labels = [line.strip() for line in f.readlines()]

# Have to do a weird fix for label map if using the COCO "starter model" from
# https://www.tensorflow.org/lite/models/object_detection/overview
# First label is '???', which has to be removed.
if labels[0] == "???":
    del labels[0]

# Load the Tensorflow Lite model.
interpreter = Interpreter(model_path=PATH_TO_CKPT)

interpreter.allocate_tensors()

# Get model details
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
height = input_details[0]["shape"][1]
width = input_details[0]["shape"][2]

floating_model = input_details[0]["dtype"] == np.float32

input_mean = 127.5
input_std = 127.5


for image_path in images:
    image_name = image_path.split('\\')[-1]

    # Load image and resize to expected shape [1xHxWx3]
    image = cv2.imread(image_path)
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    imH, imW, _ = image.shape
    image_resized = cv2.resize(image_rgb, (width, height))
    input_data = np.expand_dims(image_resized, axis=0)

    # Normalize pixel values if using a floating model (i.e. if model is non-quantized)
    if floating_model:
        input_data = (np.float32(input_data) - input_mean) / input_std

    # Perform the actual detection by running the model with the image as input
    interpreter.tensor(interpreter.get_input_details()[0]["index"])
    interpreter.set_tensor(input_details[0]["index"], input_data)
    interpreter.invoke()

    # Retrieve detection results
    boxes = interpreter.get_tensor(output_details[0]["index"])[0]
    # Bounding box coordinates of detected objects
    classes = interpreter.get_tensor(output_details[1]["index"])[0]
    # Class index of detected objects
    scores = interpreter.get_tensor(output_details[2]["index"])[0]
    # Confidence of detected objects
    # Total number of detected objects (inaccurate and not needed)
    num = interpreter.get_tensor(output_details[3]['index'])[0]

    # Loop over all detections and draw detection box if confidence is above minimum threshold
    for i in range(len(scores)):
        if (scores[i] > min_conf_threshold) and (scores[i] <= 1.0):

            # Get bounding box coordinates and draw box
            # Interpreter can return coordinates that are outside of image dimensions, need to force them to be within image using max() and min()
            ymin = int(max(1, (boxes[i][0] * imH)))
            xmin = int(max(1, (boxes[i][1] * imW)))
            ymax = int(min(imH, (boxes[i][2] * imH)))
            xmax = int(min(imW, (boxes[i][3] * imW)))

            # cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (10, 255, 0), 2)
            cropped_image = tf.image.crop_to_bounding_box(image, ymin, xmin,
                                                          ymax - ymin, xmax - xmin).numpy()

            # cv2.imshow("cropped_image Window", cropped_image)
            cv2.imwrite(os.path.join(CWD_PATH, savedir,
                                     image_name), cropped_image)
