merge cam, fix export

PaddlePaddle · Feb 2, 2023 · 967cf9b · 967cf9b
1 parent e564d74
commit 967cf9b
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 9 deletions.
diff --git a/ppdet/modeling/architectures/ppyoloe.py b/ppdet/modeling/architectures/ppyoloe.py
@@ -97,15 +97,21 @@ def _forward(self):
                     raise ValueError
             return yolo_losses
         else:
+            cam_data = {}  # record bbox scores and index before nms
             yolo_head_outs = self.yolo_head(neck_feats)
+            cam_data['scores'] = yolo_head_outs[0]
+
             if self.post_process is not None:
-                bbox, bbox_num = self.post_process(
+                bbox, bbox_num, before_nms_indexes = self.post_process(
                     yolo_head_outs, self.yolo_head.mask_anchors,
                     self.inputs['im_shape'], self.inputs['scale_factor'])
+                cam_data['before_nms_indexes'] = before_nms_indexes
             else:
-                bbox, bbox_num = self.yolo_head.post_process(
+                bbox, bbox_num, before_nms_indexes = self.yolo_head.post_process(
                     yolo_head_outs, self.inputs['scale_factor'])
-            output = {'bbox': bbox, 'bbox_num': bbox_num}
+                # data for cam
+                cam_data['before_nms_indexes'] = before_nms_indexes
+            output = {'bbox': bbox, 'bbox_num': bbox_num, 'cam_data': cam_data}
 
             return output
 
@@ -195,15 +201,21 @@ def _forward(self):
                 aux_pred=[aux_cls_scores, aux_bbox_preds])
             return loss
         else:
+            cam_data = {}  # record bbox scores and index before nms
             yolo_head_outs = self.yolo_head(neck_feats)
+            cam_data['scores'] = yolo_head_outs[0]
+
             if self.post_process is not None:
-                bbox, bbox_num = self.post_process(
+                bbox, bbox_num, before_nms_indexes = self.post_process(
                     yolo_head_outs, self.yolo_head.mask_anchors,
                     self.inputs['im_shape'], self.inputs['scale_factor'])
+                cam_data['before_nms_indexes'] = before_nms_indexes
             else:
-                bbox, bbox_num = self.yolo_head.post_process(
+                bbox, bbox_num, before_nms_indexes = self.yolo_head.post_process(
                     yolo_head_outs, self.inputs['scale_factor'])
-            output = {'bbox': bbox, 'bbox_num': bbox_num}
+                # data for cam
+                cam_data['before_nms_indexes'] = before_nms_indexes
+            output = {'bbox': bbox, 'bbox_num': bbox_num, 'cam_data': cam_data}
 
             return output
 

diff --git a/ppdet/modeling/heads/ppyoloe_head.py b/ppdet/modeling/heads/ppyoloe_head.py
@@ -480,7 +480,8 @@ def post_process(self, head_outs, scale_factor):
         pred_bboxes *= stride_tensor
         if self.exclude_post_process:
             return paddle.concat(
-                [pred_bboxes, pred_scores.transpose([0, 2, 1])], axis=-1), None
+                [pred_bboxes, pred_scores.transpose([0, 2, 1])],
+                axis=-1), None, None
         else:
             # scale bbox to origin
             scale_y, scale_x = paddle.split(scale_factor, 2, axis=-1)
@@ -490,9 +491,10 @@ def post_process(self, head_outs, scale_factor):
             pred_bboxes /= scale_factor
             if self.exclude_nms:
                 # `exclude_nms=True` just use in benchmark
-                return pred_bboxes, pred_scores
+                return pred_bboxes, pred_scores, None
             else:
-                bbox_pred, bbox_num, before_nms_indexes = self.nms(pred_bboxes, pred_scores)
+                bbox_pred, bbox_num, before_nms_indexes = self.nms(pred_bboxes,
+                                                                   pred_scores)
                 return bbox_pred, bbox_num, before_nms_indexes