Optim Load_image class

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):

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):

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']},

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']},