diff --git a/python/rapidocr/utils/load_image.py b/python/rapidocr/utils/load_image.py index 86ab953..6e7eb5d 100644 --- a/python/rapidocr/utils/load_image.py +++ b/python/rapidocr/utils/load_image.py @@ -8,7 +8,7 @@ from typing import Any, Union import cv2 import numpy as np import requests -from PIL import Image, UnidentifiedImageError +from PIL import Image, ImageOps, UnidentifiedImageError from .utils import is_url @@ -33,12 +33,14 @@ class LoadImage: def load_img(self, img: InputType) -> np.ndarray: if isinstance(img, (str, Path)): - if is_url(img): + if is_url(str(img)): img = Image.open(requests.get(img, stream=True, timeout=60).raw) else: self.verify_exist(img) img = Image.open(img) + img = self.exif_transpose(img) + try: img = self.img_to_ndarray(img) except UnidentifiedImageError as e: @@ -57,6 +59,21 @@ class LoadImage: raise LoadImageError(f"{type(img)} is not supported!") + @staticmethod + def verify_exist(file_path: Union[str, Path]): + if not Path(file_path).exists(): + raise LoadImageError(f"{file_path} does not exist.") + + @staticmethod + def exif_transpose(img: Image.Image) -> Image.Image: + try: + img_corrected = ImageOps.exif_transpose(img) + if img_corrected is None: + return img + return img_corrected + except Exception as e: + return img + def img_to_ndarray(self, img: Image.Image) -> np.ndarray: if img.mode == "1": img = img.convert("L") @@ -121,11 +138,6 @@ class LoadImage: new_img = cv2.bitwise_not(new_img) return new_img - @staticmethod - def verify_exist(file_path: Union[str, Path]): - if not Path(file_path).exists(): - raise LoadImageError(f"{file_path} does not exist.") - class LoadImageError(Exception): pass diff --git a/python/tests/test_det_cls_rec.py b/python/tests/test_det_cls_rec.py index 5fbfd73..27cd2ad 100644 --- a/python/tests/test_det_cls_rec.py +++ b/python/tests/test_det_cls_rec.py @@ -31,7 +31,7 @@ def test_only_cls(engine): img_path = tests_dir / "text_cls.jpg" result = engine(img_path, use_det=False, use_cls=True, use_rec=False) assert len(result) == 1 - assert result.cls_res[0][0] == "0" + assert result.cls_res[0][0] == "180" def test_only_rec(engine): @@ -53,7 +53,7 @@ def test_cls_rec(engine): assert result is not None assert len(result) == 1 - assert result.txts[0] == "韩国小馆" + assert result.txts[0] == "怪我咯" def test_det_cls_rec(engine): diff --git a/python/tests/test_files/img_exif_orientation.jpg b/python/tests/test_files/img_exif_orientation.jpg new file mode 100644 index 0000000..46c078f Binary files /dev/null and b/python/tests/test_files/img_exif_orientation.jpg differ diff --git a/python/tests/test_files/text_cls.jpg b/python/tests/test_files/text_cls.jpg index fcd7e12..9fc0b79 100644 Binary files a/python/tests/test_files/text_cls.jpg and b/python/tests/test_files/text_cls.jpg differ diff --git a/python/tests/test_input.py b/python/tests/test_input.py index d1c58f8..cc744ec 100644 --- a/python/tests/test_input.py +++ b/python/tests/test_input.py @@ -14,8 +14,8 @@ sys.path.append(str(root_dir)) from rapidocr import LoadImageError, RapidOCR -tests_dir = root_dir / "tests" / "test_files" -img_path = tests_dir / "ch_en_num.jpg" +test_dir = root_dir / "tests" / "test_files" +img_path = test_dir / "ch_en_num.jpg" @pytest.fixture() @@ -24,6 +24,12 @@ def engine(): return engine +def test_exif_transpose(engine): + img_path = test_dir / "img_exif_orientation.jpg" + result = engine(img_path, use_cls=False) + assert result.txts[0] == "我是中国人" + + @mark.parametrize( "img_name,gt", [ @@ -38,7 +44,7 @@ def engine(): ], ) def test_transparent_img(engine, img_name: str, gt: str): - img_path = tests_dir / img_name + img_path = test_dir / img_name result = engine(img_path) assert result.txts[0] == gt @@ -52,7 +58,7 @@ def test_long_img(engine): def test_full_black_img(engine): - img_path = tests_dir / "empty_black.jpg" + img_path = test_dir / "empty_black.jpg" result = engine(img_path) assert result.img is None assert result.boxes is None @@ -104,7 +110,7 @@ def test_input_parameters(engine): def test_input_three_ndim_two_channel(engine): - img_npy = tests_dir / "two_dim_image.npy" + img_npy = test_dir / "two_dim_image.npy" image_array = np.load(str(img_npy)) result = engine(image_array) @@ -124,7 +130,7 @@ def test_input_three_ndim_one_channel(engine): def test_mode_one_img(engine): - img_path = tests_dir / "issue_170.png" + img_path = test_dir / "issue_170.png" result = engine(img_path) assert result.txts[0] == "TEST" @@ -141,7 +147,7 @@ def test_mode_one_img(engine): ], ) def test_letterbox_like(engine, img_name, gt_len, gt_first_len): - img_path = tests_dir / img_name + img_path = test_dir / img_name result = engine(img_path) assert len(result) == gt_len