diff --git a/ppocr/data/imaug/operators.py b/ppocr/data/imaug/operators.py index 2ec51c286de2db94c89ec3023ffd51742cbe09ac..920e7fee13e4b3dfbb3d26f285e05fdcc97e8d45 100644 --- a/ppocr/data/imaug/operators.py +++ b/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