From a11329ec3e7bd989dda712b69a5f554f471b724c Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 7 Jul 2023 11:41:28 +0800 Subject: [PATCH] [Fix] fix a rec inference vis bug --- mmocr/apis/inferencers/mmocr_inferencer.py | 3 ++- tests/test_apis/test_inferencers/test_mmocr_inferencer.py | 7 +++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/mmocr/apis/inferencers/mmocr_inferencer.py b/mmocr/apis/inferencers/mmocr_inferencer.py index c531be35c..4aaebb55f 100644 --- a/mmocr/apis/inferencers/mmocr_inferencer.py +++ b/mmocr/apis/inferencers/mmocr_inferencer.py @@ -232,7 +232,8 @@ def visualize(self, inputs: InputsType, preds: PredType, **kwargs) else: return self.textrec_inferencer.visualize( - self.rec_inputs, preds['rec'][0], **kwargs) + self.rec_inputs, [pred[0] for pred in preds['rec']], + **kwargs) else: return self.textdet_inferencer.visualize(inputs, preds['det'], **kwargs) diff --git a/tests/test_apis/test_inferencers/test_mmocr_inferencer.py b/tests/test_apis/test_inferencers/test_mmocr_inferencer.py index 8628cd44f..6c8f9bd6b 100644 --- a/tests/test_apis/test_inferencers/test_mmocr_inferencer.py +++ b/tests/test_apis/test_inferencers/test_mmocr_inferencer.py @@ -150,6 +150,13 @@ def test_rec(self, mock_load): osp.join(tmp_dir, 'preds', pred_dir)) self.assert_predictions_equal(res['predictions'][i], dumped_res) + # test different batch sizes + img_dir = 'tests/data/rec_toy_dataset/imgs' + res_bs3 = inferencer(img_dir, batch_size=3, return_vis=True) + self.assertIn('visualization', res_bs3) + self.assertIn('predictions', res_bs3) + self.assertEqual( + len(res_bs3['predictions']), len(res_bs3['visualization'])) @mock.patch('mmengine.infer.infer._load_checkpoint') def test_det_rec(self, mock_load):