feat(util): nms实现

Author: zjZSTU

py/detector.py

@@ -69,8 +69,7 @@ def load_data(img_path, xml_path):
     device = "cpu"
 
     # img, data_dict = load_data('../imgs/cucumber_9.jpg', '../imgs/cucumber_9.xml')
-    # img, data_dict = load_data('../imgs/000012.jpg', '../imgs/000012.xml')
-    img, data_dict = load_data('../imgs/000007.jpg', '../imgs/000007.xml')
+    img, data_dict = load_data('../imgs/000012.jpg', '../imgs/000012.xml')
     model = file.load_model(device, S, B, C)
     # 计算
     outputs = model.forward(img.to(device)).cpu().squeeze(0)
@@ -98,6 +97,8 @@ def load_data(img_path, xml_path):
 
     # 预测边界框的缩放，回到原始图像
     pred_bboxs = util.deform_bboxs(pred_cate_bboxs, data_dict, S)
+
+    nms_rects, nms_scores, nms_cates = util.nms(pred_bboxs, pred_cate_probs, pred_cates)
     # 在原图绘制标注边界框和预测边界框
     dst = draw.plot_bboxs(data_dict['src'], data_dict['bndboxs'], data_dict['name_list'], cate_list,
                           pred_bboxs, pred_cates, pred_cate_probs)

py/lib/utils/util.py

@@ -162,12 +162,48 @@ def deform_bboxs(pred_bboxs, data_dict, S):
     return bboxs.astype(int)
 
 
-def nms(cates, probs, bboxs):
+def nms(rect_list, score_list, cate_list, thresh=0.3):
     """
-    non-maximum suppression
-    :param cates:
-    :param probs:
-    :param bboxs:
-    :return:
+    非最大抑制
+    :param rect_list: list，大小为[N, 4]
+    :param score_list： list，大小为[N]
+    :param cate_list: list, 大小为[N]
     """
-    pass
+    nms_rects = list()
+    nms_scores = list()
+    nms_cates = list()
+
+    rect_array = np.array(rect_list)
+    score_array = np.array(score_list)
+    cate_array = np.array(cate_list)
+
+    # 一次排序后即可
+    # 按分类概率从大到小排序
+    idxs = np.argsort(score_array)[::-1]
+    rect_array = rect_array[idxs]
+    score_array = score_array[idxs]
+    cate_array = cate_array[idxs]
+
+    while len(score_array) > 0:
+        # 添加分类概率最大的边界框
+        nms_rects.append(rect_array[0])
+        nms_scores.append(score_array[0])
+        nms_cates.append(cate_array[0])
+        rect_array = rect_array[1:]
+        score_array = score_array[1:]
+        cate_array = cate_array[1:]
+
+        length = len(score_array)
+        if length <= 0:
+            break
+
+        # 计算IoU
+        iou_scores = iou(np.array(nms_rects[len(nms_rects) - 1]), rect_array)
+        # print(iou_scores)
+        # 去除重叠率大于等于thresh的边界框
+        idxs = np.where(iou_scores < thresh)[0]
+        rect_array = rect_array[idxs]
+        score_array = score_array[idxs]
+        cate_array = cate_array[idxs]
+
+    return nms_rects, nms_scores, nms_cates