From 3aad4947f20dc020d61d181e75a4138910bb6930 Mon Sep 17 00:00:00 2001
From: qingqing01 <dangqingqing@baidu.com>
Date: Wed, 3 Jun 2020 21:03:17 +0800
Subject: [PATCH] Adapte softnms for YOLOv3 (#854) (#866)

* Adapte softnms for YOLOv3
---
 ppdet/modeling/anchor_heads/yolo_head.py |  4 ++-
 ppdet/modeling/ops.py                    | 39 ++++++++++++++++++------
 2 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/ppdet/modeling/anchor_heads/yolo_head.py b/ppdet/modeling/anchor_heads/yolo_head.py
index a998f789d..e8dac83bd 100644
--- a/ppdet/modeling/anchor_heads/yolo_head.py
+++ b/ppdet/modeling/anchor_heads/yolo_head.py
@@ -20,7 +20,7 @@ from paddle import fluid
 from paddle.fluid.param_attr import ParamAttr
 from paddle.fluid.regularizer import L2Decay
 
-from ppdet.modeling.ops import MultiClassNMS
+from ppdet.modeling.ops import MultiClassNMS, MultiClassSoftNMS
 from ppdet.modeling.losses.yolo_loss import YOLOv3Loss
 from ppdet.core.workspace import register
 from ppdet.modeling.ops import DropBlock
@@ -335,6 +335,8 @@ class YOLOv3Head(object):
 
         yolo_boxes = fluid.layers.concat(boxes, axis=1)
         yolo_scores = fluid.layers.concat(scores, axis=2)
+        if type(self.nms) is MultiClassSoftNMS:
+            yolo_scores = fluid.layers.transpose(yolo_scores, perm=[0, 2, 1])
         pred = self.nms(bboxes=yolo_boxes, scores=yolo_scores)
         return {'bbox': pred}
 
diff --git a/ppdet/modeling/ops.py b/ppdet/modeling/ops.py
index 4065e7d36..972fcdf16 100644
--- a/ppdet/modeling/ops.py
+++ b/ppdet/modeling/ops.py
@@ -547,8 +547,6 @@ class MultiClassSoftNMS(object):
             return dets_final
 
         def _soft_nms(bboxes, scores):
-            bboxes = np.array(bboxes)
-            scores = np.array(scores)
             class_nums = scores.shape[-1]
 
             softnms_thres = self.score_threshold
@@ -562,7 +560,8 @@ class MultiClassSoftNMS(object):
             for j in range(start_idx, class_nums):
                 inds = np.where(scores[:, j] >= softnms_thres)[0]
                 scores_j = scores[inds, j]
-                rois_j = bboxes[inds, j, :]
+                rois_j = bboxes[inds, j, :] if len(
+                    bboxes.shape) > 2 else bboxes[inds, :]
                 dets_j = np.hstack((scores_j[:, np.newaxis], rois_j)).astype(
                     np.float32, copy=False)
                 cls_rank = np.argsort(-dets_j[:, 0])
@@ -584,12 +583,34 @@ class MultiClassSoftNMS(object):
                 keep = np.where(cls_boxes[:, 0] >= image_thresh)[0]
                 pred_result = pred_result[keep, :]
 
-            res = fluid.LoDTensor()
-            res.set_lod([[0, pred_result.shape[0]]])
-            if pred_result.shape[0] == 0:
-                pred_result = np.array([[1]], dtype=np.float32)
-            res.set(pred_result, fluid.CPUPlace())
+            return pred_result
+
+        def _batch_softnms(bboxes, scores):
+            batch_offsets = bboxes.lod()
+            bboxes = np.array(bboxes)
+            scores = np.array(scores)
+            out_offsets = [0]
+            pred_res = []
+            if len(batch_offsets) > 0:
+                batch_offset = batch_offsets[0]
+                for i in range(len(batch_offset) - 1):
+                    s, e = batch_offset[i], batch_offset[i + 1]
+                    pred = _soft_nms(bboxes[s:e], scores[s:e])
+                    out_offsets.append(pred.shape[0] + out_offsets[-1])
+                    pred_res.append(pred)
+            else:
+                assert len(bboxes.shape) == 3
+                assert len(scores.shape) == 3
+                for i in range(bboxes.shape[0]):
+                    pred = _soft_nms(bboxes[i], scores[i])
+                    out_offsets.append(pred.shape[0] + out_offsets[-1])
+                    pred_res.append(pred)
 
+            res = fluid.LoDTensor()
+            res.set_lod([out_offsets])
+            if len(pred_res) == 0:
+                pred_res = np.array([[1]], dtype=np.float32)
+            res.set(np.vstack(pred_res), fluid.CPUPlace())
             return res
 
         pred_result = create_tmp_var(
@@ -599,7 +620,7 @@ class MultiClassSoftNMS(object):
             shape=[-1, 6],
             lod_level=1)
         fluid.layers.py_func(
-            func=_soft_nms, x=[bboxes, scores], out=pred_result)
+            func=_batch_softnms, x=[bboxes, scores], out=pred_result)
         return pred_result
 
 
-- 
GitLab