[Fix]: Fix the error of TTA when there is no prediction results and s…

…upport vis (open-mmlab#9840)
ytzfhqs · Feb 27, 2023 · 13aa724 · 13aa724
1 parent 28b1d34
commit 13aa724
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 6 deletions.
diff --git a/configs/centernet/centernet_tta.py b/configs/centernet/centernet_tta.py
@@ -31,6 +31,7 @@
                     test_pad_mode=['logical_or', 31],
                     test_pad_add_pix=1),
             ],
+            [dict(type='LoadAnnotations', with_bbox=True)],
             [
                 dict(
                     type='PackDetInputs',

diff --git a/configs/retinanet/retinanet_tta.py b/configs/retinanet/retinanet_tta.py
@@ -12,7 +12,7 @@
         ], [
             dict(type='RandomFlip', prob=1.),
             dict(type='RandomFlip', prob=0.)
-        ],
+        ], [dict(type='LoadAnnotations', with_bbox=True)],
                     [
                         dict(
                             type='PackDetInputs',

diff --git a/configs/yolox/yolox_tta.py b/configs/yolox/yolox_tta.py
@@ -25,6 +25,7 @@
                     pad_to_square=True,
                     pad_val=dict(img=(114.0, 114.0, 114.0))),
             ],
+            [dict(type='LoadAnnotations', with_bbox=True)],
             [
                 dict(
                     type='PackDetInputs',

diff --git a/mmdet/models/test_time_augs/det_tta.py b/mmdet/models/test_time_augs/det_tta.py
@@ -103,7 +103,7 @@ def _merge_single_sample(
         to one prediction.
 
         Args:
-            data_samples_list (List[DetDataSample]): List of predictions
+            data_samples (List[DetDataSample]): List of predictions
             of enhanced data which come form one image.
         Returns:
             List[DetDataSample]: Merged prediction.
@@ -126,10 +126,7 @@ def _merge_single_sample(
         merged_labels = torch.cat(aug_labels, dim=0)
 
         if merged_bboxes.numel() == 0:
-            det_bboxes = torch.cat([merged_bboxes, merged_scores[:, None]], -1)
-            return [
-                (det_bboxes, merged_labels),
-            ]
+            return data_samples[0]
 
         det_bboxes, keep_idxs = batched_nms(merged_bboxes, merged_scores,
                                             merged_labels, self.tta_cfg.nms)