import sys
sys.path.append("./keras-ssd")

from keras import backend as K
from keras.models import load_model
from keras.preprocessing import image
from keras.optimizers import Adam
from imageio import imread
import numpy as np
from matplotlib import pyplot as plt

from models.keras_ssd300 import ssd_300
from keras_loss_function.keras_ssd_loss import SSDLoss
from keras_layers.keras_layer_AnchorBoxes import AnchorBoxes
from keras_layers.keras_layer_DecodeDetections import DecodeDetections
from keras_layers.keras_layer_DecodeDetectionsFast import DecodeDetectionsFast
from keras_layers.keras_layer_L2Normalization import L2Normalization

from ssd_encoder_decoder.ssd_output_decoder import decode_detections, decode_detections_fast

from data_generator.object_detection_2d_data_generator import DataGenerator
from data_generator.object_detection_2d_photometric_ops import ConvertTo3Channels
from data_generator.object_detection_2d_geometric_ops import Resize
from data_generator.object_detection_2d_misc_utils import apply_inverse_transforms
from flask import Flask, request, jsonify, send_file, Response
from werkzeug.utils import secure_filename
import os
import io
import matplotlib.image as mpimg

import six.moves.urllib as urllib
import sys
import tarfile
import tensorflow as tf
import zipfile

from collections import defaultdict
from io import StringIO
from matplotlib import pyplot as plt
from PIL import Image
from vis.visualization import visualize_cam
import cv2

# the all-important app variable:
app = Flask(__name__)

def buildModel(img_height, img_width):
    tmpModel = ssd_300(image_size=(img_height, img_width, 3),
                    n_classes=20,
                    mode='inference',
                    l2_regularization=0.0005,
                    scales=[0.1, 0.2, 0.37, 0.54, 0.71, 0.88, 1.05],
                    # The scales for MS COCO are [0.07, 0.15, 0.33, 0.51, 0.69, 0.87, 1.05]
                    aspect_ratios_per_layer=[[1.0, 2.0, 0.5],
                                             [1.0, 2.0, 0.5, 3.0, 1.0 / 3.0],
                                             [1.0, 2.0, 0.5, 3.0, 1.0 / 3.0],
                                             [1.0, 2.0, 0.5, 3.0, 1.0 / 3.0],
                                             [1.0, 2.0, 0.5],
                                             [1.0, 2.0, 0.5]],
                    two_boxes_for_ar1=True,
                    steps=[8, 16, 32, 64, 100, 300],
                    offsets=[0.5, 0.5, 0.5, 0.5, 0.5, 0.5],
                    clip_boxes=False,
                    variances=[0.1, 0.1, 0.2, 0.2],
                    normalize_coords=True,
                    subtract_mean=[123, 117, 104],
                    swap_channels=[2, 1, 0],
                    confidence_thresh=0.5,
                    iou_threshold=0.45,
                    top_k=200,
                    nms_max_output_size=400)

    # 2: Load the trained weights into the model.

    # TODO: Set the path of the trained weights.
    weights_path = './keras-ssd/data/VGG_VOC0712Plus_SSD_300x300_iter_240000.h5'

    tmpModel.load_weights(weights_path, by_name=True)

    # 3: Compile the model so that Keras won't complain the next time you load it.

    adam = Adam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-08, decay=0.0)
    ssd_loss = SSDLoss(neg_pos_ratio=3, alpha=1.0)
    tmpModel.compile(optimizer=adam, loss=ssd_loss.compute_loss)
    return tmpModel


@app.route("/upload", methods=['POST'])
def uploadFiles():
    images = request.files.to_dict()
    for image in images:
        file = images[image]
        file_name = secure_filename(file.filename)
        file.save(os.path.join('upload', file_name))

    fig = detectObjects(file_name)
    img = io.BytesIO()
    fig.savefig(img, bbox_inches='tight', pad_inches = 0)
    img.seek(0)
    return send_file(img, mimetype='image/png')

def detectObjects(image_name):
    K.clear_session()
    # We'll only load one image in this example.
    img_path = os.path.join('upload', image_name)
    img_width, img_height = 300, 300
    print("Image size {} {}".format(img_height, img_width))
    model = buildModel(img_height, img_width)

    orig_images = []  # Store the images here.
    input_images = []  # Store resized versions of the images here.

    orig_images.append(imread(img_path))
    img = image.load_img(img_path, target_size=(img_height, img_width))
    img = image.img_to_array(img)
    input_images.append(img)
    input_images = np.array(input_images)

    y_pred = model.predict(input_images)

    confidence_threshold = 0.5
    y_pred_thresh = [y_pred[k][y_pred[k, :, 1] > confidence_threshold] for k in range(y_pred.shape[0])]
    np.set_printoptions(precision=2, suppress=True, linewidth=90)

    colors = plt.cm.hsv(np.linspace(0, 1, 21)).tolist()
    classes = ['background',
               'aeroplane', 'bicycle', 'bird', 'boat',
               'bottle', 'bus', 'car', 'cat',
               'chair', 'cow', 'diningtable', 'dog',
               'horse', 'motorbike', 'person', 'pottedplant',
               'sheep', 'sofa', 'train', 'tvmonitor']

    fig = plt.figure(figsize=(40, 20))
    # fig.subplots_adjust(left=0, right=2, top=1, bottom=0, wspace=0, hspace=0)
    plt.subplot(1, 2, 1)
    plt.imshow(orig_images[0])
    plt.axis("off")

    current_axis = plt.gca()

    for box in y_pred_thresh[0]:
        # Transform the predicted bounding boxes for the 300x300 image to the original image dimensions.
        xmin = box[2] * orig_images[0].shape[1] / img_width
        ymin = box[3] * orig_images[0].shape[0] / img_height
        xmax = box[4] * orig_images[0].shape[1] / img_width
        ymax = box[5] * orig_images[0].shape[0] / img_height
        color = colors[int(box[0])]
        label = '{}: {:.2f}'.format(classes[int(box[0])], box[1])
        current_axis.add_patch(
            plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, color=color, fill=False, linewidth=2))
        current_axis.text(xmin, ymin, label, size='x-large', color='white', bbox={'facecolor': color, 'alpha': 1.0})

    # show heat map
    # After try with each layer, we choose layer 23 to visualize CAM.
        heatmap = visualizeCam(model, 23, input_images[0])
        plt.subplot(1, 2, 2)
        res = cv2.resize(heatmap, dsize=(orig_images[0].shape[1], orig_images[0].shape[0]), interpolation=cv2.INTER_CUBIC)
        plt.imshow(res)
        plt.axis('off')
    return fig


def visualizeCam(model,layer, image):
    heat_map = visualize_cam(model, layer, None, image)
    return heat_map

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0')


