From 220390dfc3f21b82ddf320f77e88a98dd7e56ada Mon Sep 17 00:00:00 2001 From: SWHL Date: Fri, 10 Mar 2023 18:27:13 +0800 Subject: [PATCH] Optim Load_image class --- python/rapidocr_onnxruntime/utils.py | 154 +++++++------------------- python/rapidocr_openvino/utils.py | 155 +++++++-------------------- python/setup_onnxruntime.py | 2 +- python/setup_openvino.py | 2 +- 4 files changed, 84 insertions(+), 229 deletions(-) diff --git a/python/rapidocr_onnxruntime/utils.py b/python/rapidocr_onnxruntime/utils.py index c297d0d..7a86351 100644 --- a/python/rapidocr_onnxruntime/utils.py +++ b/python/rapidocr_onnxruntime/utils.py @@ -2,18 +2,17 @@ # @Author: SWHL # @Contact: liekkaskono@163.com import argparse -import imghdr -import tempfile import warnings -from contextlib import contextmanager +from io import BytesIO from pathlib import Path -from typing import Optional, Union +from typing import Union import cv2 import numpy as np import yaml from onnxruntime import (GraphOptimizationLevel, InferenceSession, SessionOptions, get_available_providers, get_device) +from PIL import Image, UnidentifiedImageError root_dir = Path(__file__).resolve().parent InputType = Union[str, np.ndarray, bytes, Path] @@ -94,130 +93,59 @@ class ONNXRuntimeError(Exception): class LoadImage(): def __init__(self, ): - self.valid_format = ['jpg', 'jpeg', 'png', 'tiff', 'tif', 'bmp'] - self.gif_format = ['gif'] - self.ndarray_format = ['ndarray'] + pass - def __call__(self, img: InputType): + def __call__(self, img: InputType) -> np.ndarray: if not isinstance(img, InputType.__args__): raise LoadImageError( - f'The img type {type(img)} does not in [str, np.ndarray, bytes]') + f'The img type {type(img)} does not in {InputType.__args__}') + + img = self.load_img(img) + + if img.ndim == 2: + return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) + + if img.ndim == 3 and img.shape[2] == 4: + return self.cvt_four_to_three(img) - 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: + def load_img(self, img: InputType) -> np.ndarray: + if isinstance(img, (str, Path)): + self.verify_exist(img) + try: + img = np.array(Image.open(img)) + except UnidentifiedImageError as e: + raise LoadImageError( + f'cannot identify image file {img}') from e 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 isinstance(img, bytes): + return np.array(Image.open(BytesIO(img))) - if img_content.shape[2] == 4: - # RGBA → RGB - return self._read_transparent_png(img_content) + if isinstance(img, np.ndarray): + return img - 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 + raise LoadImageError(f'{type(img)} is not supported!') @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] + def cvt_four_to_three(img: np.ndarray) -> np.ndarray: + '''RGBA → RGB + ''' + r, g, b, a = cv2.split(img) + new_img = cv2.merge((b, g, r)) - # White Background Image - white_bg_img = np.ones_like(rgb_channels, dtype=np.uint8) * 255 + not_a = cv2.bitwise_not(a) + not_a = cv2.cvtColor(not_a, cv2.COLOR_GRAY2BGR) - # 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) + new_img = cv2.bitwise_and(new_img, new_img, mask=a) + new_img = cv2.add(new_img, not_a) + return new_img - # 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) + @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): diff --git a/python/rapidocr_openvino/utils.py b/python/rapidocr_openvino/utils.py index d4c5163..ac623d1 100644 --- a/python/rapidocr_openvino/utils.py +++ b/python/rapidocr_openvino/utils.py @@ -2,17 +2,15 @@ # @Author: SWHL # @Contact: liekkaskono@163.com import argparse -import imghdr -import tempfile -import warnings -from contextlib import contextmanager +from io import BytesIO from pathlib import Path -from typing import Optional, Union +from typing import Union import cv2 import numpy as np import yaml from openvino.runtime import Core +from PIL import Image, UnidentifiedImageError root_dir = Path(__file__).resolve().parent InputType = Union[str, np.ndarray, bytes, Path] @@ -43,130 +41,59 @@ class OpenVINOInferSession(): class LoadImage(): def __init__(self, ): - self.valid_format = ['jpg', 'jpeg', 'png', 'tiff', 'tif', 'bmp'] - self.gif_format = ['gif'] - self.ndarray_format = ['ndarray'] + pass - def __call__(self, img: InputType): + def __call__(self, img: InputType) -> np.ndarray: if not isinstance(img, InputType.__args__): raise LoadImageError( - f'The img type {type(img)} does not in [str, np.ndarray, bytes]') + f'The img type {type(img)} does not in {InputType.__args__}') + + img = self.load_img(img) + + if img.ndim == 2: + return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) + + if img.ndim == 3 and img.shape[2] == 4: + return self.cvt_four_to_three(img) - 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: + def load_img(self, img: InputType) -> np.ndarray: + if isinstance(img, (str, Path)): + self.verify_exist(img) + try: + img = np.array(Image.open(img)) + except UnidentifiedImageError as e: + raise LoadImageError( + f'cannot identify image file {img}') from e 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 isinstance(img, bytes): + return np.array(Image.open(BytesIO(img))) - if img_content.shape[2] == 4: - # RGBA → RGB - return self._read_transparent_png(img_content) + if isinstance(img, np.ndarray): + return img - 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 + raise LoadImageError(f'{type(img)} is not supported!') @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] + def cvt_four_to_three(img: np.ndarray) -> np.ndarray: + '''RGBA → RGB + ''' + r, g, b, a = cv2.split(img) + new_img = cv2.merge((b, g, r)) - # White Background Image - white_bg_img = np.ones_like(rgb_channels, dtype=np.uint8) * 255 + not_a = cv2.bitwise_not(a) + not_a = cv2.cvtColor(not_a, cv2.COLOR_GRAY2BGR) - # 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) + new_img = cv2.bitwise_and(new_img, new_img, mask=a) + new_img = cv2.add(new_img, not_a) + return new_img - # 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) + @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): diff --git a/python/setup_onnxruntime.py b/python/setup_onnxruntime.py index c95dcf8..df205d2 100644 --- a/python/setup_onnxruntime.py +++ b/python/setup_onnxruntime.py @@ -71,7 +71,7 @@ setuptools.setup( include_package_data=True, install_requires=["pyclipper>=1.2.1", "onnxruntime>=1.7.0", "opencv_python>=4.5.1.48", "numpy>=1.19.3", - "six>=1.15.0", "Shapely>=1.7.1", 'PyYAML'], + "six>=1.15.0", "Shapely>=1.7.1", 'PyYAML', 'Pillow'], package_dir={'': MODULE_NAME}, packages=setuptools.find_namespace_packages(where=MODULE_NAME), package_data={'': ['*.onnx', '*.yaml']}, diff --git a/python/setup_openvino.py b/python/setup_openvino.py index caeb2bc..a0ae901 100644 --- a/python/setup_openvino.py +++ b/python/setup_openvino.py @@ -71,7 +71,7 @@ setuptools.setup( include_package_data=True, install_requires=["pyclipper>=1.2.1", "openvino>=2022.2.0", "opencv_python>=4.5.1.48", "numpy>=1.19.3", - "six>=1.15.0", "Shapely>=1.7.1", 'PyYAML'], + "six>=1.15.0", "Shapely>=1.7.1", 'PyYAML', 'Pillow'], package_dir={'': MODULE_NAME}, packages=setuptools.find_namespace_packages(where=MODULE_NAME), package_data={'': ['*.onnx', '*.yaml', '*.txt']},