diff --git a/python/rapidocr_onnxruntime/rapid_ocr_api.py b/python/rapidocr_onnxruntime/rapid_ocr_api.py index 8077fac..31a3d5d 100644 --- a/python/rapidocr_onnxruntime/rapid_ocr_api.py +++ b/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) diff --git a/python/rapidocr_onnxruntime/utils.py b/python/rapidocr_onnxruntime/utils.py index b41fef0..53249d7 100644 --- a/python/rapidocr_onnxruntime/utils.py +++ b/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: diff --git a/python/rapidocr_openvino/rapid_ocr_api.py b/python/rapidocr_openvino/rapid_ocr_api.py index 9026ca6..68ecd67 100644 --- a/python/rapidocr_openvino/rapid_ocr_api.py +++ b/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 diff --git a/python/rapidocr_openvino/utils.py b/python/rapidocr_openvino/utils.py index 0db97c5..f4caf4e 100644 --- a/python/rapidocr_openvino/utils.py +++ b/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 diff --git a/python/tests/test_all_ort.py b/python/tests/test_all_ort.py index f8c40c6..6443c23 100644 --- a/python/tests/test_all_ort.py +++ b/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 diff --git a/python/tests/test_all_vino.py b/python/tests/test_all_vino.py index 479fea0..de5d6a5 100644 --- a/python/tests/test_all_vino.py +++ b/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