Merge pull request #2357 from caopulan/fix_srn_post

fix srn_postprocess
This commit is contained in:
xiaoting 2021-03-31 11:11:56 +08:00 committed by GitHub
commit a09604f897
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 2 deletions

View File

@ -216,6 +216,7 @@ class SRNLabelDecode(BaseRecLabelDecode):
character_type='en',
use_space_char=False,
**kwargs):
self.max_text_length = kwargs['max_text_length']
super(SRNLabelDecode, self).__init__(character_dict_path,
character_type, use_space_char)
@ -229,9 +230,9 @@ class SRNLabelDecode(BaseRecLabelDecode):
preds_idx = np.argmax(pred, axis=1)
preds_prob = np.max(pred, axis=1)
preds_idx = np.reshape(preds_idx, [-1, 25])
preds_idx = np.reshape(preds_idx, [-1, self.max_text_length])
preds_prob = np.reshape(preds_prob, [-1, 25])
preds_prob = np.reshape(preds_prob, [-1, self.max_text_length])
text = self.decode(preds_idx, preds_prob)