Add more formats of input

python/rapidocr_onnxruntime/rapid_ocr_api.py

@@ -48,7 +48,8 @@ class RapidOCR():
 
         self.load_img = LoadImage()
 
-    def __call__(self, img_content: Union[str, np.ndarray, bytes], **kwargs):
+    def __call__(self,
+                 img_content: Union[str, np.ndarray, bytes, Path], **kwargs):
         if kwargs:
             box_thresh = kwargs.get('box_thresh', 0.5)
             unclip_ratio = kwargs.get('unclip_ratio', 1.6)

python/rapidocr_onnxruntime/utils.py

@@ -15,7 +15,7 @@ from onnxruntime import (GraphOptimizationLevel, InferenceSession,
                          SessionOptions, get_available_providers, get_device)
 
 root_dir = Path(__file__).resolve().parent
-InputType = Union[str, np.ndarray, bytes]
+InputType = Union[str, np.ndarray, bytes, Path]
 
 
 class OrtInferSession():
@@ -98,7 +98,7 @@ class LoadImage():
         self.ndarray_format = ['ndarray']
 
     def __call__(self, img: InputType):
-        if not isinstance(img, (str, bytes, np.ndarray)):
+        if not isinstance(img, InputType.__args__):
             raise LoadImageError(
                 f'The img type {type(img)} does not in [str, np.ndarray, bytes]')
 
@@ -110,13 +110,13 @@ class LoadImage():
         img_suffix = self._which_format(img)
 
         if img_suffix in self.gif_format:
-            if isinstance(img, str):
+            if isinstance(img, (str, Path)):
                 return self._read_gif_str(img)
 
             if isinstance(img, bytes):
                 return self._read_gif_bytes(img)
         elif img_suffix in self.valid_format:
-            if isinstance(img, str):
+            if isinstance(img, (str, Path)):
                 return self._read_valid_img_str(img)
 
             if isinstance(img, bytes):
@@ -138,7 +138,7 @@ class LoadImage():
         return img_content
 
     def _which_format(self, img_content: InputType) -> Optional[str]:
-        if isinstance(img_content, str):
+        if isinstance(img_content, (str, Path)):
             with open(img_content, 'rb') as f:
                 return imghdr.what(f).lower()
 
@@ -151,10 +151,10 @@ class LoadImage():
         if isinstance(img_content, np.ndarray):
             return self.ndarray_format[0]
 
-    def _read_gif_str(self, img: str) -> np.ndarray:
+    def _read_gif_str(self, img: Union[str, Path]) -> np.ndarray:
         @contextmanager
         def video_capture(img):
-            cap = cv2.VideoCapture(img)
+            cap = cv2.VideoCapture(str(img))
             try:
                 yield cap
             finally:
@@ -170,7 +170,7 @@ class LoadImage():
         img = frame[:, :, ::-1]
         return img
 
-    def _read_valid_img_str(self, img: str) -> np.ndarray:
+    def _read_valid_img_str(self, img: Union[str, Path]) -> np.ndarray:
         with open(img, 'rb') as f:
             img = self._read_valid_img_bytes(f.read())
         return img
@@ -179,7 +179,7 @@ class LoadImage():
         with tempfile.NamedTemporaryFile(delete=True) as fp:
             fp.write(img)
             fp.seek(0)
-            img = self._read_gif_str(fp.name)
+            img = self._read_gif_str(fp)
         return img
 
     def _read_valid_img_bytes(self, img: bytes) -> np.ndarray:

python/rapidocr_openvino/rapid_ocr_api.py

@@ -8,7 +8,8 @@ from pathlib import Path
 
 import cv2
 import numpy as np
-import yaml
+from typing import Union
+from .utils import LoadImage, read_yaml
 
 root_dir = Path(__file__).resolve().parent
 sys.path.append(str(root_dir))
@@ -20,7 +21,7 @@ class RapidOCR():
         if not Path(config_path).exists():
             raise FileExistsError(f'{config_path} does not exist!')
 
-        config = self.read_yaml(config_path)
+        config = read_yaml(config_path)
 
         global_config = config['Global']
         self.print_verbose = global_config['print_verbose']
