fix shape unmatch error

b541cc12 · 文幕地方 · 8d841440 · b541cc12
隐藏空白更改
内联并排

Showing with 2 addition and 0 deletion

ppocr/postprocess/pse_postprocess/pse_postprocess.py ppocr/postprocess/pse_postprocess/pse_postprocess.py +2 -0

未找到文件。
--- a/ppocr/postprocess/pse_postprocess/pse_postprocess.py
+++ b/ppocr/postprocess/pse_postprocess/pse_postprocess.py
@@ -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()