"""
## Importing the necessary libraries
"""
import numpy as np
import pandas as pd
from keras import Model
from keras.applications.mobilenet import MobileNet, preprocess_input
from keras.callbacks import ModelCheckpoint, EarlyStopping, ReduceLROnPlateau, Callback
from keras.layers import Conv2D, Reshape
from keras.utils import Sequence
from keras.backend import epsilon
import tensorflow as tf

from PIL import Image

import os
import random

import matplotlib.pyplot as plt
import matplotlib.patches as patches
from PIL import Image
import numpy as np
import cv2

np.random.seed(1)
"""
## Reading the training data from train.csv file
"""
train = pd.read_csv("../input/racoon-detection/train_labels_.csv")
train.head()
train.shape
row = train.sample()

img = cv2.imread(f'../input/racoon-detection/Racoon Images/images/{row.filename.values[0]}')
xmin = row.xmin.values[0]
ymin = row.ymin.values[0]
xmax = row.xmax.values[0]
ymax = row.ymax.values[0]

img = cv2.rectangle(img, (xmin, ymin), (xmax, ymax), color=(255, 0, 0), thickness=4)
plt.imshow(img)
plt.show()

"""
* Create a list variable known as 'path' which has all the path for all the training images
* Create an array 'coords' which has the resized coordinates of the bounding box for the training images
"""
IMAGE_SIZE = 128
coords = train[["width","height","xmin","ymin","xmax","ymax"]]

coords["xmin"] = coords["xmin"] * IMAGE_SIZE / coords["width"]
coords["xmax"] = coords["xmax"] * IMAGE_SIZE / coords["width"]
coords["ymin"] = coords["ymin"] * IMAGE_SIZE / coords["height"]
coords["ymax"] = coords["ymax"] * IMAGE_SIZE / coords["height"]

coords.drop(["width","height"], axis=1, inplace=True)
coords.head()
paths = train["filename"]
len(paths)
images = "../input/racoon-detection/Racoon Images/images/"

batch_images = np.zeros((len(paths), IMAGE_SIZE, IMAGE_SIZE,3), dtype=np.float32)

for i, f in enumerate(paths):
    img = Image.open(images + f)
    img = img.resize((IMAGE_SIZE, IMAGE_SIZE))
    img = img.convert('RGB')
    batch_images[i] = preprocess_input(np.array(img, dtype=np.float32))
"""
## Model Building
* Building the model using transfer learning
"""
model = MobileNet(input_shape=(IMAGE_SIZE, IMAGE_SIZE, 3), include_top=False)
for layers in model.layers:
    layers.trainable = False

x = model.layers[-1].output
x = Conv2D(4, kernel_size=4, name="coords")(x)
x = Reshape((4,))(x)

model = Model(inputs=model.inputs, outputs=x)
model.summary()
"""
## Define a custom loss function IoU which calculates Intersection Over Union
"""
def loss(gt, pred):
    intersections = 0
    unions = 0
    diff_width = np.minimum(gt[:,0] + gt[:,2], pred[:,0] + pred[:,2]) - np.maximum(gt[:,0], pred[:,0])
    diff_height = np.minimum(gt[:,1] + gt[:,3], pred[:,1] + pred[:,3]) - np.maximum(gt[:,1], pred[:,1])
    intersection = diff_width * diff_height
    
    # Compute union
    area_gt = gt[:,2] * gt[:,3]
    area_pred = pred[:,2] * pred[:,3]
    union = area_gt + area_pred - intersection

    # Compute intersection and union over multiple boxes
    for j, _ in enumerate(union):
        if union[j] > 0 and intersection[j] > 0 and union[j] >= intersection[j]:
            intersections += intersection[j]
            unions += union[j]

    # Compute IOU. Use 1e-8 to prevent division by zero
    iou = np.round(intersections / (unions + 1e-8), 4)
    iou = iou.astype(np.float32)
    return iou


def IoU(y_true, y_pred):
    iou = tf.py_function(loss, [y_true, y_pred], tf.float32)
    return iou
"""
## Compiling the model
"""
gt = coords

PATIENCE=10

model.compile(optimizer="Adam",
              loss="mse",
              metrics=[IoU])

stop = EarlyStopping(monitor='val_iou', patience=PATIENCE, mode="max")

reduce_lr = ReduceLROnPlateau(monitor='val_iou', factor=0.2, patience=PATIENCE,
                              min_lr=1e-7, verbose=1, mode="max")

model.fit(batch_images, gt,
          epochs=100,
          callbacks=[stop, reduce_lr],
          verbose=1)
"""
## Pick a test image from the given data
"""
test_img = random.choice(paths)
filename = images + test_img
unscaled = cv2.imread(filename)
unscaled = cv2.cvtColor(unscaled, cv2.COLOR_BGR2RGB)
"""
## Preprocessing of Test Image
Resizing the image to 128 * 128 and preprocess the image for the MobileNet model
"""
image_height, image_width, _ = unscaled.shape
image = cv2.resize(unscaled, (IMAGE_SIZE, IMAGE_SIZE))
feat_scaled = preprocess_input(np.array(image, dtype=np.float32))
"""
* Predict the coordinates of the bounding box for the given test image
"""
region = model.predict(x=np.array([feat_scaled]))[0]
"""
* Scaling the BBox
"""
x0 = int(region[0] * image_width / IMAGE_SIZE) 
y0 = int(region[1] * image_height / IMAGE_SIZE)

x1 = int((region[2]) * image_width / IMAGE_SIZE)
y1 = int((region[3]) * image_height / IMAGE_SIZE)
"""
## Plotting the predicted bounding box
"""
unscaled = cv2.rectangle(unscaled, (x0, y0), (x1, y1), color=(255, 0, 0), thickness=2)

# Display the image
plt.imshow(unscaled)
plt.show()