@@ -44,7 +45,10 @@ class RapidOCR():
                                               config['Cls']['class_name'])
             self.text_cls = TextClassifier(config['Cls'])
 
-    def __call__(self, img: np.ndarray, **kwargs):
+        self.load_img = LoadImage()
+
+    def __call__(self,
+                 img_content: Union[str, np.ndarray, bytes, Path], **kwargs):
         if kwargs:
             # 获得超参数
             box_thresh = kwargs.get('box_thresh', 0.5)
@@ -56,6 +60,7 @@ class RapidOCR():
             self.text_detector.postprocess_op.unclip_ratio = unclip_ratio
             self.text_score = text_score
 
+        img = self.load_img(img_content)
         h, w = img.shape[:2]
         if self.width_height_ratio == -1:
             use_limit_ratio = False

python/rapidocr_openvino/utils.py

@@ -1,13 +1,19 @@
 # -*- encoding: utf-8 -*-
 # @Author: SWHL
 # @Contact: liekkaskono@163.com
+import imghdr
+import tempfile
+from contextlib import contextmanager
 from pathlib import Path
+from typing import Optional, Union
 
+import cv2
 import numpy as np
 import yaml
 from openvino.runtime import Core
 
 root_dir = Path(__file__).resolve().parent
+InputType = Union[str, np.ndarray, bytes, Path]
 
 
 class OpenVINOInferSession():
@@ -37,3 +43,135 @@ def read_yaml(yaml_path):
     with open(yaml_path, 'rb') as f:
         data = yaml.load(f, Loader=yaml.Loader)
     return data
+
+
+class LoadImage():
+    def __init__(self, ):
+        self.valid_format = ['jpg', 'jpeg', 'png', 'tiff', 'tif', 'bmp']
+        self.gif_format = ['gif']
+        self.ndarray_format = ['ndarray']
+
+    def __call__(self, img: InputType):
+        if not isinstance(img, InputType.__args__):
+            raise LoadImageError(
+                f'The img type {type(img)} does not in [str, np.ndarray, bytes]')
+
+        img = self.load_direct_img(img)
+        img = self.convert_three_channels(img)
+        return img
+
+    def load_direct_img(self, img: InputType) -> np.ndarray:
+        img_suffix = self._which_format(img)
+
+        if img_suffix in self.gif_format:
+            if isinstance(img, (str, Path)):
+                return self._read_gif_str(img)
+
+            if isinstance(img, bytes):
+                return self._read_gif_bytes(img)
+        elif img_suffix in self.valid_format:
+            if isinstance(img, (str, Path)):
+                return self._read_valid_img_str(img)
+
+            if isinstance(img, bytes):
+                return self._read_valid_img_bytes(img)
+        elif img_suffix in self.ndarray_format:
+            return img
+        else:
+            raise LoadImageError(f'{img_suffix} is not supported!')
+
+    def convert_three_channels(self, img_content: np.ndarray) -> np.ndarray:
+        img_shape_len = len(img_content.shape)
+        if img_shape_len == 2:
+            return cv2.cvtColor(img_content, cv2.COLOR_GRAY2BGR)
+
+        if img_content.shape[2] == 4:
+            # RGBA → RGB
+            return self._read_transparent_png(img_content)
+
+        return img_content
+
+    def _which_format(self, img_content: InputType) -> Optional[str]:
+        if isinstance(img_content, (str, Path)):
+            with open(img_content, 'rb') as f:
+                return imghdr.what(f).lower()
+
+        if isinstance(img_content, bytes):
+            with tempfile.NamedTemporaryFile(delete=True) as fp:
+                fp.write(img_content)
+                fp.seek(0)
+                return imghdr.what(fp).lower()
+
+        if isinstance(img_content, np.ndarray):
+            return self.ndarray_format[0]
+
+    def _read_gif_str(self, img: Union[str, Path]) -> np.ndarray:
+        @contextmanager
+        def video_capture(img):
+            cap = cv2.VideoCapture(str(img))
+            try:
+                yield cap
+            finally:
+                cap.release()
+
+        with video_capture(img) as cap:
+            ret, frame = cap.read()
+        if not ret:
+            raise LoadImageError('OpenCV does not read the gif img')
+
+        if len(frame.shape) == 2 or frame.shape[-1] == 1:
+            frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB)
+        img = frame[:, :, ::-1]
+        return img
+
+    def _read_valid_img_str(self, img: Union[str, Path]) -> np.ndarray:
+        with open(img, 'rb') as f:
+            img = self._read_valid_img_bytes(f.read())
+        return img
+
+    def _read_gif_bytes(self, img: bytes) -> np.ndarray:
+        with tempfile.NamedTemporaryFile(delete=True) as fp:
+            fp.write(img)
+            fp.seek(0)
+            img = self._read_gif_str(fp)
+        return img
+
+    def _read_valid_img_bytes(self, img: bytes) -> np.ndarray:
+        np_arr = np.frombuffer(img, dtype=np.uint8)
+        cv_img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
+        if cv_img is None:
+            from io import BytesIO
+            from PIL import Image, UnidentifiedImageError
+
+            bytes_img = BytesIO(img)
+            try:
+                cv_img = np.array(Image.open(bytes_img))
+            except (UnidentifiedImageError, ValueError) as e:
+                raise LoadImageError('read_valid_img_bytes has errors.') from e
+        return cv_img
+
+    @staticmethod
+    def _read_transparent_png(image_4channel: np.ndarray) -> np.ndarray:
+        if image_4channel.shape[-1] != 4:
+            return image_4channel
+        alpha_channel = image_4channel[:, :, 3]
+        rgb_channels = image_4channel[:, :, :3]
+
+        # White Background Image
+        white_bg_img = np.ones_like(rgb_channels, dtype=np.uint8) * 255
+
+        # Alpha factor
+        alpha_channel = alpha_channel[:, :, np.newaxis].astype(np.float32)
+        alpha_factor = alpha_channel / 255.0
+        alpha_factor = np.concatenate(
+            (alpha_factor, alpha_factor, alpha_factor), axis=2)
+
+        # Transparent Image Rendered on White Background
+        base = rgb_channels.astype(np.float32) * alpha_factor
+        white = white_bg_img.astype(np.float32) * (1 - alpha_factor)
+        final_image = base + white
+        return final_image.astype(np.uint8)
+
+
+class LoadImageError(Exception):
+    pass

