diff --git a/python/demo.py b/python/demo.py index 6a0e4e1..200c2ed 100644 --- a/python/demo.py +++ b/python/demo.py @@ -6,7 +6,8 @@ from rapidocr import RapidOCR engine = RapidOCR() img_url = "https://img1.baidu.com/it/u=3619974146,1266987475&fm=253&fmt=auto&app=138&f=JPEG?w=500&h=516" -result = engine(img_url) +img_url = "tests/test_files/return_word_debug.jpg" +result = engine(img_url, return_word_box=True) print(result) result.vis("vis_result.jpg") diff --git a/python/rapidocr/__init__.py b/python/rapidocr/__init__.py index a81fcfd..15e216c 100644 --- a/python/rapidocr/__init__.py +++ b/python/rapidocr/__init__.py @@ -2,5 +2,6 @@ # @Author: SWHL # @Contact: liekkaskono@163.com from .main import RapidOCR -from .utils import LoadImageError, VisRes +from .utils.load_image import LoadImageError from .utils.typings import EngineType, LangCls, LangDet, LangRec, ModelType, OCRVersion +from .utils.vis_res import VisRes diff --git a/python/rapidocr/cal_rec_boxes/main.py b/python/rapidocr/cal_rec_boxes/main.py index 928312c..2b8728f 100644 --- a/python/rapidocr/cal_rec_boxes/main.py +++ b/python/rapidocr/cal_rec_boxes/main.py @@ -25,13 +25,13 @@ class CalRecBoxes: def __call__( self, imgs: List[np.ndarray], - dt_boxes: List[np.ndarray], + dt_boxes: np.ndarray, rec_res: TextRecOutput, return_single_char_box: bool = False, ) -> TextRecOutput: word_results = [] for idx, (img, box) in enumerate(zip(imgs, dt_boxes)): - if rec_res.txts is None: + if rec_res.txts is None or img.size == 0 or rec_res.word_results is None: continue h, w = img.shape[:2] diff --git a/python/rapidocr/ch_ppocr_det/main.py b/python/rapidocr/ch_ppocr_det/main.py index 146f3d5..9c2736f 100644 --- a/python/rapidocr/ch_ppocr_det/main.py +++ b/python/rapidocr/ch_ppocr_det/main.py @@ -15,7 +15,7 @@ # @Author: SWHL # @Contact: liekkaskono@163.com import time -from typing import Any, Dict, List +from typing import Any, Dict import numpy as np @@ -77,7 +77,7 @@ class TextDetector: return DetPreProcess(limit_side_len, self.limit_type, self.mean, self.std) @staticmethod - def sorted_boxes(dt_boxes: np.ndarray) -> List[np.ndarray]: + def sorted_boxes(dt_boxes: np.ndarray) -> np.ndarray: """ Sort text boxes in order from top to bottom, left to right args: @@ -100,4 +100,4 @@ class TextDetector: _boxes[j + 1] = tmp else: break - return _boxes + return np.array(_boxes) diff --git a/python/rapidocr/ch_ppocr_det/utils.py b/python/rapidocr/ch_ppocr_det/utils.py index 1311f32..8449d81 100644 --- a/python/rapidocr/ch_ppocr_det/utils.py +++ b/python/rapidocr/ch_ppocr_det/utils.py @@ -18,7 +18,7 @@ from ..utils.vis_res import VisRes class TextDetOutput: img: Optional[np.ndarray] = None boxes: Optional[np.ndarray] = None - scores: Optional[Tuple[float]] = None + scores: Optional[List[float]] = None elapse: float = 0.0 def __len__(self): diff --git a/python/rapidocr/ch_ppocr_rec/main.py b/python/rapidocr/ch_ppocr_rec/main.py index 430174a..43789f5 100644 --- a/python/rapidocr/ch_ppocr_rec/main.py +++ b/python/rapidocr/ch_ppocr_rec/main.py @@ -35,7 +35,6 @@ DEFAULT_MODEL_PATH = Path(__file__).parent.parent / "models" class TextRecognizer: def __init__(self, cfg: Dict[str, Any]): self.session = get_engine(cfg.engine_type)(cfg) - self.logger = logger # onnx has inner character, other engine get or download character_dict_path character, character_dict_path = self.get_character_dict(cfg) @@ -77,7 +76,7 @@ class TextRecognizer: file_url=dict_download_url, sha256=None, save_path=dict_path, - logger=self.logger, + logger=logger, ) ) diff --git a/python/rapidocr/ch_ppocr_rec/typings.py b/python/rapidocr/ch_ppocr_rec/typings.py index 0239d9f..a61952a 100644 --- a/python/rapidocr/ch_ppocr_rec/typings.py +++ b/python/rapidocr/ch_ppocr_rec/typings.py @@ -36,7 +36,7 @@ class TextRecInput: class TextRecOutput: imgs: Optional[List[np.ndarray]] = None txts: Optional[Tuple[str]] = None - scores: Tuple[float] = (1.0,) + scores: List[float] = field(default_factory=lambda: [1.0]) word_results: Tuple[Tuple[str, float, Optional[List[List[int]]]]] = ( ("", 1.0, None), ) diff --git a/python/rapidocr/inference_engine/onnxruntime/main.py b/python/rapidocr/inference_engine/onnxruntime/main.py index 411eacf..7ed03b7 100644 --- a/python/rapidocr/inference_engine/onnxruntime/main.py +++ b/python/rapidocr/inference_engine/onnxruntime/main.py @@ -9,7 +9,7 @@ from typing import Any, Dict, List import numpy as np from onnxruntime import GraphOptimizationLevel, InferenceSession, SessionOptions -from ...utils import DownloadFile, DownloadFileInput +from ...utils.download_file import DownloadFile, DownloadFileInput from ...utils.log import logger from ..base import FileInfo, InferSession from .provider_config import ProviderConfig @@ -17,8 +17,6 @@ from .provider_config import ProviderConfig class OrtInferSession(InferSession): def __init__(self, cfg: Dict[str, Any]): - self.logger = logger - # support custom session (PR #451) session = cfg.get("session", None) if session is not None: @@ -27,7 +25,7 @@ class OrtInferSession(InferSession): f"Expected session to be an InferenceSession, got {type(session)}" ) - self.logger.debug("Using the provided InferenceSession for inference.") + logger.debug("Using the provided InferenceSession for inference.") self.session = session return @@ -48,11 +46,11 @@ class OrtInferSession(InferSession): file_url=model_info["model_dir"], sha256=model_info["SHA256"], save_path=model_path, - logger=self.logger, + logger=logger, ) DownloadFile.run(download_params) - self.logger.info(f"Using {model_path}") + logger.info(f"Using {model_path}") model_path = Path(model_path) self._verify_model(model_path) diff --git a/python/rapidocr/inference_engine/onnxruntime/provider_config.py b/python/rapidocr/inference_engine/onnxruntime/provider_config.py index 8b47144..deff27a 100644 --- a/python/rapidocr/inference_engine/onnxruntime/provider_config.py +++ b/python/rapidocr/inference_engine/onnxruntime/provider_config.py @@ -19,8 +19,6 @@ class EP(Enum): class ProviderConfig: def __init__(self, engine_cfg: Dict[str, Any]): - self.logger = logger - self.had_providers: List[str] = get_available_providers() self.default_provider = self.had_providers[0] @@ -37,13 +35,13 @@ class ProviderConfig: results.insert(0, (EP.CUDA_EP.value, self.cuda_ep_cfg())) if self.is_dml_available(): - self.logger.info( + logger.info( "Windows 10 or above detected, try to use DirectML as primary provider" ) results.insert(0, (EP.DIRECTML_EP.value, self.dml_ep_cfg())) if self.is_cann_available(): - self.logger.info("Try to use CANNExecutionProvider to infer") + logger.info("Try to use CANNExecutionProvider to infer") results.insert(0, (EP.CANN_EP.value, self.cann_ep_cfg())) return results @@ -79,10 +77,10 @@ class ProviderConfig: for ep, check_func in providers_to_check.items(): if check_func() and first_provider != ep.value: - self.logger.warning( + logger.warning( f"{ep.value} is available, but the inference part is automatically shifted to be executed under {first_provider}. " ) - self.logger.warning(f"The available lists are {session_providers}") + logger.warning(f"The available lists are {session_providers}") def is_cuda_available(self) -> bool: if not self.cfg_use_cuda: @@ -92,7 +90,7 @@ class ProviderConfig: if get_device() == "GPU" and CUDA_EP in self.had_providers: return True - self.logger.warning( + logger.warning( f"{CUDA_EP} is not in available providers ({self.had_providers}). Use {self.default_provider} inference by default." ) install_instructions = [ @@ -113,7 +111,7 @@ class ProviderConfig: cur_os = platform.system() if cur_os != "Windows": - self.logger.warning( + logger.warning( f"DirectML is only supported in Windows OS. The current OS is {cur_os}. Use {self.default_provider} inference by default.", ) return False @@ -123,7 +121,7 @@ class ProviderConfig: int(window_build_number_str) if window_build_number_str.isdigit() else 0 ) if window_build_number < 18362: - self.logger.warning( + logger.warning( f"DirectML is only supported in Windows 10 Build 18362 and above OS. The current Windows Build is {window_build_number}. Use {self.default_provider} inference by default.", ) return False @@ -132,7 +130,7 @@ class ProviderConfig: if DML_EP in self.had_providers: return True - self.logger.warning( + logger.warning( f"{DML_EP} is not in available providers ({self.had_providers}). Use {self.default_provider} inference by default." ) install_instructions = [ @@ -152,7 +150,7 @@ class ProviderConfig: if CANN_EP in self.had_providers: return True - self.logger.warning( + logger.warning( f"{CANN_EP} is not in available providers ({self.had_providers}). Use {self.default_provider} inference by default." ) install_instructions = [ @@ -167,4 +165,4 @@ class ProviderConfig: def print_log(self, log_list: List[str]): for log_info in log_list: - self.logger.info(log_info) + logger.info(log_info) diff --git a/python/rapidocr/inference_engine/openvino.py b/python/rapidocr/inference_engine/openvino.py index a9f0484..11b905c 100644 --- a/python/rapidocr/inference_engine/openvino.py +++ b/python/rapidocr/inference_engine/openvino.py @@ -18,7 +18,6 @@ from .base import FileInfo, InferSession class OpenVINOInferSession(InferSession): def __init__(self, cfg: DictConfig): super().__init__(cfg) - self.logger = logger core = Core() @@ -38,11 +37,11 @@ class OpenVINOInferSession(InferSession): file_url=model_info["model_dir"], sha256=model_info["SHA256"], save_path=model_path, - logger=self.logger, + logger=logger, ) DownloadFile.run(download_params) - self.logger.info(f"Using {model_path}") + logger.info(f"Using {model_path}") model_path = Path(model_path) self._verify_model(model_path) @@ -85,7 +84,7 @@ class OpenVINOInferSession(InferSession): if scheduling_core_type is not None: config["SCHEDULING_CORE_TYPE"] = str(scheduling_core_type) - self.logger.info(f"Using OpenVINO config: {config}") + logger.info(f"Using OpenVINO config: {config}") return config def __call__(self, input_content: np.ndarray) -> np.ndarray: diff --git a/python/rapidocr/inference_engine/paddle.py b/python/rapidocr/inference_engine/paddle.py index c8a5679..5ae6008 100644 --- a/python/rapidocr/inference_engine/paddle.py +++ b/python/rapidocr/inference_engine/paddle.py @@ -19,7 +19,6 @@ from .base import FileInfo, InferSession class PaddleInferSession(InferSession): def __init__(self, cfg, mode: Optional[str] = None) -> None: - self.logger = logger self.mode = mode # 设置设备特定的环境变量 @@ -37,7 +36,7 @@ class PaddleInferSession(InferSession): def _set_envs(envs): for key, val in envs.items(): os.environ[key] = val - self.logger.debug(f"{key} has been set to {val}.") + logger.debug(f"{key} has been set to {val}.") if engine_cfg.get("use_cuda", False): envs = {} @@ -97,8 +96,8 @@ class PaddleInferSession(InferSession): model_info, default_model_dir, pdiparams_name ) - self.logger.info(f"Using {pdmodel_path}") - self.logger.info(f"Using {pdiparams_path}") + logger.info(f"Using {pdmodel_path}") + logger.info(f"Using {pdiparams_path}") return pdmodel_path, pdiparams_path model_dir = Path(model_dir) @@ -109,8 +108,8 @@ class PaddleInferSession(InferSession): self._verify_model(pdmodel_path) self._verify_model(pdiparams_path) - self.logger.info(f"Using {pdmodel_path}") - self.logger.info(f"Using {pdiparams_path}") + logger.info(f"Using {pdmodel_path}") + logger.info(f"Using {pdiparams_path}") return pdmodel_path, pdiparams_path def download_model( @@ -125,7 +124,7 @@ class PaddleInferSession(InferSession): file_url=model_file_url, sha256=model_info[model_file_name], save_path=model_file_path, - logger=self.logger, + logger=logger, ) ) return model_file_path @@ -136,27 +135,27 @@ class PaddleInferSession(InferSession): if cfg.engine_cfg.use_cuda: gpu_id = self.get_infer_gpuid() if gpu_id is None: - self.logger.warning( + logger.warning( "GPU is not found in current device by nvidia-smi. Please check your device or ignore it if run on jetson." ) device_id = cfg.engine_cfg.gpu_id infer_opts.enable_use_gpu(cfg.engine_cfg.gpu_mem, device_id) - self.logger.info(f"Using GPU device with ID: {device_id}") + logger.info(f"Using GPU device with ID: {device_id}") elif cfg.engine_cfg.use_npu: npu_id = cfg.engine_cfg.get("npu_id", 0) infer_opts.enable_custom_device("npu", npu_id) - self.logger.info(f"Using NPU device with ID: {npu_id}") + logger.info(f"Using NPU device with ID: {npu_id}") else: # CPU infer_opts.disable_gpu() - self.logger.info("Using CPU device") + logger.info("Using CPU device") cpu_nums = os.cpu_count() infer_num_threads = cfg.engine_cfg.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) - self.logger.info(f"Set CPU math library threads to: {infer_num_threads}") + logger.info(f"Set CPU math library threads to: {infer_num_threads}") # enable memory optim infer_opts.enable_memory_optim() @@ -174,31 +173,29 @@ class PaddleInferSession(InferSession): if cfg.engine_cfg.use_cuda: gpu_id = self.get_infer_gpuid() if gpu_id is None: - self.logger.warning( + logger.warning( "GPU is not found in current device by nvidia-smi. Please check your device or ignore it if run on jetson." ) device_id = cfg.engine_cfg.gpu_id infer_opts.enable_use_gpu(cfg.engine_cfg.gpu_mem, device_id) - self.logger.info(f"Using GPU device with ID: {device_id}") + logger.info(f"Using GPU device with ID: {device_id}") elif cfg.engine_cfg.use_npu: npu_id = cfg.engine_cfg.get("npu_id", 0) infer_opts.enable_custom_device("npu", npu_id) - self.logger.info(f"Using NPU device with ID: {npu_id}") + logger.info(f"Using NPU device with ID: {npu_id}") else: # CPU infer_opts.disable_gpu() if hasattr(infer_opts, "disable_mkldnn"): infer_opts.disable_mkldnn() - self.logger.info("Using CPU device") + logger.info("Using CPU device") cpu_nums = os.cpu_count() infer_num_threads = cfg.engine_cfg.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) - self.logger.info( - f"Set CPU math library threads to: {infer_num_threads}" - ) + logger.info(f"Set CPU math library threads to: {infer_num_threads}") if hasattr(infer_opts, "enable_new_ir"): infer_opts.enable_new_ir(True) diff --git a/python/rapidocr/main.py b/python/rapidocr/main.py index b880368..544678f 100644 --- a/python/rapidocr/main.py +++ b/python/rapidocr/main.py @@ -15,18 +15,18 @@ from .ch_ppocr_cls import TextClassifier, TextClsOutput from .ch_ppocr_det import TextDetector, TextDetOutput from .ch_ppocr_rec import TextRecInput, TextRecognizer, TextRecOutput from .cli import check_install, generate_cfg -from .utils import ( - LoadImage, - RapidOCROutput, - VisRes, - add_round_letterbox, - get_padding_h, +from .utils.load_image import LoadImage +from .utils.log import logger +from .utils.output import RapidOCROutput +from .utils.parse_parameters import ParseParams +from .utils.process_img import ( + apply_vertical_padding, get_rotate_crop_image, + map_boxes_to_original, resize_image_within_bounds, ) -from .utils.log import logger -from .utils.parse_parameters import ParseParams from .utils.typings import LangRec +from .utils.vis_res import VisRes root_dir = Path(__file__).resolve().parent DEFAULT_CFG_PATH = root_dir / "config.yaml" @@ -39,7 +39,6 @@ class RapidOCR: cfg = self._load_config(config_path, params) logger.setLevel(cfg.Global.log_level.upper()) - self.logger = logger self._initialize(cfg) @@ -109,64 +108,134 @@ class RapidOCR: ori_img = self.load_img(img_content) img, op_record = self.preprocess_img(ori_img) + det_res, cls_res, rec_res, cropped_img_list = self.run_ocr_steps(img, op_record) + return self.build_final_output( + ori_img, det_res, cls_res, rec_res, cropped_img_list, op_record + ) + def run_ocr_steps(self, img: np.ndarray, op_record: Dict[str, Any]): det_res, cls_res, rec_res = TextDetOutput(), TextClsOutput(), TextRecOutput() + if self.use_det: try: - img, det_res = self.get_det_res(img, op_record) + cropped_img_list, det_res = self.detect_and_crop(img, op_record) except RapidOCRError as e: - self.logger.warning(e) - return RapidOCROutput() + logger.warning(e) + return TextDetOutput(), TextClsOutput(), TextRecOutput(), [] + else: + cropped_img_list = [img] if self.use_cls: try: - img, cls_res = self.get_cls_res(img) + cls_img_list, cls_res = self.cls_and_rotate(cropped_img_list) except RapidOCRError as e: - self.logger.warning(e) - return RapidOCROutput() + logger.warning(e) + return det_res, TextClsOutput(), TextRecOutput(), [] + else: + cls_img_list = cropped_img_list if self.use_rec: - rec_res = self.get_rec_res(img) + try: + rec_res = self.recognize_txt(cls_img_list) + except RapidOCRError as e: + logger.warning(e) + return det_res, cls_res, TextRecOutput(), [] - return self.finalize_results(ori_img, det_res, cls_res, rec_res, img, op_record) + return det_res, cls_res, rec_res, cropped_img_list - def finalize_results( + def build_final_output( self, ori_img: np.ndarray, det_res: TextDetOutput, cls_res: TextClsOutput, rec_res: TextRecOutput, - img: List[np.ndarray], + cropped_img_list: List[np.ndarray], op_record: Dict[str, Any], ) -> Union[TextDetOutput, TextClsOutput, TextRecOutput, RapidOCROutput]: - raw_h, raw_w = ori_img.shape[:2] + ori_h, ori_w = ori_img.shape[:2] + + if det_res.boxes is not None: + det_res.boxes = map_boxes_to_original( + det_res.boxes, op_record, ori_h, ori_w + ) + + # 过滤识别结果为空的值 + if ( + rec_res.txts is not None + and det_res.boxes is not None + and det_res.scores is not None + ): + empty_ids = {i for i, v in enumerate(rec_res.txts) if not v.strip()} + + det_res.boxes = np.array( + [v for i, v in enumerate(det_res.boxes) if i not in empty_ids] + ) + det_res.scores = [ + v for i, v in enumerate(det_res.scores) if i not in empty_ids + ] + + rec_res.txts = [v for i, v in enumerate(rec_res.txts) if i not in empty_ids] + rec_res.word_results = [ + v for i, v in enumerate(rec_res.word_results) if i not in empty_ids + ] + + # 仅分类结果 + if ( + det_res.boxes is None + and rec_res.txts is None + and cls_res.cls_res is not None + ): + return cls_res + + # 无有效输出 + if det_res.boxes is None and rec_res.txts is None: + return RapidOCROutput() + + # 仅识别结果（无检测） + if det_res.boxes is None and rec_res.txts is not None: + return rec_res + + # 仅检测结果（无识别） + if det_res.boxes is not None and rec_res.txts is None: + return det_res if ( self.return_word_box and det_res.boxes is not None - and all(v for v in rec_res.word_results) + and all(rec_res.word_results) ): rec_res.word_results = self.calc_word_boxes( - img, det_res, rec_res, op_record, raw_h, raw_w + cropped_img_list, det_res.boxes, rec_res, op_record, ori_h, ori_w ) - if det_res.boxes is not None: - det_res.boxes = self._get_origin_points( - det_res.boxes, op_record, raw_h, raw_w - ) - return self.get_final_res(ori_img, det_res, cls_res, rec_res) + ocr_res = RapidOCROutput( + img=ori_img, + boxes=det_res.boxes, + txts=rec_res.txts, + scores=rec_res.scores, + word_results=rec_res.word_results, + elapse_list=[det_res.elapse, cls_res.elapse, rec_res.elapse], + viser=VisRes( + text_score=self.cfg.Global.text_score, + lang_type=self.cfg.Rec.lang_type, + font_path=self.cfg.Global.font_path, + ), + ) + + ocr_res = self.filter_by_text_score(ocr_res) + return ocr_res if len(ocr_res) > 0 else RapidOCROutput() def calc_word_boxes( self, img: List[np.ndarray], - det_res: TextDetOutput, + dt_boxes: np.ndarray, rec_res: TextRecOutput, op_record: Dict[str, Any], raw_h: int, raw_w: int, ) -> Any: rec_res = self.cal_rec_boxes( - img, det_res.boxes, rec_res, self.return_single_char_box + img, dt_boxes, rec_res, self.return_single_char_box ) origin_words = [] @@ -176,8 +245,8 @@ class RapidOCR: if bbox is None: continue - origin_words_points = self._get_origin_points( - [bbox], op_record, raw_h, raw_w + origin_words_points = map_boxes_to_original( + np.array([bbox]).astype(np.float64), op_record, raw_h, raw_w ) origin_words_points = origin_words_points.astype(np.int32).tolist()[0] origin_words_item.append((txt, score, origin_words_points)) @@ -188,14 +257,14 @@ class RapidOCR: def update_params( self, - use_det, - use_cls, - use_rec, - return_word_box, - return_single_char_box, - text_score, - box_thresh, - unclip_ratio, + use_det: Optional[bool] = None, + use_cls: Optional[bool] = None, + use_rec: Optional[bool] = None, + return_word_box: bool = False, + return_single_char_box: bool = False, + text_score: float = 0.5, + box_thresh: float = 0.5, + unclip_ratio: float = 1.6, ): self.use_det = self.use_det if use_det is None else use_det self.use_cls = self.use_cls if use_cls is None else use_cls @@ -215,28 +284,30 @@ class RapidOCR: op_record["preprocess"] = {"ratio_h": ratio_h, "ratio_w": ratio_w} return img, op_record - def get_det_res( + def detect_and_crop( self, img: np.ndarray, op_record: Dict[str, Any] ) -> Tuple[List[np.ndarray], TextDetOutput]: - img, op_record = self._add_letterbox(img, op_record) + img, op_record = apply_vertical_padding( + img, op_record, self.width_height_ratio, self.min_height + ) det_res = self.text_det(img) + if det_res.boxes is None: raise RapidOCRError("The text detection result is empty") - img_list = self.get_crop_img_list(img, det_res) - return img_list, det_res + img_crop_list = self.crop_text_regions(img, det_res.boxes) + return img_crop_list, det_res - def get_crop_img_list( - self, img: np.ndarray, det_res: TextDetOutput + def crop_text_regions( + self, img: np.ndarray, det_boxes: np.ndarray ) -> List[np.ndarray]: img_crop_list = [] - for box in det_res.boxes: - tmp_box = copy.deepcopy(box) - img_crop = get_rotate_crop_image(img, tmp_box) + for box in det_boxes: + img_crop = get_rotate_crop_image(img, copy.deepcopy(box)) img_crop_list.append(img_crop) return img_crop_list - def get_cls_res( + def cls_and_rotate( self, img: List[np.ndarray] ) -> Tuple[List[np.ndarray], TextClsOutput]: cls_res = self.text_cls(img) @@ -244,112 +315,36 @@ class RapidOCR: raise RapidOCRError("The text classifier is empty") return cls_res.img_list, cls_res - def get_rec_res(self, img: List[np.ndarray]) -> TextRecOutput: + def recognize_txt(self, img: List[np.ndarray]) -> TextRecOutput: rec_input = TextRecInput(img=img, return_word_box=self.return_word_box) - return self.text_rec(rec_input) - def get_final_res( - self, - ori_img: np.ndarray, - det_res: TextDetOutput, - cls_res: TextClsOutput, - rec_res: TextRecOutput, - ) -> Union[TextDetOutput, TextClsOutput, TextRecOutput, RapidOCROutput]: - dt_boxes = det_res.boxes - txt_res = rec_res.txts + rec_res = self.text_rec(rec_input) + if rec_res.txts is None: + raise RapidOCRError("The text recognize result is empty") - if dt_boxes is None and txt_res is None and cls_res.cls_res is not None: - return cls_res - - if dt_boxes is None and txt_res is None: - return RapidOCROutput() - - if dt_boxes is None and txt_res is not None: - return rec_res - - if dt_boxes is not None and txt_res is None: - return det_res - - ocr_res = RapidOCROutput( - img=ori_img, - boxes=det_res.boxes, - txts=rec_res.txts, - scores=rec_res.scores, - word_results=rec_res.word_results, - elapse_list=[det_res.elapse, cls_res.elapse, rec_res.elapse], - viser=VisRes( - text_score=self.cfg.Global.text_score, - lang_type=self.cfg.Rec.lang_type, - font_path=self.cfg.Global.font_path, - ), - ) - ocr_res = self.filter_by_text_score(ocr_res) - if len(ocr_res) <= 0: - return RapidOCROutput() - - return ocr_res + return rec_res def filter_by_text_score(self, ocr_res: RapidOCROutput) -> RapidOCROutput: - filter_boxes, filter_txts, filter_scores = [], [], [] - for box, txt, score in zip(ocr_res.boxes, ocr_res.txts, ocr_res.scores): - if float(score) >= self.text_score: - filter_boxes.append(box) - filter_txts.append(txt) - filter_scores.append(score) + filter_boxes, filter_txts, filter_scores, filter_words = [], [], [], [] + for i, (box, txt, score) in enumerate( + zip(ocr_res.boxes, ocr_res.txts, ocr_res.scores) + ): + if score < self.text_score: + continue + + if ocr_res.word_results[i]: + filter_words.append(ocr_res.word_results[i]) + + filter_boxes.append(box) + filter_txts.append(txt) + filter_scores.append(score) ocr_res.boxes = np.array(filter_boxes) ocr_res.txts = tuple(filter_txts) ocr_res.scores = tuple(filter_scores) + ocr_res.word_results = tuple(filter_words) return ocr_res - def _add_letterbox( - self, img: np.ndarray, op_record: Dict[str, Any] - ) -> Tuple[np.ndarray, Dict[str, Any]]: - h, w = img.shape[:2] - - if self.width_height_ratio == -1: - use_limit_ratio = False - else: - use_limit_ratio = w / h > self.width_height_ratio - - if h <= self.min_height or use_limit_ratio: - padding_h = get_padding_h(h, w, self.width_height_ratio, self.min_height) - block_img = add_round_letterbox(img, (padding_h, padding_h, 0, 0)) - op_record["padding_1"] = {"top": padding_h, "left": 0} - return block_img, op_record - - op_record["padding_1"] = {"top": 0, "left": 0} - return img, op_record - - def _get_origin_points( - self, - dt_boxes: List[np.ndarray], - op_record: Dict[str, Any], - raw_h: int, - raw_w: int, - ) -> np.ndarray: - dt_boxes_array = np.array(dt_boxes).astype(np.float32) - for op in reversed(list(op_record.keys())): - v = op_record[op] - if "padding" in op: - top, left = v.get("top"), v.get("left") - dt_boxes_array[:, :, 0] -= left - dt_boxes_array[:, :, 1] -= top - elif "preprocess" in op: - ratio_h = v.get("ratio_h") - ratio_w = v.get("ratio_w") - dt_boxes_array[:, :, 0] *= ratio_w - dt_boxes_array[:, :, 1] *= ratio_h - - dt_boxes_array = np.where(dt_boxes_array < 0, 0, dt_boxes_array) - dt_boxes_array[..., 0] = np.where( - dt_boxes_array[..., 0] > raw_w, raw_w, dt_boxes_array[..., 0] - ) - dt_boxes_array[..., 1] = np.where( - dt_boxes_array[..., 1] > raw_h, raw_h, dt_boxes_array[..., 1] - ) - return dt_boxes_array - class RapidOCRError(Exception): pass diff --git a/python/rapidocr/utils/__init__.py b/python/rapidocr/utils/__init__.py index ff42bbd..0ecdd4f 100644 --- a/python/rapidocr/utils/__init__.py +++ b/python/rapidocr/utils/__init__.py @@ -1,14 +1,3 @@ # -*- encoding: utf-8 -*- # @Author: SWHL # @Contact: liekkaskono@163.com -from .download_file import DownloadFile, DownloadFileException, DownloadFileInput -from .load_image import LoadImage, LoadImageError -from .output import RapidOCROutput -from .parse_parameters import ParseParams -from .process_img import ( - add_round_letterbox, - get_padding_h, - get_rotate_crop_image, - resize_image_within_bounds, -) -from .vis_res import VisRes diff --git a/python/rapidocr/utils/process_img.py b/python/rapidocr/utils/process_img.py index 4d71249..46e8a16 100644 --- a/python/rapidocr/utils/process_img.py +++ b/python/rapidocr/utils/process_img.py @@ -1,12 +1,56 @@ # -*- encoding: utf-8 -*- # @Author: SWHL # @Contact: liekkaskono@163.com -from typing import Tuple +from typing import Any, Dict, Tuple import cv2 import numpy as np +def map_boxes_to_original( + dt_boxes: np.ndarray, op_record: Dict[str, Any], ori_h: int, ori_w: int +) -> np.ndarray: + for op in reversed(list(op_record.keys())): + v = op_record[op] + if "padding" in op: + top, left = v.get("top"), v.get("left") + dt_boxes[:, :, 0] -= left + dt_boxes[:, :, 1] -= top + elif "preprocess" in op: + ratio_h = v.get("ratio_h") + ratio_w = v.get("ratio_w") + dt_boxes[:, :, 0] *= ratio_w + dt_boxes[:, :, 1] *= ratio_h + + dt_boxes = np.where(dt_boxes < 0, 0, dt_boxes) + dt_boxes[..., 0] = np.where(dt_boxes[..., 0] > ori_w, ori_w, dt_boxes[..., 0]) + dt_boxes[..., 1] = np.where(dt_boxes[..., 1] > ori_h, ori_h, dt_boxes[..., 1]) + return dt_boxes + + +def apply_vertical_padding( + img: np.ndarray, + op_record: Dict[str, Any], + width_height_ratio: float, + min_height: float, +) -> Tuple[np.ndarray, Dict[str, Any]]: + h, w = img.shape[:2] + + if width_height_ratio == -1: + use_limit_ratio = False + else: + use_limit_ratio = w / h > width_height_ratio + + if h <= min_height or use_limit_ratio: + padding_h = get_padding_h(h, w, width_height_ratio, min_height) + block_img = add_round_letterbox(img, (padding_h, padding_h, 0, 0)) + op_record["padding_1"] = {"top": padding_h, "left": 0} + return block_img, op_record + + op_record["padding_1"] = {"top": 0, "left": 0} + return img, op_record + + def get_padding_h(h: int, w: int, width_height_ratio: float, min_height: float) -> int: new_h = max(int(w / width_height_ratio), min_height) * 2 padding_h = int(abs(new_h - h) / 2) diff --git a/python/rapidocr/utils/vis_res.py b/python/rapidocr/utils/vis_res.py index f6f3c22..f228897 100644 --- a/python/rapidocr/utils/vis_res.py +++ b/python/rapidocr/utils/vis_res.py @@ -31,8 +31,6 @@ class VisRes: lang_type: Optional[LangRec] = None, font_path: Optional[str] = None, ): - self.logger = logger - self.text_score = text_score self.load_img = LoadImage() @@ -44,9 +42,9 @@ class VisRes: img_content: InputType, dt_boxes: np.ndarray, txts: Optional[Union[List[str], Tuple[str]]] = None, - scores: Optional[Tuple[float]] = None, + scores: Optional[List[float]] = None, ) -> np.ndarray: - self.logger.info(f"Using {self.font_path} to visualize results.") + logger.info(f"Using {self.font_path} to visualize results.") if txts is None: return self.draw_dt_boxes(img_content, dt_boxes, scores) @@ -56,7 +54,7 @@ class VisRes: self, img_content: InputType, dt_boxes: np.ndarray, - scores: Optional[Tuple[float]] = None, + scores: Optional[List[float]] = None, ) -> np.ndarray: img = self.load_img(img_content) if scores is None: @@ -90,7 +88,7 @@ class VisRes: file_url=default_info["path"], sha256=default_info["SHA256"], save_path=DEFAULT_FONT_PATH, - logger=self.logger, + logger=logger, ) if lang_type is None: @@ -106,7 +104,7 @@ class VisRes: font_url, font_sha256 = font_info["path"], font_info["SHA256"] if font_url is None: - self.logger.warning( + logger.warning( "Font file for %s is not found in the supported font list. Default font file will be used.", lang_type, ) @@ -119,7 +117,7 @@ class VisRes: file_url=font_url, sha256=font_sha256, save_path=save_font_path, - logger=self.logger, + logger=logger, verbose=False, ) DownloadFile.run(input_param) @@ -131,7 +129,7 @@ class VisRes: self, imgs: Sequence[InputType], txts: Union[List[str], Tuple[str]], - scores: Tuple[float], + scores: List[float], ) -> np.ndarray: result_imgs = [] for img, txt, score in zip(imgs, txts, scores): @@ -178,7 +176,7 @@ class VisRes: img_content: InputType, dt_boxes: np.ndarray, txts: Union[List[str], Tuple[str]], - scores: Optional[Tuple[float]] = None, + scores: Optional[List[float]] = None, ) -> np.ndarray: image = Image.fromarray(self.load_img(img_content)) h, w = image.height, image.width diff --git a/python/tests/test_files/check_return_word_len.jpeg b/python/tests/test_files/check_return_word_len.jpeg new file mode 100644 index 0000000..583553c Binary files /dev/null and b/python/tests/test_files/check_return_word_len.jpeg differ diff --git a/python/tests/test_return_word.py b/python/tests/test_return_word.py index 0814c5e..1c8cc84 100644 --- a/python/tests/test_return_word.py +++ b/python/tests/test_return_word.py @@ -13,7 +13,7 @@ sys.path.append(str(root_dir)) from rapidocr import RapidOCR -tests_dir = root_dir / "tests" / "test_files" +test_dir = root_dir / "tests" / "test_files" @pytest.fixture() @@ -22,6 +22,13 @@ def engine(): return engine +def test_txts_equal_words(engine): + img_path = test_dir / "check_return_word_len.jpeg" + result = engine(img_path, return_word_box=True) + + assert len(result.txts) == len(result.word_results) + + @mark.parametrize( "img_name,words", [ @@ -36,7 +43,7 @@ def engine(): ], ) def test_cn_word_ocr(engine, img_name: str, words: List[str]): - img_path = tests_dir / img_name + img_path = test_dir / img_name result = engine(img_path, return_word_box=True) txts, _, _ = list(zip(*result.word_results[0])) assert txts == words @@ -47,14 +54,14 @@ def test_cn_word_ocr(engine, img_name: str, words: List[str]): [("issue_170.png", "TEST"), ("return_word_debug.jpg", "3F1")], ) def test_en_word_ocr(engine, img_name: str, words: str): - img_path = tests_dir / img_name + img_path = test_dir / img_name result = engine(img_path, return_word_box=True) txts, _, _ = list(zip(*result.word_results[0])) assert txts[0] == words def test_en_return_single_char_box(engine): - img_path = tests_dir / "en.jpg" + img_path = test_dir / "en.jpg" result = engine(img_path, return_word_box=True, return_single_char_box=True) txts, _, _ = list(zip(*result.word_results[0])) assert txts[0] == "3"