From 94cb456854da915b03fae265122866562b2b9947 Mon Sep 17 00:00:00 2001
From: Leif <4603009@qq.com>
Date: Thu, 21 Jan 2021 21:17:28 +0800
Subject: [PATCH] fix image autoRotate & add limitation on image size

---
 PPOCRLabel/PPOCRLabel.py      |  8 ++++++--
 PPOCRLabel/libs/autoDialog.py | 11 +++++++++--
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/PPOCRLabel/PPOCRLabel.py b/PPOCRLabel/PPOCRLabel.py
index 9de2467661..4d9c52740a 100644
--- a/PPOCRLabel/PPOCRLabel.py
+++ b/PPOCRLabel/PPOCRLabel.py
@@ -24,6 +24,7 @@
 from functools import partial
 from collections import defaultdict
 import json
+import cv2
 
 
 __dir__ = os.path.dirname(os.path.abspath(__file__))
@@ -1242,10 +1243,13 @@ def loadFile(self, filePath=None):
         #     if unicodeFilePath in self.mImgList:
                 
         if unicodeFilePath and os.path.exists(unicodeFilePath):
-            self.imageData = read(unicodeFilePath, None)
             self.canvas.verified = False
 
-            image = QImage.fromData(self.imageData)
+            cvimg = cv2.imdecode(np.fromfile(unicodeFilePath, dtype=np.uint8), 1)
+            height, width, depth = cvimg.shape
+            cvimg = cv2.cvtColor(cvimg, cv2.COLOR_BGR2RGB)
+            image = QImage(cvimg.data, width, height, width * depth, QImage.Format_RGB888)
+
             if image.isNull():
                 self.errorMessage(u'Error opening file',
                                   u"<p>Make sure <i>%s</i> is a valid image file." % unicodeFilePath)
diff --git a/PPOCRLabel/libs/autoDialog.py b/PPOCRLabel/libs/autoDialog.py
index 09fbeb56b3..3374e92cc5 100644
--- a/PPOCRLabel/libs/autoDialog.py
+++ b/PPOCRLabel/libs/autoDialog.py
@@ -7,6 +7,8 @@
     from PyQt4.QtCore import *
 
 import json
+import cv2
+import numpy as np
 
 from libs.utils import newIcon
 
@@ -34,11 +36,16 @@ def run(self):
                 if self.handle == 0:
                     self.listValue.emit(Imgpath)
                     if self.model == 'paddle':
-                        self.result_dic = self.ocr.ocr(Imgpath, cls=True, det=True)
+                        h, w, _ = cv2.imdecode(np.fromfile(Imgpath, dtype=np.uint8), 1).shape
+                        if h > 32 and w > 32:
+                            self.result_dic = self.ocr.ocr(Imgpath, cls=True, det=True)
+                        else:
+                            print('The size of', Imgpath, 'is too small to be recognised')
+                            self.result_dic = None
 
                     # 结果保存
                     if self.result_dic is None or len(self.result_dic) == 0:
-                        print('Can not recognise file  is :  ', Imgpath)
+                        print('Can not recognise file', Imgpath)
                         pass
                     else:
                         strs = ''