#!/usr/bin/env python
# _*_ coding: utf-8 _*_

# https://www.jianshu.com/p/d452b5615850

import cv2
import numpy as np
import os

# gets the absolute path to the file
abspath = os.path.abspath(os.path.dirname(__file__))

"""
    Non-max Suppression Algorithm

    @param list  Object candidate bounding boxes
    @param list  Confidence score of bounding boxes
    @param float IoU threshold

    @return Rest boxes after nms operation
"""
def nms(bounding_boxes, confidence_score, threshold):
    # If no bounding boxes, return empty list
    if len(bounding_boxes) == 0:
        return [], []

    # Bounding boxes
    boxes = np.array(bounding_boxes)

    # coordinates of bounding boxes
    start_x = boxes[:, 0]
    start_y = boxes[:, 1]
    end_x = boxes[:, 2]
    end_y = boxes[:, 3]

    # Confidence scores of bounding boxes
    score = np.array(confidence_score)

    # Picked bounding boxes
    picked_boxes = []
    picked_score = []

    # Compute areas of bounding boxes
    areas = (end_x - start_x + 1) * (end_y - start_y + 1)

    # Sort by confidence score of bounding boxes
    order = np.argsort(score) # 在score中的从小到大的序号序列

    # Iterate bounding boxes
    while order.size > 0: # 每次循环之后order同样是有序的
        # The index of largest confidence score
        index = order[-1]

        # Pick the bounding box with largest confidence score
        picked_boxes.append(bounding_boxes[index])
        picked_score.append(confidence_score[index])

        # Compute ordinates of intersection-over-union(IOU),计算重叠部分的坐标
        x1 = np.maximum(start_x[index], start_x[order[:-1]])
        x2 = np.minimum(end_x[index], end_x[order[:-1]])
        y1 = np.maximum(start_y[index], start_y[order[:-1]])
        y2 = np.minimum(end_y[index], end_y[order[:-1]])

        # Compute areas of intersection-over-union，计算重叠区域的长和宽，并取正数
        w = np.maximum(0.0, x2 - x1 + 1)
        h = np.maximum(0.0, y2 - y1 + 1)
        intersection = w * h

        # Compute the ratio between intersection and union，计算交并比
        ratio = intersection / (areas[index] + areas[order[:-1]] - intersection)

        left = np.where(ratio < threshold) # 返回交并比小于阈值的索引
        order = order[left]

    return picked_boxes, picked_score


def main():
    # Image name
    image_name = abspath + '/data/nms.jpg'

    # Bounding boxes
    bounding_boxes = [(187, 82, 337, 317), (150, 67, 305, 282), (246, 121, 368, 304)]
    confidence_score = [0.9, 0.75, 0.8]

    # Read image
    image = cv2.imread(image_name)

    # Copy image as original
    org = image.copy()

    # Draw parameters
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 1
    thickness = 2

    # IoU threshold
    # threshold = 0.4
    threshold = 0.5

    # Draw bounding boxes and confidence score
    for (start_x, start_y, end_x, end_y), confidence in zip(bounding_boxes, confidence_score):
        (w, h), baseline = cv2.getTextSize(str(confidence), font, font_scale, thickness)
        cv2.rectangle(org, (start_x, start_y - (2 * baseline + 5)), (start_x + w, start_y), (0, 255, 255), -1)
        cv2.rectangle(org, (start_x, start_y), (end_x, end_y), (0, 255, 255), 2)
        cv2.putText(org, str(confidence), (start_x, start_y), font, font_scale, (0, 0, 0), thickness)

    # Run non-max suppression algorithm
    picked_boxes, picked_score = nms(bounding_boxes, confidence_score, threshold)

    # Draw bounding boxes and confidence score after non-maximum supression
    for (start_x, start_y, end_x, end_y), confidence in zip(picked_boxes, picked_score):
        (w, h), baseline = cv2.getTextSize(str(confidence), font, font_scale, thickness)
        cv2.rectangle(image, (start_x, start_y - (2 * baseline + 5)), (start_x + w, start_y), (0, 255, 255), -1)
        cv2.rectangle(image, (start_x, start_y), (end_x, end_y), (0, 255, 255), 2)
        cv2.putText(image, str(confidence), (start_x, start_y), font, font_scale, (0, 0, 0), thickness)

    # Show image
    cv2.imshow('Original', org)
    cv2.imshow('NMS', image)
    cv2.waitKey(0)


if __name__ == "__main__":
    main()
    