From ef9202445af461e7853a42aece05ec4312f965a5 Mon Sep 17 00:00:00 2001 From: SWHL Date: Sat, 18 May 2024 16:31:23 +0800 Subject: [PATCH] refactor(rapidocr_paddle): Decoupling utils.py files --- python/rapidocr_paddle/utils.py | 549 ------------------ python/rapidocr_paddle/utils/__init__.py | 19 + python/rapidocr_paddle/utils/infer_engine.py | 121 ++++ python/rapidocr_paddle/utils/load_image.py | 123 ++++ python/rapidocr_paddle/utils/logger.py | 19 + .../rapidocr_paddle/utils/parse_parameters.py | 187 ++++++ python/rapidocr_paddle/utils/vis_res.py | 143 +++++ 7 files changed, 612 insertions(+), 549 deletions(-) delete mode 100644 python/rapidocr_paddle/utils.py create mode 100644 python/rapidocr_paddle/utils/__init__.py create mode 100644 python/rapidocr_paddle/utils/infer_engine.py create mode 100644 python/rapidocr_paddle/utils/load_image.py create mode 100644 python/rapidocr_paddle/utils/logger.py create mode 100644 python/rapidocr_paddle/utils/parse_parameters.py create mode 100644 python/rapidocr_paddle/utils/vis_res.py diff --git a/python/rapidocr_paddle/utils.py b/python/rapidocr_paddle/utils.py deleted file mode 100644 index b8135be..0000000 --- a/python/rapidocr_paddle/utils.py +++ /dev/null @@ -1,549 +0,0 @@ -# -*- encoding: utf-8 -*- -# @Author: SWHL -# @Contact: liekkaskono@163.com -import argparse -import math -import os -import platform -import random -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 paddle -import yaml -from paddle import inference -from PIL import Image, ImageDraw, ImageFont, UnidentifiedImageError - -root_dir = Path(__file__).resolve().parent -InputType = Union[str, np.ndarray, bytes, Path, Image.Image] - - -class PaddleInferSession: - def __init__(self, config, mode: Optional[str] = None) -> None: - self.mode = mode - - model_dir = Path(config["model_path"]) - pdmodel_path = model_dir / "inference.pdmodel" - pdiparams_path = model_dir / "inference.pdiparams" - - self._verify_model(pdmodel_path) - self._verify_model(pdiparams_path) - - infer_opts = inference.Config(str(pdmodel_path), str(pdiparams_path)) - - if config["use_cuda"]: - gpu_id = self.get_infer_gpuid() - if gpu_id is None: - warnings.warn( - "GPU is not found in current device by nvidia-smi. Please check your device or ignore it if run on jetson." - ) - infer_opts.enable_use_gpu(config["gpu_mem"], config["gpu_id"]) - else: - infer_opts.disable_gpu() - - cpu_nums = os.cpu_count() - infer_num_threads = config.get("cpu_math_library_num_threads", -1) - if infer_num_threads != -1 and 1 <= infer_num_threads <= cpu_nums: - infer_opts.set_cpu_math_library_num_threads(infer_num_threads) - - # enable memory optim - infer_opts.enable_memory_optim() - infer_opts.disable_glog_info() - infer_opts.delete_pass("conv_transpose_eltwiseadd_bn_fuse_pass") - infer_opts.delete_pass("matmul_transpose_reshape_fuse_pass") - infer_opts.switch_use_feed_fetch_ops(False) - infer_opts.switch_ir_optim(True) - - self.predictor = inference.create_predictor(infer_opts) - - def __call__(self, img: np.ndarray): - input_tensor = self.get_input_tensors() - output_tensors = self.get_output_tensors() - - input_tensor.copy_from_cpu(img) - self.predictor.run() - - outputs = [] - for output_tensor in output_tensors: - output = output_tensor.copy_to_cpu() - outputs.append(output) - - self.predictor.try_shrink_memory() - return outputs - - @staticmethod - def _verify_model(model_path): - 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.") - - def get_input_tensors( - self, - ): - input_names = self.predictor.get_input_names() - for name in input_names: - input_tensor = self.predictor.get_input_handle(name) - return input_tensor - - def get_output_tensors( - self, - ): - output_names = self.predictor.get_output_names() - if self.mode == "rec": - output_name = "softmax_0.tmp_0" - if output_name in output_names: - return [self.predictor.get_output_handle(output_name)] - - output_tensors = [] - for output_name in output_names: - output_tensor = self.predictor.get_output_handle(output_name) - output_tensors.append(output_tensor) - return output_tensors - - @staticmethod - def get_infer_gpuid(): - sysstr = platform.system() - if sysstr == "Windows": - return 0 - - if not paddle.device.is_compiled_with_rocm: - cmd = "env | grep CUDA_VISIBLE_DEVICES" - else: - cmd = "env | grep HIP_VISIBLE_DEVICES" - env_cuda = os.popen(cmd).readlines() - - if len(env_cuda) == 0: - return 0 - - gpu_id = env_cuda[0].strip().split("=")[1] - return int(gpu_id[0]) - - -class PaddleInferError(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("--cpu_math_library_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) - det_group.add_argument("--det_gpu_id", type=int, default=0) - det_group.add_argument("--det_gpu_mem", type=int, default=500) - 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) - cls_group.add_argument("--cls_gpu_id", type=int, default=0) - cls_group.add_argument("--cls_gpu_mem", type=int, default=500) - 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) - rec_group.add_argument("--rec_gpu_id", type=int, default=0) - rec_group.add_argument("--rec_gpu_mem", type=int, default=500) - 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 = ["cpu_math_library_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." - ) diff --git a/python/rapidocr_paddle/utils/__init__.py b/python/rapidocr_paddle/utils/__init__.py new file mode 100644 index 0000000..84031f4 --- /dev/null +++ b/python/rapidocr_paddle/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 PaddleInferSession +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_paddle/utils/infer_engine.py b/python/rapidocr_paddle/utils/infer_engine.py new file mode 100644 index 0000000..0a4de3c --- /dev/null +++ b/python/rapidocr_paddle/utils/infer_engine.py @@ -0,0 +1,121 @@ +# -*- encoding: utf-8 -*- +# @Author: SWHL +# @Contact: liekkaskono@163.com +import os +import platform +from pathlib import Path +from typing import Optional + +import numpy as np +import paddle +from paddle import inference + +from .logger import get_logger + + +class PaddleInferSession: + def __init__(self, config, mode: Optional[str] = None) -> None: + self.logger = get_logger("PaddleInferSession") + self.mode = mode + + model_dir = Path(config["model_path"]) + pdmodel_path = model_dir / "inference.pdmodel" + pdiparams_path = model_dir / "inference.pdiparams" + + self._verify_model(pdmodel_path) + self._verify_model(pdiparams_path) + + infer_opts = inference.Config(str(pdmodel_path), str(pdiparams_path)) + + if config["use_cuda"]: + gpu_id = self.get_infer_gpuid() + if gpu_id is None: + self.logger.warning( + "GPU is not found in current device by nvidia-smi. Please check your device or ignore it if run on jetson." + ) + infer_opts.enable_use_gpu(config["gpu_mem"], config["gpu_id"]) + else: + infer_opts.disable_gpu() + + cpu_nums = os.cpu_count() + infer_num_threads = config.get("cpu_math_library_num_threads", -1) + if infer_num_threads != -1 and 1 <= infer_num_threads <= cpu_nums: + infer_opts.set_cpu_math_library_num_threads(infer_num_threads) + + # enable memory optim + infer_opts.enable_memory_optim() + infer_opts.disable_glog_info() + infer_opts.delete_pass("conv_transpose_eltwiseadd_bn_fuse_pass") + infer_opts.delete_pass("matmul_transpose_reshape_fuse_pass") + infer_opts.switch_use_feed_fetch_ops(False) + infer_opts.switch_ir_optim(True) + + self.predictor = inference.create_predictor(infer_opts) + + def __call__(self, img: np.ndarray): + input_tensor = self.get_input_tensors() + output_tensors = self.get_output_tensors() + + input_tensor.copy_from_cpu(img) + self.predictor.run() + + outputs = [] + for output_tensor in output_tensors: + output = output_tensor.copy_to_cpu() + outputs.append(output) + + self.predictor.try_shrink_memory() + return outputs + + @staticmethod + def _verify_model(model_path): + 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.") + + def get_input_tensors( + self, + ): + input_names = self.predictor.get_input_names() + for name in input_names: + input_tensor = self.predictor.get_input_handle(name) + return input_tensor + + def get_output_tensors( + self, + ): + output_names = self.predictor.get_output_names() + if self.mode == "rec": + output_name = "softmax_0.tmp_0" + if output_name in output_names: + return [self.predictor.get_output_handle(output_name)] + + output_tensors = [] + for output_name in output_names: + output_tensor = self.predictor.get_output_handle(output_name) + output_tensors.append(output_tensor) + return output_tensors + + @staticmethod + def get_infer_gpuid(): + sysstr = platform.system() + if sysstr == "Windows": + return 0 + + if not paddle.device.is_compiled_with_rocm: + cmd = "env | grep CUDA_VISIBLE_DEVICES" + else: + cmd = "env | grep HIP_VISIBLE_DEVICES" + env_cuda = os.popen(cmd).readlines() + + if len(env_cuda) == 0: + return 0 + + gpu_id = env_cuda[0].strip().split("=")[1] + return int(gpu_id[0]) + + +class PaddleInferError(Exception): + pass diff --git a/python/rapidocr_paddle/utils/load_image.py b/python/rapidocr_paddle/utils/load_image.py new file mode 100644 index 0000000..0566053 --- /dev/null +++ b/python/rapidocr_paddle/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_paddle/utils/logger.py b/python/rapidocr_paddle/utils/logger.py new file mode 100644 index 0000000..ffd1cd0 --- /dev/null +++ b/python/rapidocr_paddle/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_paddle/utils/parse_parameters.py b/python/rapidocr_paddle/utils/parse_parameters.py new file mode 100644 index 0000000..dcf4579 --- /dev/null +++ b/python/rapidocr_paddle/utils/parse_parameters.py @@ -0,0 +1,187 @@ +# -*- 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("--cpu_math_library_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) + det_group.add_argument("--det_gpu_id", type=int, default=0) + det_group.add_argument("--det_gpu_mem", type=int, default=500) + 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) + cls_group.add_argument("--cls_gpu_id", type=int, default=0) + cls_group.add_argument("--cls_gpu_mem", type=int, default=500) + 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) + rec_group.add_argument("--rec_gpu_id", type=int, default=0) + rec_group.add_argument("--rec_gpu_mem", type=int, default=500) + 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 = ["cpu_math_library_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_paddle/utils/vis_res.py b/python/rapidocr_paddle/utils/vis_res.py new file mode 100644 index 0000000..405beab --- /dev/null +++ b/python/rapidocr_paddle/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." + )