diff --git a/python/rapidocr_onnxruntime/main.py b/python/rapidocr_onnxruntime/main.py index 410639f..ecd9344 100644 --- a/python/rapidocr_onnxruntime/main.py +++ b/python/rapidocr_onnxruntime/main.py @@ -15,6 +15,7 @@ from .utils import ( LoadImage, UpdateParameters, VisRes, + get_logger, init_args, read_yaml, update_model_path, @@ -22,6 +23,7 @@ from .utils import ( root_dir = Path(__file__).resolve().parent DEFAULT_CFG_PATH = root_dir / "config.yaml" +logger = get_logger("RapidOCR") class RapidOCR: @@ -246,10 +248,10 @@ def main(): use_cls=use_cls, use_rec=use_rec, ) - print(result) + logger.info(result) if args.print_cost: - print(elapse_list) + logger.info(elapse_list) if args.vis_res: vis = VisRes() @@ -260,7 +262,7 @@ def main(): boxes, *_ = list(zip(*result)) vis_img = vis(args.img_path, boxes) cv2.imwrite(str(save_path), vis_img) - print(f"The vis result has saved in {save_path}") + logger.info("The vis result has saved in %s", save_path) elif use_det and use_rec: font_path = Path(args.vis_font_path) @@ -270,7 +272,7 @@ def main(): boxes, txts, scores = list(zip(*result)) vis_img = vis(args.img_path, boxes, txts, scores, font_path=font_path) cv2.imwrite(str(save_path), vis_img) - print(f"The vis result has saved in {save_path}") + logger.info("The vis result has saved in %s", save_path) if __name__ == "__main__": diff --git a/python/rapidocr_onnxruntime/utils.py b/python/rapidocr_onnxruntime/utils.py deleted file mode 100644 index 5ca74a6..0000000 --- a/python/rapidocr_onnxruntime/utils.py +++ /dev/null @@ -1,688 +0,0 @@ -# -*- encoding: utf-8 -*- -# @Author: SWHL -# @Contact: liekkaskono@163.com -import argparse -import logging -import math -import os -import platform -import random -import traceback -import warnings -from io import BytesIO -from pathlib import Path -from typing import Dict, List, Optional, Tuple, Union - -import cv2 -import numpy as np -import yaml -from onnxruntime import ( - GraphOptimizationLevel, - InferenceSession, - SessionOptions, - get_available_providers, - get_device, -) -from PIL import Image, ImageDraw, ImageFont, UnidentifiedImageError - -root_dir = Path(__file__).resolve().parent -InputType = Union[str, np.ndarray, bytes, Path, Image.Image] - -CPU_EP = "CPUExecutionProvider" -CUDA_EP = "CUDAExecutionProvider" -DIRECTML_EP = "DmlExecutionProvider" - - -class OrtInferSession: - def __init__(self, config): - self.logger = get_logger("OrtInferSession") - - sess_opt = SessionOptions() - sess_opt.log_severity_level = 4 - sess_opt.enable_cpu_mem_arena = False - sess_opt.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL - - cpu_nums = os.cpu_count() - intra_op_num_threads = config.get("intra_op_num_threads", -1) - if intra_op_num_threads != -1 and 1 <= intra_op_num_threads <= cpu_nums: - sess_opt.intra_op_num_threads = intra_op_num_threads - - inter_op_num_threads = config.get("inter_op_num_threads", -1) - if inter_op_num_threads != -1 and 1 <= inter_op_num_threads <= cpu_nums: - sess_opt.inter_op_num_threads = inter_op_num_threads - - model_path = config.get("model_path", None) - self._verify_model(model_path) - - self.cfg_use_cuda = config.get("use_cuda", None) - self.cfg_use_dml = config.get("use_dml", None) - EP_list = self._get_ep_list() - - self.session = InferenceSession( - model_path, sess_options=sess_opt, providers=EP_list - ) - - self._verify_providers() - - def _get_ep_list(self) -> List[Tuple[str, str]]: - had_providers: List[str] = get_available_providers() - - cpu_provider_opts = { - "arena_extend_strategy": "kSameAsRequested", - } - EP_list = [(CPU_EP, cpu_provider_opts)] - - cuda_provider_opts = { - "device_id": 0, - "arena_extend_strategy": "kNextPowerOfTwo", - "cudnn_conv_algo_search": "EXHAUSTIVE", - "do_copy_in_default_stream": True, - } - self.use_cuda = self._check_cuda_condition(had_providers) - if self.use_cuda: - EP_list.insert(0, (CUDA_EP, cuda_provider_opts)) - - self.use_directml = self._check_dml_condition(had_providers) - if self.use_directml: - print( - "Windows 10 or above detected, try to use DirectML as primary provider" - ) - directml_options = ( - cuda_provider_opts if self.use_cuda else cpu_provider_opts - ) - EP_list.insert(0, (DIRECTML_EP, directml_options)) - return EP_list - - def _check_cuda_condition(self, had_providers: List[str]) -> bool: - if not self.cfg_use_cuda: - return False - - cur_device = get_device() - if cur_device != "GPU" or CUDA_EP not in had_providers: - self.logger.warning( - "%s is not in available providers (%s)", CUDA_EP, had_providers - ) - self.logger.info("If you want to use GPU acceleration, you must do:") - self.logger.info( - "First, uninstall all onnxruntime pakcages in current environment." - ) - self.logger.info( - "Second, install onnxruntime-gpu by `pip install onnxruntime-gpu`." - ) - self.logger.info( - "\tNote the onnxruntime-gpu version must match your cuda and cudnn version." - ) - self.logger.info( - "\tYou can refer this link: https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html" - ) - self.logger.info( - "Third, ensure %s is in available providers list. e.g. ['CUDAExecutionProvider', 'CPUExecutionProvider']", - CUDA_EP, - ) - return False - - return True - - def _check_dml_condition(self, had_providers: List[str]) -> bool: - if not self.cfg_use_dml: - return False - - cur_os = platform.system() - if cur_os != "Windows": - self.logger.warning( - "DirectML is only supported in Windows OS. The current OS is %s", cur_os - ) - return False - - cur_window_version = int(platform.release().split(".")[0]) - if cur_window_version < 10: - self.logger.warning( - "DirectML is only supported in Windows 10 and above OS. The current Windows version is %s", - cur_window_version, - ) - return False - - if DIRECTML_EP not in had_providers: - self.logger.warning( - "%s is not in available providers (%s). ", DIRECTML_EP, had_providers - ) - self.logger.info("If you want to use DirectML acceleration, you must do:") - self.logger.info( - "First, uninstall all onnxruntime pakcages in current environment." - ) - self.logger.info( - "Second, install onnxruntime-directml by `pip install onnxruntime-directml`" - ) - self.logger.info( - "Third, ensure %s is in available providers list. e.g. ['DmlExecutionProvider', 'CPUExecutionProvider']", - DIRECTML_EP, - ) - return False - - return True - - def _verify_providers(self) -> None: - session_providers = self.session.get_providers() - first_provider = session_providers[0] - - if self.use_cuda and first_provider != CUDA_EP: - self.logger.warning( - "%s is not avaiable for current env, the inference part is automatically shifted to be executed under %s.", - CUDA_EP, - first_provider, - ) - - if self.use_directml and first_provider != DIRECTML_EP: - self.logger.warning( - "%s is not available for the current environment, the inference part is automatically shifted to be executed under %s.", - DIRECTML_EP, - first_provider, - ) - - def __call__(self, input_content: np.ndarray) -> np.ndarray: - input_dict = dict(zip(self.get_input_names(), [input_content])) - try: - return self.session.run(self.get_output_names(), input_dict) - except Exception as e: - error_info = traceback.format_exc() - raise ONNXRuntimeError(error_info) from e - - def get_input_names( - self, - ): - return [v.name for v in self.session.get_inputs()] - - def get_output_names( - self, - ): - return [v.name for v in self.session.get_outputs()] - - def get_character_list(self, key: str = "character"): - meta_dict = self.session.get_modelmeta().custom_metadata_map - return meta_dict[key].splitlines() - - def have_key(self, key: str = "character") -> bool: - meta_dict = self.session.get_modelmeta().custom_metadata_map - if key in meta_dict.keys(): - return True - return False - - @staticmethod - def _verify_model(model_path: Union[str, Path, None]): - if model_path is None: - raise ValueError("model_path is None!") - - model_path = Path(model_path) - - if not model_path.exists(): - raise FileNotFoundError(f"{model_path} does not exists.") - - if not model_path.is_file(): - raise FileExistsError(f"{model_path} is not a file.") - - -class ONNXRuntimeError(Exception): - pass - - -class LoadImage: - def __init__( - self, - ): - pass - - def __call__(self, img: InputType) -> np.ndarray: - if not isinstance(img, InputType.__args__): - raise LoadImageError( - f"The img type {type(img)} does not in {InputType.__args__}" - ) - - origin_img_type = type(img) - img = self.load_img(img) - img = self.convert_img(img, origin_img_type) - return img - - def load_img(self, img: InputType) -> np.ndarray: - if isinstance(img, (str, Path)): - self.verify_exist(img) - try: - img = self.img_to_ndarray(Image.open(img)) - except UnidentifiedImageError as e: - raise LoadImageError(f"cannot identify image file {img}") from e - return img - - if isinstance(img, bytes): - img = self.img_to_ndarray(Image.open(BytesIO(img))) - return img - - if isinstance(img, np.ndarray): - return img - - if isinstance(img, Image.Image): - return self.img_to_ndarray(img) - - raise LoadImageError(f"{type(img)} is not supported!") - - def img_to_ndarray(self, img: Image.Image) -> np.ndarray: - if img.mode == "1": - img = img.convert("L") - return np.array(img) - return np.array(img) - - def convert_img(self, img: np.ndarray, origin_img_type): - if img.ndim == 2: - return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) - - if img.ndim == 3: - channel = img.shape[2] - if channel == 1: - return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) - - if channel == 2: - return self.cvt_two_to_three(img) - - if channel == 3: - if issubclass(origin_img_type, (str, Path, bytes, Image.Image)): - return cv2.cvtColor(img, cv2.COLOR_RGB2BGR) - return img - - if channel == 4: - return self.cvt_four_to_three(img) - - raise LoadImageError( - f"The channel({channel}) of the img is not in [1, 2, 3, 4]" - ) - - raise LoadImageError(f"The ndim({img.ndim}) of the img is not in [2, 3]") - - @staticmethod - def cvt_two_to_three(img: np.ndarray) -> np.ndarray: - """gray + alpha → BGR""" - img_gray = img[..., 0] - img_bgr = cv2.cvtColor(img_gray, cv2.COLOR_GRAY2BGR) - - img_alpha = img[..., 1] - not_a = cv2.bitwise_not(img_alpha) - not_a = cv2.cvtColor(not_a, cv2.COLOR_GRAY2BGR) - - new_img = cv2.bitwise_and(img_bgr, img_bgr, mask=img_alpha) - new_img = cv2.add(new_img, not_a) - return new_img - - @staticmethod - def cvt_four_to_three(img: np.ndarray) -> np.ndarray: - """RGBA → BGR""" - r, g, b, a = cv2.split(img) - new_img = cv2.merge((b, g, r)) - - not_a = cv2.bitwise_not(a) - not_a = cv2.cvtColor(not_a, cv2.COLOR_GRAY2BGR) - - new_img = cv2.bitwise_and(new_img, new_img, mask=a) - - mean_color = np.mean(new_img) - if mean_color <= 0.0: - new_img = cv2.add(new_img, not_a) - else: - 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 - - -def read_yaml(yaml_path): - with open(yaml_path, "rb") as f: - data = yaml.load(f, Loader=yaml.Loader) - return data - - -def update_model_path(config): - key = "model_path" - config["Det"][key] = str(root_dir / config["Det"][key]) - config["Rec"][key] = str(root_dir / config["Rec"][key]) - config["Cls"][key] = str(root_dir / config["Cls"][key]) - return config - - -def init_args(): - parser = argparse.ArgumentParser() - parser.add_argument("-img", "--img_path", type=str, default=None, required=True) - parser.add_argument("-p", "--print_cost", action="store_true", default=False) - - global_group = parser.add_argument_group(title="Global") - global_group.add_argument("--text_score", type=float, default=0.5) - - global_group.add_argument("--no_det", action="store_true", default=False) - global_group.add_argument("--no_cls", action="store_true", default=False) - global_group.add_argument("--no_rec", action="store_true", default=False) - - global_group.add_argument("--print_verbose", action="store_true", default=False) - global_group.add_argument("--min_height", type=int, default=30) - global_group.add_argument("--width_height_ratio", type=int, default=8) - - global_group.add_argument("--intra_op_num_threads", type=int, default=-1) - global_group.add_argument("--inter_op_num_threads", type=int, default=-1) - - det_group = parser.add_argument_group(title="Det") - det_group.add_argument( - "--det_use_cuda", - action="store_true", - default=False, - help="Whether to use cuda. The prerequisite is: first uninstall all onnxruntime packages and install only the onnxruntime-gpu library.", - ) - det_group.add_argument( - "--det_use_dml", - action="store_true", - default=False, - help="Whether to use DirectML. The prerequisite is: the operating system is Windows 10+. First uninstall all onnxruntime packages and install only the onnxruntime-directml library.", - ) - det_group.add_argument("--det_model_path", type=str, default=None) - det_group.add_argument("--det_limit_side_len", type=float, default=736) - det_group.add_argument( - "--det_limit_type", type=str, default="min", choices=["max", "min"] - ) - det_group.add_argument("--det_thresh", type=float, default=0.3) - det_group.add_argument("--det_box_thresh", type=float, default=0.5) - det_group.add_argument("--det_unclip_ratio", type=float, default=1.6) - det_group.add_argument( - "--det_donot_use_dilation", action="store_true", default=False - ) - det_group.add_argument( - "--det_score_mode", type=str, default="fast", choices=["slow", "fast"] - ) - - cls_group = parser.add_argument_group(title="Cls") - cls_group.add_argument( - "--cls_use_cuda", - action="store_true", - default=False, - help="Whether to use cuda. The prerequisite is: first uninstall all onnxruntime packages and install only the onnxruntime-gpu library.", - ) - cls_group.add_argument( - "--cls_use_dml", - action="store_true", - default=False, - help="Whether to use DirectML. The prerequisite is: the operating system is Windows 10+. First uninstall all onnxruntime packages and install only the onnxruntime-directml library.", - ) - cls_group.add_argument("--cls_model_path", type=str, default=None) - cls_group.add_argument("--cls_image_shape", type=list, default=[3, 48, 192]) - cls_group.add_argument("--cls_label_list", type=list, default=["0", "180"]) - cls_group.add_argument("--cls_batch_num", type=int, default=6) - cls_group.add_argument("--cls_thresh", type=float, default=0.9) - - rec_group = parser.add_argument_group(title="Rec") - rec_group.add_argument( - "--rec_use_cuda", - action="store_true", - default=False, - help="Whether to use cuda. The prerequisite is: first uninstall all onnxruntime packages and install only the onnxruntime-gpu library.", - ) - rec_group.add_argument( - "--rec_use_dml", - action="store_true", - default=False, - help="Whether to use DirectML. The prerequisite is: the operating system is Windows 10+. First uninstall all onnxruntime packages and install only the onnxruntime-directml library.", - ) - rec_group.add_argument("--rec_model_path", type=str, default=None) - rec_group.add_argument("--rec_keys_path", type=str, default=None) - rec_group.add_argument("--rec_img_shape", type=list, default=[3, 48, 320]) - rec_group.add_argument("--rec_batch_num", type=int, default=6) - - vis_group = parser.add_argument_group(title="Visual Result") - vis_group.add_argument("-vis", "--vis_res", action="store_true", default=False) - vis_group.add_argument( - "--vis_font_path", - type=str, - default=None, - help="When -vis is True, the font_path must have value.", - ) - vis_group.add_argument( - "--vis_save_path", - type=str, - default=".", - help="The directory of saving the vis image.", - ) - - args = parser.parse_args() - return args - - -class UpdateParameters: - def __init__(self) -> None: - pass - - def parse_kwargs(self, **kwargs): - global_dict, det_dict, cls_dict, rec_dict = {}, {}, {}, {} - for k, v in kwargs.items(): - if k.startswith("det"): - k = k.split("det_")[1] - if k == "donot_use_dilation": - k = "use_dilation" - v = not v - - det_dict[k] = v - elif k.startswith("cls"): - cls_dict[k] = v - elif k.startswith("rec"): - rec_dict[k] = v - else: - global_dict[k] = v - return global_dict, det_dict, cls_dict, rec_dict - - def __call__(self, config, **kwargs): - global_dict, det_dict, cls_dict, rec_dict = self.parse_kwargs(**kwargs) - new_config = { - "Global": self.update_global_params(config["Global"], global_dict), - "Det": self.update_params(config["Det"], det_dict, "det_", None), - "Cls": self.update_params( - config["Cls"], - cls_dict, - "cls_", - ["cls_label_list", "cls_model_path", "cls_use_cuda"], - ), - "Rec": self.update_params( - config["Rec"], rec_dict, "rec_", ["rec_model_path", "rec_use_cuda"] - ), - } - - update_params = ["intra_op_num_threads", "inter_op_num_threads"] - new_config = self.update_global_to_module( - config, update_params, src="Global", dsts=["Det", "Cls", "Rec"] - ) - return new_config - - def update_global_to_module( - self, config, params: List[str], src: str, dsts: List[str] - ): - for dst in dsts: - for param in params: - config[dst].update({param: config[src][param]}) - return config - - def update_global_params(self, config, global_dict): - if global_dict: - config.update(global_dict) - return config - - def update_params( - self, - config, - param_dict: Dict[str, str], - prefix: str, - need_remove_prefix: Optional[List[str]] = None, - ): - if not param_dict: - return config - - filter_dict = self.remove_prefix(param_dict, prefix, need_remove_prefix) - model_path = filter_dict.get("model_path", None) - if not model_path: - filter_dict["model_path"] = str(root_dir / config["model_path"]) - - config.update(filter_dict) - return config - - @staticmethod - def remove_prefix( - config: Dict[str, str], - prefix: str, - need_remove_prefix: Optional[List[str]] = None, - ) -> Dict[str, str]: - if not need_remove_prefix: - return config - - new_rec_dict = {} - for k, v in config.items(): - if k in need_remove_prefix: - k = k.split(prefix)[1] - new_rec_dict[k] = v - return new_rec_dict - - -class VisRes: - def __init__(self, text_score: float = 0.5): - self.text_score = text_score - self.load_img = LoadImage() - - def __call__( - self, - img_content: InputType, - dt_boxes: np.ndarray, - txts: Optional[Union[List[str], Tuple[str]]] = None, - scores: Optional[Tuple[float]] = None, - font_path: Optional[str] = None, - ) -> np.ndarray: - if txts is None and scores is None: - return self.draw_dt_boxes(img_content, dt_boxes) - return self.draw_ocr_box_txt(img_content, dt_boxes, txts, scores, font_path) - - def draw_dt_boxes(self, img_content: InputType, dt_boxes: np.ndarray) -> np.ndarray: - img = self.load_img(img_content) - - for idx, box in enumerate(dt_boxes): - color = self.get_random_color() - - points = np.array(box) - cv2.polylines(img, np.int32([points]), 1, color=color, thickness=1) - - start_point = round(points[0][0]), round(points[0][1]) - cv2.putText( - img, f"{idx}", start_point, cv2.FONT_HERSHEY_SIMPLEX, 1, color, 3 - ) - return img - - def draw_ocr_box_txt( - self, - img_content: InputType, - dt_boxes: np.ndarray, - txts: Optional[Union[List[str], Tuple[str]]] = None, - scores: Optional[Tuple[float]] = None, - font_path: Optional[str] = None, - ) -> np.ndarray: - font_path = self.get_font_path(font_path) - - image = Image.fromarray(self.load_img(img_content)) - h, w = image.height, image.width - if image.mode == "L": - image = image.convert("RGB") - - img_left = image.copy() - img_right = Image.new("RGB", (w, h), (255, 255, 255)) - - random.seed(0) - draw_left = ImageDraw.Draw(img_left) - draw_right = ImageDraw.Draw(img_right) - for idx, (box, txt) in enumerate(zip(dt_boxes, txts)): - if scores is not None and float(scores[idx]) < self.text_score: - continue - - color = self.get_random_color() - - box_list = np.array(box).reshape(8).tolist() - draw_left.polygon(box_list, fill=color) - draw_right.polygon(box_list, outline=color) - - box_height = self.get_box_height(box) - box_width = self.get_box_width(box) - if box_height > 2 * box_width: - font_size = max(int(box_width * 0.9), 10) - font = ImageFont.truetype(font_path, font_size, encoding="utf-8") - cur_y = box[0][1] - - for c in txt: - draw_right.text( - (box[0][0] + 3, cur_y), c, fill=(0, 0, 0), font=font - ) - cur_y += self.get_char_size(font, c) - else: - font_size = max(int(box_height * 0.8), 10) - font = ImageFont.truetype(font_path, font_size, encoding="utf-8") - draw_right.text([box[0][0], box[0][1]], txt, fill=(0, 0, 0), font=font) - - img_left = Image.blend(image, img_left, 0.5) - img_show = Image.new("RGB", (w * 2, h), (255, 255, 255)) - img_show.paste(img_left, (0, 0, w, h)) - img_show.paste(img_right, (w, 0, w * 2, h)) - return np.array(img_show) - - @staticmethod - def get_font_path(font_path: Optional[Union[str, Path]] = None) -> str: - if font_path is None or not Path(font_path).exists(): - raise FileNotFoundError( - f"The {font_path} does not exists! \n" - f"You could download the file in the https://drive.google.com/file/d/1evWVX38EFNwTq_n5gTFgnlv8tdaNcyIA/view?usp=sharing" - ) - return str(font_path) - - @staticmethod - def get_random_color() -> Tuple[int, int, int]: - return ( - random.randint(0, 255), - random.randint(0, 255), - random.randint(0, 255), - ) - - @staticmethod - def get_box_height(box: List[List[float]]) -> float: - return math.sqrt((box[0][0] - box[3][0]) ** 2 + (box[0][1] - box[3][1]) ** 2) - - @staticmethod - def get_box_width(box: List[List[float]]) -> float: - return math.sqrt((box[0][0] - box[1][0]) ** 2 + (box[0][1] - box[1][1]) ** 2) - - @staticmethod - def get_char_size(font, char_str: str) -> float: - # compatible with Pillow v9 and v10. - if hasattr(font, "getsize"): - get_size_func = getattr(font, "getsize") - return get_size_func(char_str)[1] - - if hasattr(font, "getlength"): - get_size_func = getattr(font, "getlength") - return get_size_func(char_str) - - raise ValueError( - "The Pillow ImageFont instance has not getsize or getlength func." - ) - - -def get_logger(name: str): - logger = logging.getLogger(name) - logger.setLevel(logging.DEBUG) - - fmt = "%(asctime)s - %(name)s - %(levelname)s: %(message)s" - format_str = logging.Formatter(fmt) - - sh = logging.StreamHandler() - sh.setLevel(logging.DEBUG) - - logger.addHandler(sh) - sh.setFormatter(format_str) - return logger diff --git a/python/rapidocr_onnxruntime/utils/__init__.py b/python/rapidocr_onnxruntime/utils/__init__.py new file mode 100644 index 0000000..cf5ca69 --- /dev/null +++ b/python/rapidocr_onnxruntime/utils/__init__.py @@ -0,0 +1,19 @@ +# -*- encoding: utf-8 -*- +# @Author: SWHL +# @Contact: liekkaskono@163.com +from pathlib import Path +from typing import Dict, Union + +import yaml + +from .infer_engine import OrtInferSession +from .load_image import LoadImage, LoadImageError +from .logger import get_logger +from .parse_parameters import UpdateParameters, init_args, update_model_path +from .vis_res import VisRes + + +def read_yaml(yaml_path: Union[str, Path]) -> Dict[str, Dict]: + with open(yaml_path, "rb") as f: + data = yaml.load(f, Loader=yaml.Loader) + return data diff --git a/python/rapidocr_onnxruntime/utils/infer_engine.py b/python/rapidocr_onnxruntime/utils/infer_engine.py new file mode 100644 index 0000000..13b3e14 --- /dev/null +++ b/python/rapidocr_onnxruntime/utils/infer_engine.py @@ -0,0 +1,220 @@ +# -*- encoding: utf-8 -*- +# @Author: SWHL +# @Contact: liekkaskono@163.com +import os +import platform +import traceback +from pathlib import Path +from typing import List, Tuple, Union + +import numpy as np +from onnxruntime import ( + GraphOptimizationLevel, + InferenceSession, + SessionOptions, + get_available_providers, + get_device, +) + +from .logger import get_logger + +CPU_EP = "CPUExecutionProvider" +CUDA_EP = "CUDAExecutionProvider" +DIRECTML_EP = "DmlExecutionProvider" + + +class OrtInferSession: + def __init__(self, config): + self.logger = get_logger("OrtInferSession") + + model_path = config.get("model_path", None) + self._verify_model(model_path) + + self.cfg_use_cuda = config.get("use_cuda", None) + self.cfg_use_dml = config.get("use_dml", None) + + self.had_providers: List[str] = get_available_providers() + EP_list = self._get_ep_list() + + sess_opt = self._init_sess_opts(config) + self.session = InferenceSession( + model_path, sess_options=sess_opt, providers=EP_list + ) + self._verify_providers() + + @staticmethod + def _init_sess_opts(config): + sess_opt = SessionOptions() + sess_opt.log_severity_level = 4 + sess_opt.enable_cpu_mem_arena = False + sess_opt.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL + + cpu_nums = os.cpu_count() + intra_op_num_threads = config.get("intra_op_num_threads", -1) + if intra_op_num_threads != -1 and 1 <= intra_op_num_threads <= cpu_nums: + sess_opt.intra_op_num_threads = intra_op_num_threads + + inter_op_num_threads = config.get("inter_op_num_threads", -1) + if inter_op_num_threads != -1 and 1 <= inter_op_num_threads <= cpu_nums: + sess_opt.inter_op_num_threads = inter_op_num_threads + + return sess_opt + + def _get_ep_list(self) -> List[Tuple[str, str]]: + cpu_provider_opts = { + "arena_extend_strategy": "kSameAsRequested", + } + EP_list = [(CPU_EP, cpu_provider_opts)] + + cuda_provider_opts = { + "device_id": 0, + "arena_extend_strategy": "kNextPowerOfTwo", + "cudnn_conv_algo_search": "EXHAUSTIVE", + "do_copy_in_default_stream": True, + } + self.use_cuda = self._check_cuda() + if self.use_cuda: + EP_list.insert(0, (CUDA_EP, cuda_provider_opts)) + + self.use_directml = self._check_dml() + if self.use_directml: + self.logger.info( + "Windows 10 or above detected, try to use DirectML as primary provider" + ) + directml_options = ( + cuda_provider_opts if self.use_cuda else cpu_provider_opts + ) + EP_list.insert(0, (DIRECTML_EP, directml_options)) + return EP_list + + def _check_cuda(self) -> bool: + if not self.cfg_use_cuda: + return False + + cur_device = get_device() + if cur_device == "GPU" and CUDA_EP in self.had_providers: + return True + + self.logger.warning( + "%s is not in available providers (%s), default use of %s inference.", + CUDA_EP, + self.had_providers, + self.had_providers[0], + ) + self.logger.info("If you want to use GPU acceleration, you must do:") + self.logger.info( + "First, uninstall all onnxruntime pakcages in current environment." + ) + self.logger.info( + "Second, install onnxruntime-gpu by `pip install onnxruntime-gpu`." + ) + self.logger.info( + "\tNote the onnxruntime-gpu version must match your cuda and cudnn version." + ) + self.logger.info( + "\tYou can refer this link: https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html" + ) + self.logger.info( + "Third, ensure %s is in available providers list. e.g. ['CUDAExecutionProvider', 'CPUExecutionProvider']", + CUDA_EP, + ) + return False + + def _check_dml(self) -> bool: + if not self.cfg_use_dml: + return False + + cur_os = platform.system() + if cur_os != "Windows": + self.logger.warning( + "DirectML is only supported in Windows OS. The current OS is %s", cur_os + ) + return False + + cur_window_version = int(platform.release().split(".")[0]) + if cur_window_version < 10: + self.logger.warning( + "DirectML is only supported in Windows 10 and above OS. The current Windows version is %s", + cur_window_version, + ) + return False + + if DIRECTML_EP in self.had_providers: + return True + + self.logger.warning( + "%s is not in available providers (%s), default use of %s inference.", + DIRECTML_EP, + self.had_providers, + self.had_providers[0], + ) + self.logger.info("If you want to use DirectML acceleration, you must do:") + self.logger.info( + "First, uninstall all onnxruntime pakcages in current environment." + ) + self.logger.info( + "Second, install onnxruntime-directml by `pip install onnxruntime-directml`" + ) + self.logger.info( + "Third, ensure %s is in available providers list. e.g. ['DmlExecutionProvider', 'CPUExecutionProvider']", + DIRECTML_EP, + ) + return False + + def _verify_providers(self) -> None: + session_providers = self.session.get_providers() + first_provider = session_providers[0] + + if self.use_cuda and first_provider != CUDA_EP: + self.logger.warning( + "%s is not avaiable for current env, the inference part is automatically shifted to be executed under %s.", + CUDA_EP, + first_provider, + ) + + if self.use_directml and first_provider != DIRECTML_EP: + self.logger.warning( + "%s is not available for current env, the inference part is automatically shifted to be executed under %s.", + DIRECTML_EP, + first_provider, + ) + + def __call__(self, input_content: np.ndarray) -> np.ndarray: + input_dict = dict(zip(self.get_input_names(), [input_content])) + try: + return self.session.run(self.get_output_names(), input_dict) + except Exception as e: + error_info = traceback.format_exc() + raise ONNXRuntimeError(error_info) from e + + def get_input_names(self): + return [v.name for v in self.session.get_inputs()] + + def get_output_names(self): + return [v.name for v in self.session.get_outputs()] + + def get_character_list(self, key: str = "character"): + meta_dict = self.session.get_modelmeta().custom_metadata_map + return meta_dict[key].splitlines() + + def have_key(self, key: str = "character") -> bool: + meta_dict = self.session.get_modelmeta().custom_metadata_map + if key in meta_dict.keys(): + return True + return False + + @staticmethod + def _verify_model(model_path: Union[str, Path, None]): + if model_path is None: + raise ValueError("model_path is None!") + + model_path = Path(model_path) + if not model_path.exists(): + raise FileNotFoundError(f"{model_path} does not exists.") + + if not model_path.is_file(): + raise FileExistsError(f"{model_path} is not a file.") + + +class ONNXRuntimeError(Exception): + pass diff --git a/python/rapidocr_onnxruntime/utils/load_image.py b/python/rapidocr_onnxruntime/utils/load_image.py new file mode 100644 index 0000000..0566053 --- /dev/null +++ b/python/rapidocr_onnxruntime/utils/load_image.py @@ -0,0 +1,123 @@ +# -*- encoding: utf-8 -*- +# @Author: SWHL +# @Contact: liekkaskono@163.com +from io import BytesIO +from pathlib import Path +from typing import Union + +import cv2 +import numpy as np +from PIL import Image, UnidentifiedImageError + +root_dir = Path(__file__).resolve().parent +InputType = Union[str, np.ndarray, bytes, Path, Image.Image] + + +class LoadImage: + def __init__(self): + pass + + def __call__(self, img: InputType) -> np.ndarray: + if not isinstance(img, InputType.__args__): + raise LoadImageError( + f"The img type {type(img)} does not in {InputType.__args__}" + ) + + origin_img_type = type(img) + img = self.load_img(img) + img = self.convert_img(img, origin_img_type) + return img + + def load_img(self, img: InputType) -> np.ndarray: + if isinstance(img, (str, Path)): + self.verify_exist(img) + try: + img = self.img_to_ndarray(Image.open(img)) + except UnidentifiedImageError as e: + raise LoadImageError(f"cannot identify image file {img}") from e + return img + + if isinstance(img, bytes): + img = self.img_to_ndarray(Image.open(BytesIO(img))) + return img + + if isinstance(img, np.ndarray): + return img + + if isinstance(img, Image.Image): + return self.img_to_ndarray(img) + + raise LoadImageError(f"{type(img)} is not supported!") + + def img_to_ndarray(self, img: Image.Image) -> np.ndarray: + if img.mode == "1": + img = img.convert("L") + return np.array(img) + return np.array(img) + + def convert_img(self, img: np.ndarray, origin_img_type): + if img.ndim == 2: + return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) + + if img.ndim == 3: + channel = img.shape[2] + if channel == 1: + return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) + + if channel == 2: + return self.cvt_two_to_three(img) + + if channel == 3: + if issubclass(origin_img_type, (str, Path, bytes, Image.Image)): + return cv2.cvtColor(img, cv2.COLOR_RGB2BGR) + return img + + if channel == 4: + return self.cvt_four_to_three(img) + + raise LoadImageError( + f"The channel({channel}) of the img is not in [1, 2, 3, 4]" + ) + + raise LoadImageError(f"The ndim({img.ndim}) of the img is not in [2, 3]") + + @staticmethod + def cvt_two_to_three(img: np.ndarray) -> np.ndarray: + """gray + alpha → BGR""" + img_gray = img[..., 0] + img_bgr = cv2.cvtColor(img_gray, cv2.COLOR_GRAY2BGR) + + img_alpha = img[..., 1] + not_a = cv2.bitwise_not(img_alpha) + not_a = cv2.cvtColor(not_a, cv2.COLOR_GRAY2BGR) + + new_img = cv2.bitwise_and(img_bgr, img_bgr, mask=img_alpha) + new_img = cv2.add(new_img, not_a) + return new_img + + @staticmethod + def cvt_four_to_three(img: np.ndarray) -> np.ndarray: + """RGBA → BGR""" + r, g, b, a = cv2.split(img) + new_img = cv2.merge((b, g, r)) + + not_a = cv2.bitwise_not(a) + not_a = cv2.cvtColor(not_a, cv2.COLOR_GRAY2BGR) + + new_img = cv2.bitwise_and(new_img, new_img, mask=a) + + mean_color = np.mean(new_img) + if mean_color <= 0.0: + new_img = cv2.add(new_img, not_a) + else: + 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/rapidocr_onnxruntime/utils/logger.py b/python/rapidocr_onnxruntime/utils/logger.py new file mode 100644 index 0000000..ffd1cd0 --- /dev/null +++ b/python/rapidocr_onnxruntime/utils/logger.py @@ -0,0 +1,19 @@ +# -*- encoding: utf-8 -*- +# @Author: SWHL +# @Contact: liekkaskono@163.com +import logging + + +def get_logger(name: str): + logger = logging.getLogger(name) + logger.setLevel(logging.DEBUG) + + fmt = "%(asctime)s - %(name)s - %(levelname)s: %(message)s" + format_str = logging.Formatter(fmt) + + sh = logging.StreamHandler() + sh.setLevel(logging.DEBUG) + + logger.addHandler(sh) + sh.setFormatter(format_str) + return logger diff --git a/python/rapidocr_onnxruntime/utils/parse_parameters.py b/python/rapidocr_onnxruntime/utils/parse_parameters.py new file mode 100644 index 0000000..395866c --- /dev/null +++ b/python/rapidocr_onnxruntime/utils/parse_parameters.py @@ -0,0 +1,215 @@ +# -*- encoding: utf-8 -*- +# @Author: SWHL +# @Contact: liekkaskono@163.com +import argparse +from pathlib import Path +from typing import Dict, List, Optional, Union + +import numpy as np +from PIL import Image + +root_dir = Path(__file__).resolve().parent.parent +InputType = Union[str, np.ndarray, bytes, Path, Image.Image] + + +def update_model_path(config): + key = "model_path" + config["Det"][key] = str(root_dir / config["Det"][key]) + config["Rec"][key] = str(root_dir / config["Rec"][key]) + config["Cls"][key] = str(root_dir / config["Cls"][key]) + return config + + +def init_args(): + parser = argparse.ArgumentParser() + parser.add_argument("-img", "--img_path", type=str, default=None, required=True) + parser.add_argument("-p", "--print_cost", action="store_true", default=False) + + global_group = parser.add_argument_group(title="Global") + global_group.add_argument("--text_score", type=float, default=0.5) + + global_group.add_argument("--no_det", action="store_true", default=False) + global_group.add_argument("--no_cls", action="store_true", default=False) + global_group.add_argument("--no_rec", action="store_true", default=False) + + global_group.add_argument("--print_verbose", action="store_true", default=False) + global_group.add_argument("--min_height", type=int, default=30) + global_group.add_argument("--width_height_ratio", type=int, default=8) + + global_group.add_argument("--intra_op_num_threads", type=int, default=-1) + global_group.add_argument("--inter_op_num_threads", type=int, default=-1) + + det_group = parser.add_argument_group(title="Det") + det_group.add_argument( + "--det_use_cuda", + action="store_true", + default=False, + help="Whether to use cuda. The prerequisite is: first uninstall all onnxruntime packages and install only the onnxruntime-gpu library.", + ) + det_group.add_argument( + "--det_use_dml", + action="store_true", + default=False, + help="Whether to use DirectML. The prerequisite is: the operating system is Windows 10+. First uninstall all onnxruntime packages and install only the onnxruntime-directml library.", + ) + det_group.add_argument("--det_model_path", type=str, default=None) + det_group.add_argument("--det_limit_side_len", type=float, default=736) + det_group.add_argument( + "--det_limit_type", type=str, default="min", choices=["max", "min"] + ) + det_group.add_argument("--det_thresh", type=float, default=0.3) + det_group.add_argument("--det_box_thresh", type=float, default=0.5) + det_group.add_argument("--det_unclip_ratio", type=float, default=1.6) + det_group.add_argument( + "--det_donot_use_dilation", action="store_true", default=False + ) + det_group.add_argument( + "--det_score_mode", type=str, default="fast", choices=["slow", "fast"] + ) + + cls_group = parser.add_argument_group(title="Cls") + cls_group.add_argument( + "--cls_use_cuda", + action="store_true", + default=False, + help="Whether to use cuda. The prerequisite is: first uninstall all onnxruntime packages and install only the onnxruntime-gpu library.", + ) + cls_group.add_argument( + "--cls_use_dml", + action="store_true", + default=False, + help="Whether to use DirectML. The prerequisite is: the operating system is Windows 10+. First uninstall all onnxruntime packages and install only the onnxruntime-directml library.", + ) + cls_group.add_argument("--cls_model_path", type=str, default=None) + cls_group.add_argument("--cls_image_shape", type=list, default=[3, 48, 192]) + cls_group.add_argument("--cls_label_list", type=list, default=["0", "180"]) + cls_group.add_argument("--cls_batch_num", type=int, default=6) + cls_group.add_argument("--cls_thresh", type=float, default=0.9) + + rec_group = parser.add_argument_group(title="Rec") + rec_group.add_argument( + "--rec_use_cuda", + action="store_true", + default=False, + help="Whether to use cuda. The prerequisite is: first uninstall all onnxruntime packages and install only the onnxruntime-gpu library.", + ) + rec_group.add_argument( + "--rec_use_dml", + action="store_true", + default=False, + help="Whether to use DirectML. The prerequisite is: the operating system is Windows 10+. First uninstall all onnxruntime packages and install only the onnxruntime-directml library.", + ) + rec_group.add_argument("--rec_model_path", type=str, default=None) + rec_group.add_argument("--rec_keys_path", type=str, default=None) + rec_group.add_argument("--rec_img_shape", type=list, default=[3, 48, 320]) + rec_group.add_argument("--rec_batch_num", type=int, default=6) + + vis_group = parser.add_argument_group(title="Visual Result") + vis_group.add_argument("-vis", "--vis_res", action="store_true", default=False) + vis_group.add_argument( + "--vis_font_path", + type=str, + default=None, + help="When -vis is True, the font_path must have value.", + ) + vis_group.add_argument( + "--vis_save_path", + type=str, + default=".", + help="The directory of saving the vis image.", + ) + + args = parser.parse_args() + return args + + +class UpdateParameters: + def __init__(self) -> None: + pass + + def parse_kwargs(self, **kwargs): + global_dict, det_dict, cls_dict, rec_dict = {}, {}, {}, {} + for k, v in kwargs.items(): + if k.startswith("det"): + k = k.split("det_")[1] + if k == "donot_use_dilation": + k = "use_dilation" + v = not v + + det_dict[k] = v + elif k.startswith("cls"): + cls_dict[k] = v + elif k.startswith("rec"): + rec_dict[k] = v + else: + global_dict[k] = v + return global_dict, det_dict, cls_dict, rec_dict + + def __call__(self, config, **kwargs): + global_dict, det_dict, cls_dict, rec_dict = self.parse_kwargs(**kwargs) + new_config = { + "Global": self.update_global_params(config["Global"], global_dict), + "Det": self.update_params(config["Det"], det_dict, "det_", None), + "Cls": self.update_params( + config["Cls"], + cls_dict, + "cls_", + ["cls_label_list", "cls_model_path", "cls_use_cuda"], + ), + "Rec": self.update_params( + config["Rec"], rec_dict, "rec_", ["rec_model_path", "rec_use_cuda"] + ), + } + + update_params = ["intra_op_num_threads", "inter_op_num_threads"] + new_config = self.update_global_to_module( + config, update_params, src="Global", dsts=["Det", "Cls", "Rec"] + ) + return new_config + + def update_global_to_module( + self, config, params: List[str], src: str, dsts: List[str] + ): + for dst in dsts: + for param in params: + config[dst].update({param: config[src][param]}) + return config + + def update_global_params(self, config, global_dict): + if global_dict: + config.update(global_dict) + return config + + def update_params( + self, + config, + param_dict: Dict[str, str], + prefix: str, + need_remove_prefix: Optional[List[str]] = None, + ): + if not param_dict: + return config + + filter_dict = self.remove_prefix(param_dict, prefix, need_remove_prefix) + model_path = filter_dict.get("model_path", None) + if not model_path: + filter_dict["model_path"] = str(root_dir / config["model_path"]) + + config.update(filter_dict) + return config + + @staticmethod + def remove_prefix( + config: Dict[str, str], + prefix: str, + need_remove_prefix: Optional[List[str]] = None, + ) -> Dict[str, str]: + if not need_remove_prefix: + return config + + new_rec_dict = {} + for k, v in config.items(): + if k in need_remove_prefix: + k = k.split(prefix)[1] + new_rec_dict[k] = v + return new_rec_dict diff --git a/python/rapidocr_onnxruntime/utils/vis_res.py b/python/rapidocr_onnxruntime/utils/vis_res.py new file mode 100644 index 0000000..405beab --- /dev/null +++ b/python/rapidocr_onnxruntime/utils/vis_res.py @@ -0,0 +1,143 @@ +# -*- encoding: utf-8 -*- +# @Author: SWHL +# @Contact: liekkaskono@163.com +import math +import random +from pathlib import Path +from typing import List, Optional, Tuple, Union + +import cv2 +import numpy as np +from PIL import Image, ImageDraw, ImageFont + +from .load_image import LoadImage + +root_dir = Path(__file__).resolve().parent +InputType = Union[str, np.ndarray, bytes, Path, Image.Image] + + +class VisRes: + def __init__(self, text_score: float = 0.5): + self.text_score = text_score + self.load_img = LoadImage() + + def __call__( + self, + img_content: InputType, + dt_boxes: np.ndarray, + txts: Optional[Union[List[str], Tuple[str]]] = None, + scores: Optional[Tuple[float]] = None, + font_path: Optional[str] = None, + ) -> np.ndarray: + if txts is None and scores is None: + return self.draw_dt_boxes(img_content, dt_boxes) + return self.draw_ocr_box_txt(img_content, dt_boxes, txts, scores, font_path) + + def draw_dt_boxes(self, img_content: InputType, dt_boxes: np.ndarray) -> np.ndarray: + img = self.load_img(img_content) + + for idx, box in enumerate(dt_boxes): + color = self.get_random_color() + + points = np.array(box) + cv2.polylines(img, np.int32([points]), 1, color=color, thickness=1) + + start_point = round(points[0][0]), round(points[0][1]) + cv2.putText( + img, f"{idx}", start_point, cv2.FONT_HERSHEY_SIMPLEX, 1, color, 3 + ) + return img + + def draw_ocr_box_txt( + self, + img_content: InputType, + dt_boxes: np.ndarray, + txts: Optional[Union[List[str], Tuple[str]]] = None, + scores: Optional[Tuple[float]] = None, + font_path: Optional[str] = None, + ) -> np.ndarray: + font_path = self.get_font_path(font_path) + + image = Image.fromarray(self.load_img(img_content)) + h, w = image.height, image.width + if image.mode == "L": + image = image.convert("RGB") + + img_left = image.copy() + img_right = Image.new("RGB", (w, h), (255, 255, 255)) + + random.seed(0) + draw_left = ImageDraw.Draw(img_left) + draw_right = ImageDraw.Draw(img_right) + for idx, (box, txt) in enumerate(zip(dt_boxes, txts)): + if scores is not None and float(scores[idx]) < self.text_score: + continue + + color = self.get_random_color() + + box_list = np.array(box).reshape(8).tolist() + draw_left.polygon(box_list, fill=color) + draw_right.polygon(box_list, outline=color) + + box_height = self.get_box_height(box) + box_width = self.get_box_width(box) + if box_height > 2 * box_width: + font_size = max(int(box_width * 0.9), 10) + font = ImageFont.truetype(font_path, font_size, encoding="utf-8") + cur_y = box[0][1] + + for c in txt: + draw_right.text( + (box[0][0] + 3, cur_y), c, fill=(0, 0, 0), font=font + ) + cur_y += self.get_char_size(font, c) + else: + font_size = max(int(box_height * 0.8), 10) + font = ImageFont.truetype(font_path, font_size, encoding="utf-8") + draw_right.text([box[0][0], box[0][1]], txt, fill=(0, 0, 0), font=font) + + img_left = Image.blend(image, img_left, 0.5) + img_show = Image.new("RGB", (w * 2, h), (255, 255, 255)) + img_show.paste(img_left, (0, 0, w, h)) + img_show.paste(img_right, (w, 0, w * 2, h)) + return np.array(img_show) + + @staticmethod + def get_font_path(font_path: Optional[Union[str, Path]] = None) -> str: + if font_path is None or not Path(font_path).exists(): + raise FileNotFoundError( + f"The {font_path} does not exists! \n" + f"You could download the file in the https://drive.google.com/file/d/1evWVX38EFNwTq_n5gTFgnlv8tdaNcyIA/view?usp=sharing" + ) + return str(font_path) + + @staticmethod + def get_random_color() -> Tuple[int, int, int]: + return ( + random.randint(0, 255), + random.randint(0, 255), + random.randint(0, 255), + ) + + @staticmethod + def get_box_height(box: List[List[float]]) -> float: + return math.sqrt((box[0][0] - box[3][0]) ** 2 + (box[0][1] - box[3][1]) ** 2) + + @staticmethod + def get_box_width(box: List[List[float]]) -> float: + return math.sqrt((box[0][0] - box[1][0]) ** 2 + (box[0][1] - box[1][1]) ** 2) + + @staticmethod + def get_char_size(font, char_str: str) -> float: + # compatible with Pillow v9 and v10. + if hasattr(font, "getsize"): + get_size_func = getattr(font, "getsize") + return get_size_func(char_str)[1] + + if hasattr(font, "getlength"): + get_size_func = getattr(font, "getlength") + return get_size_func(char_str) + + raise ValueError( + "The Pillow ImageFont instance has not getsize or getlength func." + )