手写NMS和魔改（Pytorch版本）

NMS方法的总结可以参考我之前的文章：
https://blog.csdn.net/qq_34919792/article/details/108186234

非极大值抑制（Non-Maximum Suppression，NMS），顾名思义就是抑制不是极大值的元素。在检测中，我们通过将IOU大于一定阈值的框做一个筛选，只保留置信度最高的框。

网上比较经典的实现思路

def py_cpu_nms(dets, thresh): 
	"""Pure Python NMS baseline.""" 
	#x1、y1、x2、y2、以及score赋值 
	x1 = dets[:, 0] 
	y1 = dets[:, 1] 
	x2 = dets[:, 2] 
	y2 = dets[:, 3] 
	scores = dets[:, 4] 

	#每一个检测框的面积 
	areas = (x2 - x1 + 1) * (y2 - y1 + 1) 
	#按照score置信度降序排序 
	order = scores.argsort()[::-1] 

	keep = [] #保留的结果框集合
	while order.size > 0: 
		i = order[0] 
		keep.append(i) #保留该类剩余box中得分最高的一个 
		#得到相交区域,左上及右下 
		xx1 = np.maximum(x1[i], x1[order[1:]]) 
		yy1 = np.maximum(y1[i], y1[order[1:]]) 
		xx2 = np.minimum(x2[i], x2[order[1:]]) 
		yy2 = np.minimum(y2[i], y2[order[1:]]) 

		#计算相交的面积,不重叠时面积为0 
		w = np.maximum(0.0, xx2 - xx1 + 1) 
		h = np.maximum(0.0, yy2 - yy1 + 1) 
		inter = w * h 
		#计算IoU：重叠面积 /（面积1+面积2-重叠面积） 
		ovr = inter / (areas[i] + areas[order[1:]] - inter) 
		#保留IoU小于阈值的box 
		inds = np.where(ovr <= thresh)[0] 
		order = order[inds + 1] #因为ovr数组的长度比order数组少一个,所以这里要将所有下标后移一位 
	
	return keep

其实在NMS可以做pytorch版本，速度和cpython加速后一样，即和pytorch自带的库一样，但是更适合改。

def IOU(box_a, box_b):
    inter = intersect(box_a, box_b)
    area_a = ((box_a[:, 2]-box_a[:, 0]) *
              (box_a[:, 3]-box_a[:, 1])).unsqueeze(1).expand_as(inter)  # [A,B]
    area_b = ((box_b[:, 2]-box_b[:, 0]) *
              (box_b[:, 3]-box_b[:, 1])).unsqueeze(0).expand_as(inter)  # [A,B]
    union = area_a + area_b - inter
    return inter / union  # [A,B]


def torch_nms(boxes, scores, iou_threshold):
    _, idx = scores.sort(0, descending=True) # descending表示降序
    boxes_idx = boxes[idx]
    iou = IOU(boxes_idx, boxes_idx).triu_(diagonal=1) #取上三角矩阵，不包含对角线
    B = iou
    while 1:
        A = B
        maxA, _ = torch.max(A, dim=0)
        E = (maxA <= iou_threshold).float().unsqueeze(1).expand_as(A)
        B = iou.mul(E)
        if A.equal(B) == True:
            break
    keep= idx[maxA <= iou_threshold]
    return keep

这样的代码其实很好改进，比如可以把IOU去改成GIOU或者DIOU或者自己魔改魔改，也可以去改进NMS，提供个DIOU的改进示例。

def DIOU(box_a, box_b, delta = 0.9):
    inter = intersect(box_a, box_b)
    area_a = ((box_a[:, :, 2]-box_a[:, :, 0]) *
              (box_a[:, :, 3]-box_a[:, :, 1])).unsqueeze(2).expand_as(inter)  # [A,B]
    area_b = ((box_b[:, :, 2]-box_b[:, :, 0]) *
              (box_b[:, :, 3]-box_b[:, :, 1])).unsqueeze(1).expand_as(inter)  # [A,B]
    union = area_a + area_b - inter
    x1 = ((box_a[:, :, 2]+box_a[:, :, 0]) / 2).unsqueeze(2).expand_as(inter)
    y1 = ((box_a[:, :, 3]+box_a[:, :, 1]) / 2).unsqueeze(2).expand_as(inter)
    x2 = ((box_b[:, :, 2]+box_b[:, :, 0]) / 2).unsqueeze(1).expand_as(inter)
    y2 = ((box_b[:, :, 3]+box_b[:, :, 1]) / 2).unsqueeze(1).expand_as(inter)

    t1 = box_a[:, :, 1].unsqueeze(2).expand_as(inter)
    b1 = box_a[:, :, 3].unsqueeze(2).expand_as(inter)
    l1 = box_a[:, :, 0].unsqueeze(2).expand_as(inter)
    r1 = box_a[:, :, 2].unsqueeze(2).expand_as(inter)

    t2 = box_b[:, :, 1].unsqueeze(1).expand_as(inter)
    b2 = box_b[:, :, 3].unsqueeze(1).expand_as(inter)
    l2 = box_b[:, :, 0].unsqueeze(1).expand_as(inter)
    r2 = box_b[:, :, 2].unsqueeze(1).expand_as(inter)
    cr = torch.max(r1, r2)
    cl = torch.min(l1, l2)
    ct = torch.min(t1, t2)
    cb = torch.max(b1, b2)
    D = (((x2 - x1)**2 + (y2 - y1)**2) / ((cr-cl)**2 + (cb-ct)**2 + 1e-7))
    out = inter / union - D ** delta
    return out if use_batch else out.squeeze(0)


def torch_nms(boxes, scores, iou_threshold):
    _, idx = scores.sort(0, descending=True) # descending表示降序
    boxes_idx = boxes[idx]
    iou = DIOU(boxes_idx, boxes_idx).triu_(diagonal=1) #取上三角矩阵，不包含对角线
    B = iou
    while 1:
        A = B
        maxA, _ = torch.max(A, dim=0)
        E = (maxA <= iou_threshold).float().unsqueeze(1).expand_as(A)
        B = iou.mul(E)
        if A.equal(B) == True:
            break
    keep= idx[maxA <= iou_threshold]
    return keep