Merge pull request #6781 from WenmuZhou/xlm

fix shape unmatch error in PSE postprocess
2026-09-24 23:33:08 +08:00 · 2022-07-05 15:43:46 +08:00
parent bb95527f7d b541cc12be
commit e97daf7d97
1 changed files with 2 additions and 0 deletions
@@ -58,6 +58,8 @@ class PSEPostProcess(object):

        kernels = (pred > self.thresh).astype('float32')
        text_mask = kernels[:, 0, :, :]
+        text_mask = paddle.unsqueeze(text_mask, axis=1)
+
        kernels[:, 0:, :, :] = kernels[:, 0:, :, :] * text_mask

        score = score.numpy()