2D Non-Maximum Suppression

nms_2d.py

# sid rajaram
import numpy as np
import time

# This approach assumes there are class scores in the incoming lists as well.
# Selects best score and then suppresses.
# class score + bounding box = (p, x, y, w, h)
# p: classification score / probability
# x,y: location
# w,h: dimensions


iou_threshold = 0.45

def iou(box_a, box_b):

    box_a_top_right_corner = [box_a[1]+box_a[3], box_a[2]+box_a[4]]
    box_b_top_right_corner = [box_b[1]+box_b[3], box_b[2]+box_b[4]]

    box_a_area = (box_a[3]) * (box_a[4])
    box_b_area = (box_b[3]) * (box_b[4])

    xi = max(box_a[1], box_b[1])
    yi = max(box_a[2], box_b[2])

    corner_x_i = min(box_a_top_right_corner[0], box_b_top_right_corner[0])
    corner_y_i = min(box_a_top_right_corner[1], box_b_top_right_corner[1])

    intersection_area = max(0, corner_x_i - xi) * max(0, corner_y_i - yi)

    iou = intersection_area / float(box_a_area + box_b_area - intersection_area + 1e-5)

    return iou


def nms(original_boxes):

    boxes_probability_sorted = original_boxes[np.flip(np.argsort(original_boxes[:, 0]))]
    box_indices = np.arange(0, len(boxes_probability_sorted))
    suppressed_box_indices = []
    tmp_suppress = []

    while len(box_indices) > 0:

        if box_indices[0] not in suppressed_box_indices:
            selected_box = box_indices[0]
            tmp_suppress = []

            for i in range(len(box_indices)):
                if box_indices[i] != selected_box:
                    selected_iou = iou(boxes_probability_sorted[selected_box], boxes_probability_sorted[box_indices[i]])
                    if selected_iou > iou_threshold:
                        suppressed_box_indices.append(box_indices[i])
                        tmp_suppress.append(i)

        box_indices = np.delete(box_indices, tmp_suppress, axis=0)
        box_indices = box_indices[1:]

    preserved_boxes = np.delete(boxes_probability_sorted, suppressed_box_indices, axis=0)
    return preserved_boxes, suppressed_box_indices


if __name__ == "__main__":

    # some random test bounding boxes
    box_0 = np.array([0.7, 1.0, 1.0, 2.0, 2.0])
    box_1 = np.array([0.92, 1, 0.5, 2.5, 2.0])
    box_2 = np.array([0.80, 1.5, 0.5, 2.5, 2.0])

    box_3 = np.array([0.89, 5, 2, 4.0, 2.0])
    box_4 = np.array([0.85, 5.5, 2, 4.0, 2.0])
    box_5 = np.array([0.92, 5, 1.5, 4.0, 2.0])

    box_6 = np.array([0.75, 1, 5, 2.0, 3.0])

    box_7 = np.array([0.80, 5, 6, 2.0, 4.0])
    box_8 = np.array([0.82, 5.5, 6, 2.5, 3.5])
    box_9 = np.array([0.97, 5.5, 6, 2.0, 3.5])

    boxes = np.array([box_0, box_1, box_2, box_3,
                      box_4, box_5, box_6,
                      box_7, box_8, box_9])

    print("{} Input Bounding Boxes (p,x,y,w,h):".format(len(boxes)))
    print(boxes)
    print()

    start = time.time()
    p, s = nms(boxes)
    end = time.time()

    print("{} seconds".format(end-start))
    print("{} Post-NMS Bounding Boxes (p,x,y,w,h):".format(len(p)))
    print(p)
    print()