diff --git a/python/rapidocr/main.py b/python/rapidocr/main.py index ceadbbc..a2890dd 100644 --- a/python/rapidocr/main.py +++ b/python/rapidocr/main.py @@ -23,9 +23,11 @@ from .utils.process_img import ( apply_vertical_padding, get_rotate_crop_image, map_boxes_to_original, + map_img_to_original, resize_image_within_bounds, ) from .utils.typings import LangRec +from .utils.utils import filter_by_indices from .utils.vis_res import VisRes root_dir = Path(__file__).resolve().parent @@ -159,28 +161,27 @@ class RapidOCR: det_res.boxes, op_record, ori_h, ori_w ) + ratio_h = op_record["preprocess"]["ratio_h"] + ratio_w = op_record["preprocess"]["ratio_w"] + cropped_img_list = map_img_to_original(cropped_img_list, ratio_h, ratio_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()} + valid_ids = [i for i, v in enumerate(rec_res.txts) if 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 - ] + det_res.boxes = filter_by_indices(det_res.boxes, valid_ids) + det_res.scores = filter_by_indices(det_res.scores, valid_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 - ] - rec_res.scores = [ - v for i, v in enumerate(rec_res.scores) if i not in empty_ids - ] + rec_res.imgs = filter_by_indices(rec_res.imgs, valid_ids) + rec_res.txts = filter_by_indices(rec_res.txts, valid_ids) + rec_res.word_results = filter_by_indices(rec_res.word_results, valid_ids) + rec_res.scores = filter_by_indices(rec_res.scores, valid_ids) + + cropped_img_list = filter_by_indices(cropped_img_list, valid_ids) # 仅分类结果 if ( @@ -243,17 +244,9 @@ class RapidOCR: origin_words = [] for word_line in rec_res.word_results: - origin_words_item = [] - for txt, score, bbox in word_line: - if bbox is None: - continue - - 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)) - + origin_words_item = [ + (txt, score, bbox) for txt, score, bbox in word_line if bbox is not None + ] if origin_words_item: origin_words.append(tuple(origin_words_item)) return tuple(origin_words) diff --git a/python/rapidocr/utils/process_img.py b/python/rapidocr/utils/process_img.py index 46e8a16..13a2804 100644 --- a/python/rapidocr/utils/process_img.py +++ b/python/rapidocr/utils/process_img.py @@ -1,12 +1,24 @@ # -*- encoding: utf-8 -*- # @Author: SWHL # @Contact: liekkaskono@163.com -from typing import Any, Dict, Tuple +from typing import Any, Dict, List, Tuple import cv2 import numpy as np +def map_img_to_original( + imgs: List[np.ndarray], ratio_h: float, ratio_w: float +) -> List[np.ndarray]: + results = [] + for img in imgs: + img_h, img_w = img.shape[:2] + ori_img_h, ori_img_w = round(img_h * ratio_h), round(img_w * ratio_w) + resize_img = cv2.resize(img, (ori_img_w, ori_img_h)) + results.append(resize_img) + return results + + def map_boxes_to_original( dt_boxes: np.ndarray, op_record: Dict[str, Any], ori_h: int, ori_w: int ) -> np.ndarray: diff --git a/python/rapidocr/utils/utils.py b/python/rapidocr/utils/utils.py index c0f63a3..54ea641 100644 --- a/python/rapidocr/utils/utils.py +++ b/python/rapidocr/utils/utils.py @@ -5,13 +5,26 @@ import hashlib import importlib from pathlib import Path from sys import platform -from typing import Tuple, Union +from typing import Any, List, Tuple, Union from urllib.parse import urlparse import cv2 import numpy as np +def filter_by_indices( + data: Union[np.ndarray, List[Any], Tuple[Any]], + indices: Union[np.ndarray, List[int], Tuple[int, ...]], +) -> Union[np.ndarray, List[Any], Tuple[Any]]: + if isinstance(data, np.ndarray): + return data[indices] + + if isinstance(data, (list, tuple)): + return [data[i] for i in indices] + + raise TypeError(f"Unsupported data type: {type(data)}") + + def mkdir(dir_path): Path(dir_path).mkdir(parents=True, exist_ok=True) @@ -29,8 +42,16 @@ def quads_to_rect_bbox(bbox: np.ndarray) -> Tuple[float, float, float, float]: return float(x_min), float(y_min), float(x_max), float(y_max) +def is_chinese_char(ch: str) -> bool: + return ( + "\u4e00" <= ch <= "\u9fff" # 汉字 + or "\u3000" <= ch <= "\u303f" # CJK 标点(如 。 、 “” 《》 ……) + or "\uff00" <= ch <= "\uffef" # 全角符号(如 ,.!?【】) + ) + + def has_chinese_char(text: str) -> bool: - return any("\u4e00" <= ch <= "\u9fff" for ch in text) + return any(is_chinese_char(ch) for ch in text) def get_file_sha256(file_path: Union[str, Path], chunk_size: int = 65536) -> str: