add ignore_orientation

ppocr/data/imaug/operators.py

@@ -28,9 +28,14 @@ import numpy as np
 class DecodeImage(object):
     """ decode image """
 
-    def __init__(self, img_mode='RGB', channel_first=False, **kwargs):
+    def __init__(self,
+                 img_mode='RGB',
+                 channel_first=False,
+                 ignore_orientation=False,
+                 **kwargs):
         self.img_mode = img_mode
         self.channel_first = channel_first
+        self.ignore_orientation = ignore_orientation
 
     def __call__(self, data):
         img = data['image']
@@ -41,7 +46,11 @@ class DecodeImage(object):
             assert type(img) is bytes and len(
                 img) > 0, "invalid input 'img' in DecodeImage"
         img = np.frombuffer(img, dtype='uint8')
-        img = cv2.imdecode(img, 1)
+        if self.ignore_orientation:
+            img = cv2.imdecode(img, cv2.IMREAD_IGNORE_ORIENTATION |
+                               cv2.IMREAD_COLOR)
+        else:
+            img = cv2.imdecode(img, 1)
         if img is None:
             return None
         if self.img_mode == 'GRAY':
@@ -60,14 +69,9 @@ class DecodeImage(object):
 class NRTRDecodeImage(object):
     """ decode image """
 
-    def __init__(self,
-                 img_mode='RGB',
-                 channel_first=False,
-                 ignore_orientation=False,
-                 **kwargs):
+    def __init__(self, img_mode='RGB', channel_first=False, **kwargs):
         self.img_mode = img_mode
         self.channel_first = channel_first
-        self.ignore_orientation = ignore_orientation
 
     def __call__(self, data):
         img = data['image']
@@ -79,11 +83,7 @@ class NRTRDecodeImage(object):
                 img) > 0, "invalid input 'img' in DecodeImage"
         img = np.frombuffer(img, dtype='uint8')
 
-        if self.ignore_orientation:
-            img = cv2.imdecode(img, cv2.IMREAD_IGNORE_ORIENTATION |
-                               cv2.IMREAD_COLOR)
-        else:
-            img = cv2.imdecode(img, 1)
+        img = cv2.imdecode(img, 1)
 
         if img is None:
             return None