diff --git a/mmocr/models/textdet/dense_heads/head_mixin.py b/mmocr/models/textdet/dense_heads/head_mixin.py index fc553bd3..67f8821a 100644 --- a/mmocr/models/textdet/dense_heads/head_mixin.py +++ b/mmocr/models/textdet/dense_heads/head_mixin.py @@ -59,7 +59,9 @@ class HeadMixin: boundaries = self.resize_boundary( boundaries, 1.0 / self.downsample_ratio / img_metas[0]['scale_factor']) - results = dict(boundary_result=boundaries) + results = dict( + boundary_result=boundaries, filename=img_metas[0]['filename']) + return results def loss(self, pred_maps, **kwargs): diff --git a/tests/test_models/test_detector.py b/tests/test_models/test_detector.py index 91522254..9ed75a15 100644 --- a/tests/test_models/test_detector.py +++ b/tests/test_models/test_detector.py @@ -380,7 +380,8 @@ def test_textsnake(cfg_file): one_meta = img_metas[0] result = detector.bbox_head.get_boundary(maps, [one_meta], False) - assert len(result) == 1 + assert 'boundary_result' in result + assert 'filename' in result # Test show result results = {'boundary_result': [[0, 0, 1, 0, 1, 1, 0, 1, 0.9]]} @@ -519,7 +520,7 @@ def test_drrg(cfg_file): detector.bbox_head.out_conv.bias.data.fill_(0.) outs = detector.bbox_head.single_test(maps) boundaries = detector.bbox_head.get_boundary(*outs, img_metas, True) - assert len(boundaries) + assert len(boundaries) == 1 # Test show result results = {'boundary_result': [[0, 0, 1, 0, 1, 1, 0, 1, 0.9]]}