from PIL import Image, ImageFont, ImageDraw
import numpy as np
from keras import backend as K
from keras.models import load_model
from keras.layers import Input
from yolo3.model import yolo_eval, yolo_body, tiny_yolo_body
from yolo3.utils import letterbox_image
import os
import tensorflow as tf
from timeit import default_timer as timer
from yolo3.utils import rgb2gray

def sigmoid(x):
    return 1./(1.+np.exp(-x))

sigmoid = np.vectorize(sigmoid)

#loading of model and printing of summary
num_classes = 2
anchors = np.array(
[np.array([55, 116]), np.array([67, 120]), np.array([69, 144]), np.array([85, 132]), np.array([107, 184]), np.array([203, 206])])
anchor_mask = [[3, 4, 5], [1, 2, 3]]
t_yolo = tiny_yolo_body(Input(shape=(256, 256, 3)), 6 // 2, num_classes)
t_yolo.load_weights("app4/trained/round2/trained_weights_final.h5")
t_yolo.summary()

#preparations of two images
image = Image.open("img1.png")
boxed_image = letterbox_image(image, tuple(reversed((416, 416))))
image_datas = rgb2gray(boxed_image, triplet=True)
image2 = Image.open("img2.png")
boxed_image2 = letterbox_image(image2, tuple(reversed((416, 416))))
image_datas2 = rgb2gray(boxed_image2, triplet=True)
image_data = np.array([image_datas, image_datas2])

#inference
features = t_yolo.predict(image_data)
length = len(features)
proto_box = []
proto_scores = []

#processing of results
for idx, val in enumerate(features):
    anchors = np.array([np.array([55,116]), np.array([67,120]), np.array([69,144]), np.array([85,132]), np.array([107,184]), np.array([203,216])])
    anchor_mask = [[3,4,5], [1,2,3]]
    input_shape = np.asarray(np.shape(features[0])[1 : 3]) * 32

    first = anchors[anchor_mask[idx]]
    image_size = (256, 256)

    num_anchors = len(first)
    anchors_tensor = np.reshape(first, [1, 1, 1, num_anchors, 2])
    grid_shape = np.shape(val)[1 : 3]
    b = np.reshape(np.arange(0, stop=grid_shape[0]), [-1, 1, 1, 1])
    grid_y = np.tile(np.reshape(np.arange(0, stop=grid_shape[0]), [-1, 1, 1, 1]), [1, grid_shape[1], 1, 1])
    grid_x = np.tile(np.reshape(np.arange(0, stop=grid_shape[1]), [1, -1, 1, 1]),
                    [grid_shape[0], 1, 1, 1])
    grid = np.concatenate([grid_x, grid_y], axis=3)

    feats = np.reshape(
            val, [-1, grid_shape[0], grid_shape[1], num_anchors, num_classes + 5])

    box_xy = (sigmoid(feats[..., :2]) + grid) / grid_shape[::-1]
    pre_box_wh = feats[..., 2:4] * anchors_tensor / input_shape[::-1]
    box_wh = np.exp(feats[..., 2:4]) * anchors_tensor / input_shape[::-1]
    box_confidence = sigmoid(feats[..., 4:5])
    box_class_probs = sigmoid(feats[..., 5:])

    box_yx = box_xy[..., ::-1]
    box_hw = box_wh[..., ::-1]
    image_shape = np.array([256, 256])
    new_shape = np.round((image_shape * np.min(input_shape/image_shape)))
    offset = (input_shape-new_shape)/2./input_shape
    scale = input_shape/new_shape
    box_yx = (box_yx - offset) * scale
    box_hw *= scale

    box_mins = box_yx - (box_hw / 2.)
    box_maxes = box_yx + (box_hw / 2.)

    boxes = np.concatenate([
            box_mins[..., 0:1],  # y_min
            box_mins[..., 1:2],  # x_min
            box_maxes[..., 0:1],  # y_max
            box_maxes[..., 1:2]  # x_max
        ], axis=4)

    # Scale boxes back to original image shape.
    scaler = np.concatenate([image_shape, image_shape])
    boxes *= scaler
    #here at original implementation is loosing of data, because batch size is ignored
    boxes = np.reshape(boxes, [boxes.shape[0], -1, 4])
    box_scores = box_confidence * box_class_probs
    box_scores = np.reshape(box_scores, [box_scores.shape[0], -1, num_classes])
    proto_box.append(boxes)
    proto_scores.append(box_scores)

proto_box = np.concatenate(proto_box, axis=1)
proto_scores = np.concatenate(proto_scores, axis=1)
mask = proto_scores >= 0.6
_boxes = []

#there is need for non maxima supression algorithm implementation
for idx, batch in enumerate(proto_scores):
    final_classes = []
    final_boxes = []
    final_scores = []
    for c in range(num_classes):
        class_boxes = proto_box[idx, mask[idx, :, c]]
        class_box_scores = proto_scores[idx, :, c][mask[idx, :, c]]
        classes = np.ones_like(class_box_scores, dtype="int32") * c
        final_boxes.append(class_boxes)
        final_scores.append(class_box_scores)
    final_boxes = np.concatenate(final_boxes, axis=0)
    final_scores = np.concatenate(final_scores, axis=0)
    _boxes.append(final_boxes)