目标检测nms python代码_目标检测中IOU和NMS的python实现

IOU：两个框的交并比

import numpy as np

def compute_iou(box1, box2, wh=False):

"""

compute the iou of two boxes.

Args:

box1, box2: [xmin, ymin, xmax, ymax] (wh=False) or [xcenter, ycenter, w, h] (wh=True)

wh: the format of coordinate.

Return:

iou: iou of box1 and box2.

"""

if wh == False:

xmin1, ymin1, xmax1, ymax1 = box1

xmin2, ymin2, xmax2, ymax2 = box2

else:

xmin1, ymin1 = int(box1[0]-box1[2]/2.0), int(box1[1]-box1[3]/2.0)

xmax1, ymax1 = int(box1[0]+box1[2]/2.0), int(box1[1]+box1[3]/2.0)

xmin2, ymin2 = int(box2[0]-box2[2]/2.0), int(box2[1]-box2[3]/2.0)

xmax2, ymax2 = int(box2[0]+box2[2]/2.0), int(box2[1]+box2[3]/2.0)

## 获取矩形框交集对应的左上角和右下角的坐标(intersection)

xx1 = np.max([xmin1, xmin2])

yy1 = np.max([ymin1, ymin2])

xx2 = np.min([xmax1, xmax2])

yy2 = np.min([ymax1, ymax2])

## 计算两个矩形框面积

area1 = (xmax1-xmin1) * (ymax1-ymin1)

area2 = (xmax2-xmin2) * (ymax2-ymin2)

inter_area = (np.max([0, xx2-xx1])) * (np.max([0, yy2-yy1]))　#计算交集面积

iou = inter_area / (area1+area2-inter_area+1e-6)　＃计算交并比

return iou

NMS：非极大值抑制

for object in all objects:

(1) 获取当前目标类别下所有bbx的信息

(2) 将bbx按照confidence从高到低排序,并记录当前confidence最大的bbx

(3) 计算最大confidence对应的bbx与剩下所有的bbx的IOU,移除所有大于IOU阈值的bbx

(4) 对剩下的bbx，循环执行(2)和(3)直到所有的bbx均满足要求(即不能再移除bbx)

NMS是对所有的类别分别执行的。举个栗子，假设最后预测出的矩形框有2类(分别为cup, pen),在NMS之前，每个类别可能都会有不只一个bbx被预测出来，这个时候我们需要对这两个类别分别执行一次NMS过程。

import numpy as np

class Bounding_box:

def __init__(self, x1, y1, x2, y2, score):

self.x1 = x1

self.y1 = y1

self.x2 = x2

self.y2 = y2

self.score = score

def get_iou(boxa, boxb):

max_x = max(boxa.x1, boxb.x1)

max_y = max(boxa.y1, boxb.y1)

min_x = min(boxa.x2, boxb.x2)

min_y = min(boxa.y2, boxb.y2)

if min_x <= max_x or min_y <= max_y:

return 0

area_i = (min_x - max_x) * (min_y - max_y)

area_a = (boxa.x2 - boxa.x1) * (boxa.y2 - boxa.y1)

area_b = (boxb.x2 - boxb.x1) * (boxb.y2 - boxb.y1)

area_u = area_a + area_b - area_i

return float(area_i) / float(area_u)

def NMS(box_lists, k):

box_lists = sorted(box_lists, key=lambda x: x.score, reverse=True)

NMS_lists = [box_lists[0]]

temp_lists = []

for i in range(k):

for j in range(1, len(box_lists)):

iou = get_iou(NMS_lists[i], box_lists[j])

if iou < 0.7:

temp_lists.append(box_lists[j])

if len(temp_lists) == 0:

return NMS_lists

box_lists = temp_lists

temp_lists = []

NMS_lists.append(box_lists[0])

return NMS_lists

box1 = Bounding_box(13, 22, 268, 367, 0.124648176)

box2 = Bounding_box(18, 27, 294, 400, 0.35818103)

box3 = Bounding_box(234, 123, 466, 678, 0.13638769)

box_lists = [box1, box2, box3]

NMS_list = NMS(box_lists, 2)

print NMS_list

print NMS_list[0].x1