python/tests/test_all_ort.py

@@ -54,3 +54,10 @@ def test_input_bytes():
         result, _ = rapid_ocr(f.read())
     assert result[0][1] == '正品促销'
     assert len(result) == 17
+
+
+def test_input_path():
+    image_path = tests_dir / 'ch_en_num.jpg'
+    result, _ = rapid_ocr(image_path)
+    assert result[0][1] == '正品促销'
+    assert len(result) == 17

python/tests/test_all_vino.py

@@ -13,6 +13,7 @@ root_dir = Path(__file__).resolve().parent.parent
 sys.path.append(str(root_dir))
 
 from rapidocr_openvino import RapidOCR
+from rapidocr_openvino.utils import LoadImageError
 
 rapid_ocr = RapidOCR()
 
@@ -29,13 +30,35 @@ def test_normal():
 
 def test_empty():
     img = None
-    with pytest.raises(AttributeError) as exc_info:
+    with pytest.raises(LoadImageError) as exc_info:
         rapid_ocr(img)
-        raise AttributeError
-    assert exc_info.type is AttributeError
+        raise LoadImageError
+    assert exc_info.type is LoadImageError
 
 
 def test_zeros():
     img = np.zeros([640, 640, 3])
     result, _ = rapid_ocr(img)
     assert result is None
+
+
+def test_input_str():
+    image_path = tests_dir / 'ch_en_num.jpg'
+    result, _ = rapid_ocr(str(image_path))
+    assert result[0][1] == '正品促销'
+    assert len(result) == 17
+
+
+def test_input_bytes():
+    image_path = tests_dir / 'ch_en_num.jpg'
+    with open(image_path, 'rb') as f:
+        result, _ = rapid_ocr(f.read())
+    assert result[0][1] == '正品促销'
+    assert len(result) == 17
+
+
+def test_input_path():
+    image_path = tests_dir / 'ch_en_num.jpg'
+    result, _ = rapid_ocr(image_path)
+    assert result[0][1] == '正品促销'
+    assert len(result) == 17