improve the function of simple_test_bboxes

mmdet/models/roi_heads/bbox_heads/sabl_head.py

@@ -274,6 +274,15 @@ def side_aware_split(self, feat):
         feat = torch.cat([feat_fl, feat_fr], dim=-1)
         return feat
 
+    def bbox_pred_split(self, bbox_pred, num_proposals_per_img):
+        """Split batch bbox prediction back to each image."""
+        bucket_cls_preds, bucket_offset_preds = bbox_pred
+        bucket_cls_preds = bucket_cls_preds.split(num_proposals_per_img, 0)
+        bucket_offset_preds = bucket_offset_preds.split(
+            num_proposals_per_img, 0)
+        bbox_pred = tuple(zip(bucket_cls_preds, bucket_offset_preds))
+        return bbox_pred
+
     def reg_forward(self, reg_x):
         outs = self.side_aware_feature_extractor(reg_x)
         edge_offset_preds = []

mmdet/models/roi_heads/test_mixins.py

@@ -67,16 +67,17 @@ def simple_test_bboxes(self,
         num_proposals_per_img = tuple(len(p) for p in proposals)
         rois = rois.split(num_proposals_per_img, 0)
         cls_score = cls_score.split(num_proposals_per_img, 0)
-        # the type of bbox_pred in sabl is tuple not Tensor
-        if isinstance(bbox_pred, torch.Tensor):
-            # some detector with_reg is False, bbox_pred will be None
-            bbox_pred = bbox_pred.split(
-                num_proposals_per_img,
-                0) if bbox_pred is not None else [None, None]
+
+        # some detector with_reg is False, bbox_pred will be None
+        if bbox_pred is not None:
+            # the bbox prediction of some detectors like SABL is not Tensor
+            if isinstance(bbox_pred, torch.Tensor):
+                bbox_pred = bbox_pred.split(num_proposals_per_img, 0)
+            else:
+                bbox_pred = self.bbox_head.bbox_pred_split(
+                    bbox_pred, num_proposals_per_img)
         else:
-            # some detector with_reg is False, bbox_pred will be None
-            bbox_pred = (
-                bbox_pred, ) if bbox_pred is not None else [None, None]
+            bbox_pred = (None, ) * len(proposals)
 
         # apply bbox post-processing to each image individually
         det_bboxes = []