diff --git a/ppdet/modeling/anchor_heads/yolo_head.py b/ppdet/modeling/anchor_heads/yolo_head.py index dfcec0d87f7e124550ecac6cf97299173e9a03c6..132045deaee2e22f4a1f962479bc9634f03d9d5e 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 @@ -333,6 +333,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 4065e7d36ad2bec57b5d08978bd52276e8a74bec..972fcdf16399422d247a056f364f8ee0a4db3904 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