不规则四边形回归框计算IOU与NMS

前言

在复现DPR论文过程中，遇到了计算IOU和NMS的问题。传统目标检测大多基于anchor，给出的预测也是矩形框，所以需要进行相应的修改，以套用到不规则四边形回归框上。

注1：编者水平有限，如有谬误，欢迎指正。若要转载，请注明出处，谢谢。
注2：实现参考了TextBoxes++部分源码和SSD pytorch部分源码，并非完全原创。
联系方式：
邮箱：[email protected]
QQ：1156356625

---

• 计算IOU

import shapely
from shapely.geometry import Polygon,MultiPoint  #多边形

def get_iou(a,b):
	'''
	:param a: box a [x0,y0,x1,y1,x2,y2,x3,y3]
	:param b: box b [x0,y0,x1,y1,x2,y2,x3,y3]
	:return: iou of bbox a and bbox b
	'''
	a = a.reshape(4, 2)
	poly1 = Polygon(a).convex_hull
	
	b = b.reshape(4, 2)
	poly2 = Polygon(b).convex_hull
	
	if not poly1.intersects(poly2):  # 如果两四边形不相交
		iou = 0
	else:
		try:
			inter_area = poly1.intersection(poly2).area  # 相交面积
			union_area = poly1.area + poly2.area - inter_area
			if union_area == 0:
				iou = 0
			else:
				iou = inter_area / union_area
		except shapely.geos.TopologicalError:
			print('shapely.geos.TopologicalError occured, iou set to 0')
			iou = 0
	return iou

这里基本用的是TextBoxes++的实现源码，去掉了一些冗余。

• 计算NMS

def NMS_quadrangle(bboxes, overlap=0.5, top_k=200):
	'''
	An implement of quadrangle NMS
	:param bboxes: top k bboxes(sorted with confidence)
	:param overlap: IOU threshold
	:return:bboxes after NMS
	'''
	# bboxes = bboxes[:-1].transpose(0,1)
    top_k = min(bboxes.size()[0], top_k)
	keep = torch.zeros(top_k)
	count = 0
	idx = torch.Tensor(range(top_k)[::-1]).int()
	if not bboxes.numel():
		return bboxes
	while idx.numel() > 0:
		a = idx[-1]
		keep[count] = a
		count += 1
		if not idx.numel():
			break
			
		idx = idx[:-1]
		iou = torch.Tensor([get_iou(bboxes[a],bboxes[b]) for b in idx])
		
		idx = idx[iou.le(overlap)]
	bboxes =bboxes.index_select(dim=0,index=keep[:count].long())
	return torch.where(bboxes < 0,torch.full_like(bboxes,0),bboxes)

基于anchor的NMS比较好写成tensor操作，但是不规则四边形，暂时没有想到好的处理方式。这部分代码只在inference过程测试过，速度要求不高，如果需要并入训练过程可能需要考虑计算iou部分如何做成矩阵运算。