Fix Pad Shape (#10997)

mmdet/models/data_preprocessors/data_preprocessor.py

@@ -170,10 +170,10 @@ def _get_pad_shape(self, data: dict) -> List[tuple]:
                 'or a list of tensor, but got a tensor with shape: '
                 f'{_batch_inputs.shape}')
             pad_h = int(
-                np.ceil(_batch_inputs.shape[1] /
+                np.ceil(_batch_inputs.shape[2] /
                         self.pad_size_divisor)) * self.pad_size_divisor
             pad_w = int(
-                np.ceil(_batch_inputs.shape[2] /
+                np.ceil(_batch_inputs.shape[3] /
                         self.pad_size_divisor)) * self.pad_size_divisor
             batch_pad_shape = [(pad_h, pad_w)] * _batch_inputs.shape[0]
         else:

tests/test_models/test_data_preprocessors/test_data_preprocessor.py

@@ -34,6 +34,29 @@ def test_init(self):
         with self.assertRaises(AssertionError):
             DetDataPreprocessor(bgr_to_rgb=True, rgb_to_bgr=True)
 
+    def test_pad_shape(self):
+        processor = DetDataPreprocessor()
+        data = dict(inputs=torch.randint(0, 256, (2, 3, 10, 15)))
+        batch_pad_shape = processor._get_pad_shape(data)
+        self.assertEqual(batch_pad_shape, [(10, 15), (10, 15)])
+
+        data = dict(inputs=[torch.randint(0, 256, (3, 11, 10))])
+        self.assertEqual(processor._get_pad_shape(data), [(11, 10)])
+
+        # batch with different image sizes
+        data = dict(inputs=[
+            torch.randint(0, 256, (3, 10, 16)),
+            torch.randint(0, 256, (3, 15, 20))
+        ])
+
+        self.assertEqual(processor._get_pad_shape(data), [(10, 16), (15, 20)])
+
+        # test with pad divisor
+        processor = DetDataPreprocessor(pad_size_divisor=10)
+        data = dict(inputs=torch.randint(0, 256, (2, 3, 52, 65)))
+        self.assertAlmostEqual(
+            processor._get_pad_shape(data), [(60, 70), (60, 70)])
+
     def test_forward(self):
         processor = DetDataPreprocessor(mean=[0, 0, 0], std=[1, 1, 1])