[Enhance] Replace data_dict calling 'img' key to support MMDet3D (ope…

…n-mmlab#514) * remove dict calling img key for compatibility * fix unit test * infer batch size using len(result) to be consistent with mmcv
bowenroom · Apr 29, 2021 · ce56e68 · ce56e68
1 parent 6ccb1c0
commit ce56e68
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/mmseg/apis/test.py b/mmseg/apis/test.py
@@ -97,7 +97,7 @@ def single_gpu_test(model,
                 result = np2tmp(result)
             results.append(result)
 
-        batch_size = data['img'][0].size(0)
+        batch_size = len(result)
         for _ in range(batch_size):
             prog_bar.update()
     return results

diff --git a/mmseg/models/segmentors/base.py b/mmseg/models/segmentors/base.py
@@ -155,7 +155,7 @@ def train_step(self, data_batch, optimizer, **kwargs):
         outputs = dict(
             loss=loss,
             log_vars=log_vars,
-            num_samples=len(data_batch['img'].data))
+            num_samples=len(data_batch['img_metas']))
 
         return outputs