手动实现目标检测非极大值抑制

0 187 90 586 607 0.94
0 120 116 600 370 0.68
1 511 185 961 418 0.58
0 340 145 568 478 0.47
1 524 70 920 565 0.92

代码如下：

import cv2
import numpy as np
from utilsW.utils import cvShow

def draw_box(image, bboxes_list):
    overlay_color = {
        '0': (0, 255, 0),
        '1': (255, 0, 0)
    }

    for coord in bboxes_list:
        #类别
        class_name = coord[0]
        start_point = (int(coord[1]), int(coord[2]))
        end_point = (int(coord[3]), int(coord[4]))
        prob = float(coord[5])
        text_start_point = (int(coord[1]), int(coord[2]) - 10)

        image = cv2.rectangle(image, start_point, end_point, overlay_color[class_name], 2)
        image = cv2.putText(image, str(prob), text_start_point, None, 0.8, overlay_color[class_name], 2)

    cv2.imshow("im", image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

def IOU(bboxes1, bboxes2):
    bboxes1 = [int(i) for i in bboxes1]
    bboxes2 = [int(i) for i in bboxes2]

    xA = max(bboxes1[0], bboxes2[0])
    yA = max(bboxes1[1], bboxes2[1])
    xB = min(bboxes1[2], bboxes2[2])
    yB = min(bboxes1[3], bboxes2[3])

    intersection_area = max(0, xB - xA + 1) * max(0, yB - yA + 1)

    box1_area = (bboxes1[2] - bboxes1[0] + 1) * (bboxes1[3] - bboxes1[1] + 1)
    box2_area = (bboxes2[2] - bboxes2[0] + 1) * (bboxes2[3] - bboxes2[1] + 1)

    iou = intersection_area / float(box1_area + box2_area - intersection_area)

    return iou


def NMS(src_copy, file, conf_threshold, iou_threshold):
    req_bboxes, final_boxes = [], []

    for coord in file:
        prob = float(coord[5])
        if prob > conf_threshold:
            req_bboxes.append(coord)

    # sorting the bounding boxes based on probability score
    bboxes_sorted = sorted(req_bboxes, reverse=True, key=lambda x: x[5])

    while len(bboxes_sorted) > 0:
        # removing the best probability bounding box
        box = bboxes_sorted.pop(0)

        for b in bboxes_sorted:
            # comparing with the same class
            if box[0] == b[0]:
                iou = IOU(box[1:-1], b[1:-1])
                if iou >= iou_threshold:
                    # if IOU is large then discard the box with lowest probability
                    bboxes_sorted.remove(b)
        print(len(bboxes_sorted))

        final_boxes.append(box)

    return final_boxes
if __name__ == '__main__':
    #1 设置初始参数
    conf_threshold = 0.5
    iou_threshold = 0.4

    #2 读取图片
    src = cv2.imread("zoraya.jpg")
    src_copy = src.copy()
    cvShow("src", src)

    #3 读取坐标框信息
    with open("coordinates.txt") as F:
        file = F.readlines()

    file = [i[:-1].split(' ') for i in file]

    #现在图片上画框，看一下效果
    draw_box(src, file)
    #非极大值抑制筛选点
    final_boxes = NMS(src_copy, file, conf_threshold, iou_threshold)
    draw_box(src_copy, final_boxes)