################################################################################################################################

import numpy      as np

import tensorflow as tf

################################################################################################################################

# This function is designed to work with Keras datagenerators as a model metric.

# SOURCE1: https://github.com/Balupurohit23/IOU-for-bounding-box-regression-in-Keras/blob/master/iou_metric.py

# SOURCE2: https://www.kaggle.com/vbookshelf/keras-iou-metric-implemented-without-tensor-drama

# I made use of the two sources because they each offered something I needed. I essentially used the framework of

# Source 2, but with the code of Source 1 because the format aligned with my data.

# I also reassigned the 'astype' numpy code to fit with tensorflow as well as converting the keras functions to np.

def tf_iou(y_true, y_pred):

    # iou as metric for bounding box regression

    # input must be as [x1, y1, x2, y2]

    results = []

    for i in range(0,y_true.shape[0]):

        y_true = tf.dtypes.cast(y_true, tf.float32)

        y_pred = tf.dtypes.cast(y_pred, tf.float32)

        # AOG = Area of Groundtruth box

        AoG = np.abs(np.transpose(y_true)[2] - np.transpose(y_true)[0] + 1) * np.abs(np.transpose(y_true)[3] -\

                                                                                     np.transpose(y_true)[1] + 1)

        # AOP = Area of Predicted box

        AoP = np.abs(np.transpose(y_pred)[2] - np.transpose(y_pred)[0] + 1) * np.abs(np.transpose(y_pred)[3] - \

                                                                                     np.transpose(y_pred)[1] + 1)

        # overlaps are the co-ordinates of intersection box

        overlap_0 = np.maximum(np.transpose(y_true)[0], np.transpose(y_pred)[0])

        overlap_1 = np.maximum(np.transpose(y_true)[1], np.transpose(y_pred)[1])

        overlap_2 = np.minimum(np.transpose(y_true)[2], np.transpose(y_pred)[2])

        overlap_3 = np.minimum(np.transpose(y_true)[3], np.transpose(y_pred)[3])

        # intersection area

        intersection = (overlap_2 - overlap_0 + 1) * (overlap_3 - overlap_1 + 1)

        # area of union of both boxes

        union = AoG + AoP - intersection

        # iou calculation

        iou = intersection / union

        # bounding values of iou to (0,1)

        epsilon = np.finfo(np.float32).eps

        iou     = np.clip(iou, 0.0 + epsilon, 1.0 - epsilon)

        iou     = tf.dtypes.cast(iou, tf.float32)

        # append the result to a list at the end of each loop

        results.append(iou)

    # return the mean IoU score for the batch

    return np.mean(results)

################################################################################################################################