From fd520ffdad27cefa043884348240b7e18a67a71c Mon Sep 17 00:00:00 2001 From: SWHL Date: Thu, 17 Oct 2024 22:24:31 +0800 Subject: [PATCH] Fix issue 232 (#235) * fix(rapidocr_onnxruntime): Fixed padding_h variable before assignment * feat(rapidocr_onnxruntime): Add resize function * feat(rapidocr_onnxruntime): Add limit of max_side_len and min_side_len * tests: Add __init__ --- python/rapidocr_onnxruntime/config.yaml | 2 + python/rapidocr_onnxruntime/main.py | 84 ++++++++++++++--- python/rapidocr_onnxruntime/utils/__init__.py | 1 + .../utils/parse_parameters.py | 2 + .../rapidocr_onnxruntime/utils/process_img.py | 87 ++++++++++++++++++ python/requirements_ort.txt | 3 +- python/tests/__init__.py | 3 + python/tests/base_module.py | 23 +++++ python/tests/test_files/short.png | Bin 0 -> 1803 bytes python/tests/test_ort.py | 15 +++ 10 files changed, 206 insertions(+), 14 deletions(-) create mode 100644 python/rapidocr_onnxruntime/utils/process_img.py create mode 100644 python/tests/__init__.py create mode 100644 python/tests/test_files/short.png diff --git a/python/rapidocr_onnxruntime/config.yaml b/python/rapidocr_onnxruntime/config.yaml index 84425f1..f92d330 100644 --- a/python/rapidocr_onnxruntime/config.yaml +++ b/python/rapidocr_onnxruntime/config.yaml @@ -6,6 +6,8 @@ Global: print_verbose: false min_height: 30 width_height_ratio: 8 + max_side_len: 2000 + min_side_len: 30 intra_op_num_threads: &intra_nums -1 inter_op_num_threads: &inter_nums -1 diff --git a/python/rapidocr_onnxruntime/main.py b/python/rapidocr_onnxruntime/main.py index f53a7d7..43c4b93 100644 --- a/python/rapidocr_onnxruntime/main.py +++ b/python/rapidocr_onnxruntime/main.py @@ -3,7 +3,7 @@ # @Contact: liekkaskono@163.com import copy from pathlib import Path -from typing import Any, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Tuple, Union import cv2 import numpy as np @@ -15,9 +15,12 @@ from .utils import ( LoadImage, UpdateParameters, VisRes, + add_round_letterbox, get_logger, + increase_min_side, init_args, read_yaml, + reduce_max_side, update_model_path, ) @@ -54,6 +57,8 @@ class RapidOCR: self.text_rec = TextRecognizer(config["Rec"]) self.load_img = LoadImage() + self.max_side_len = global_config["max_side_len"] + self.min_side_len = global_config["min_side_len"] def __call__( self, @@ -78,11 +83,16 @@ class RapidOCR: img = self.load_img(img_content) + raw_h, raw_w = img.shape[:2] + op_record = {} + img, ratio_h, ratio_w = self.preprocess(img) + op_record["preprocess"] = {"ratio_h": ratio_h, "ratio_w": ratio_w} + dt_boxes, cls_res, rec_res = None, None, None det_elapse, cls_elapse, rec_elapse = 0.0, 0.0, 0.0 if use_det: - img, padding_h = self.maybe_add_letterbox(img) + img, op_record = self.maybe_add_letterbox(img, op_record) dt_boxes, det_elapse = self.auto_text_det(img) if dt_boxes is None: return None, None @@ -95,16 +105,30 @@ class RapidOCR: if use_rec: rec_res, rec_elapse = self.text_rec(img) - if dt_boxes is not None and padding_h > 0: - for box in dt_boxes: - box[:, 1] -= padding_h + if dt_boxes is not None and rec_res is not None: + dt_boxes = self._get_origin_points(dt_boxes, op_record, raw_h, raw_w) ocr_res = self.get_final_res( dt_boxes, cls_res, rec_res, det_elapse, cls_elapse, rec_elapse ) return ocr_res - def maybe_add_letterbox(self, img: np.ndarray) -> Tuple[np.ndarray, int]: + def preprocess(self, img: np.ndarray) -> Tuple[np.ndarray, float, float]: + h, w = img.shape[:2] + max_value = max(h, w) + ratio_h = ratio_w = 1.0 + if max_value > self.max_side_len: + img, ratio_h, ratio_w = reduce_max_side(img, self.max_side_len) + + h, w = img.shape[:2] + min_value = min(h, w) + if min_value < self.min_side_len: + img, ratio_h, ratio_w = increase_min_side(img, self.min_side_len) + return img, ratio_h, ratio_w + + def maybe_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: @@ -113,13 +137,18 @@ class RapidOCR: use_limit_ratio = w / h > self.width_height_ratio if h <= self.min_height or use_limit_ratio: - new_h = max(int(w / self.width_height_ratio), self.min_height) * 2 - padding_h = int(abs(new_h - h) / 2) - block_img = cv2.copyMakeBorder( - img, padding_h, padding_h, 0, 0, cv2.BORDER_CONSTANT, value=(0, 0, 0) - ) - return block_img, padding_h - return img, 0 + padding_h = self._get_padding_h(h, w) + 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(self, h: int, w: int) -> int: + new_h = max(int(w / self.width_height_ratio), self.min_height) * 2 + padding_h = int(abs(new_h - h) / 2) + return padding_h def auto_text_det( self, img: np.ndarray @@ -201,6 +230,35 @@ class RapidOCR: break return _boxes + 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) + for op in reversed(op_record): + 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 + def get_final_res( self, dt_boxes: Optional[List[np.ndarray]], diff --git a/python/rapidocr_onnxruntime/utils/__init__.py b/python/rapidocr_onnxruntime/utils/__init__.py index cf5ca69..150da63 100644 --- a/python/rapidocr_onnxruntime/utils/__init__.py +++ b/python/rapidocr_onnxruntime/utils/__init__.py @@ -10,6 +10,7 @@ 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 .process_img import add_round_letterbox, increase_min_side, reduce_max_side from .vis_res import VisRes diff --git a/python/rapidocr_onnxruntime/utils/parse_parameters.py b/python/rapidocr_onnxruntime/utils/parse_parameters.py index 7c39c45..527a5c9 100644 --- a/python/rapidocr_onnxruntime/utils/parse_parameters.py +++ b/python/rapidocr_onnxruntime/utils/parse_parameters.py @@ -35,6 +35,8 @@ def init_args(): 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("--max_side_len", type=int, default=2000) + global_group.add_argument("--min_side_len", type=int, default=30) global_group.add_argument("--intra_op_num_threads", type=int, default=-1) global_group.add_argument("--inter_op_num_threads", type=int, default=-1) diff --git a/python/rapidocr_onnxruntime/utils/process_img.py b/python/rapidocr_onnxruntime/utils/process_img.py new file mode 100644 index 0000000..7d3e194 --- /dev/null +++ b/python/rapidocr_onnxruntime/utils/process_img.py @@ -0,0 +1,87 @@ +# -*- encoding: utf-8 -*- +# @Author: SWHL +# @Contact: liekkaskono@163.com +from typing import Tuple + +import cv2 +import numpy as np + + +def reduce_max_side( + img: np.ndarray, max_side_len: int = 2000 +) -> Tuple[np.ndarray, float, float]: + h, w = img.shape[:2] + + ratio = 1.0 + if max(h, w) > max_side_len: + if h > w: + ratio = float(max_side_len) / h + else: + ratio = float(max_side_len) / w + + resize_h = int(h * ratio) + resize_w = int(w * ratio) + + resize_h = int(round(resize_h / 32) * 32) + resize_w = int(round(resize_w / 32) * 32) + + try: + if int(resize_w) <= 0 or int(resize_h) <= 0: + raise ResizeImgError("resize_w or resize_h is less than or equal to 0") + img = cv2.resize(img, (resize_w, resize_h)) + except Exception as exc: + raise ResizeImgError() from exc + + ratio_h = h / resize_h + ratio_w = w / resize_w + return img, ratio_h, ratio_w + + +def increase_min_side( + img: np.ndarray, min_side_len: int = 30 +) -> Tuple[np.ndarray, float, float]: + h, w = img.shape[:2] + + ratio = 1.0 + if min(h, w) < min_side_len: + if h < w: + ratio = float(min_side_len) / h + else: + ratio = float(min_side_len) / w + + resize_h = int(h * ratio) + resize_w = int(w * ratio) + + resize_h = int(round(resize_h / 32) * 32) + resize_w = int(round(resize_w / 32) * 32) + + try: + if int(resize_w) <= 0 or int(resize_h) <= 0: + raise ResizeImgError("resize_w or resize_h is less than or equal to 0") + img = cv2.resize(img, (resize_w, resize_h)) + except Exception as exc: + raise ResizeImgError() from exc + + ratio_h = h / resize_h + ratio_w = w / resize_w + return img, ratio_h, ratio_w + + +def add_round_letterbox( + img: np.ndarray, + padding_tuple: Tuple[int, int, int, int], +) -> np.ndarray: + padded_img = cv2.copyMakeBorder( + img, + padding_tuple[0], + padding_tuple[1], + padding_tuple[2], + padding_tuple[3], + cv2.BORDER_CONSTANT, + value=(0, 0, 0), + ) + return padded_img + + +class ResizeImgError(Exception): + pass diff --git a/python/requirements_ort.txt b/python/requirements_ort.txt index 96a8951..11139a3 100644 --- a/python/requirements_ort.txt +++ b/python/requirements_ort.txt @@ -5,4 +5,5 @@ six>=1.15.0 Shapely>=1.7.1,!=2.0.4 # python3.12 2.0.4 bug PyYAML Pillow -onnxruntime>=1.7.0 \ No newline at end of file +onnxruntime>=1.7.0 +tqdm \ No newline at end of file diff --git a/python/tests/__init__.py b/python/tests/__init__.py new file mode 100644 index 0000000..0ecdd4f --- /dev/null +++ b/python/tests/__init__.py @@ -0,0 +1,3 @@ +# -*- encoding: utf-8 -*- +# @Author: SWHL +# @Contact: liekkaskono@163.com diff --git a/python/tests/base_module.py b/python/tests/base_module.py index 73026d2..e762b8f 100644 --- a/python/tests/base_module.py +++ b/python/tests/base_module.py @@ -4,8 +4,11 @@ import importlib import sys from pathlib import Path +from typing import Union +import requests import yaml +from tqdm import tqdm class BaseModule: @@ -30,3 +33,23 @@ class BaseModule: with open(yaml_path, "rb") as f: data = yaml.load(f, Loader=yaml.Loader) return data + + +def download_file(url: str, save_path: Union[str, Path]): + response = requests.get(url, stream=True, timeout=60) + status_code = response.status_code + + if status_code != 200: + raise DownloadModelError("Something went wrong while downloading models") + + total_size_in_bytes = int(response.headers.get("content-length", 1)) + block_size = 1024 # 1 Kibibyte + with tqdm(total=total_size_in_bytes, unit="iB", unit_scale=True) as pb: + with open(save_path, "wb") as file: + for data in response.iter_content(block_size): + pb.update(len(data)) + file.write(data) + + +class DownloadModelError(Exception): + pass diff --git a/python/tests/test_files/short.png b/python/tests/test_files/short.png new file mode 100644 index 0000000000000000000000000000000000000000..744bf341b4380bbebc7ebc851431c09df9865e37 GIT binary patch literal 1803 zcmeAS@N?(olHy`uVBq!ia0vp^jzG-K!2~4Nhm}hJDaPU;cPEB*=VV?2IkQqDJkxxA z8MJ_G4hA+xAqHlU5+Gz?lwx27vl$q|7^UItI7SVq8lahsh72H+Ie|E>J(Gb2s=fe7 zgMbGR15JX^!FibliNz(UKtA&VMua^JnBe*x7cj%w+(2~sd^Q;5Swi(tboki)RIJnirk#MVyg;UC9n!B zAR8pCucQE0Qj%?}6yY17;GAESs$iyPpl6_D$EBd4U{jQmW)z9|8>y;bpKhp88yV>WRp=I1=9MH?=;jqGLkxkLMfuL^+7WFhI$72*t7vX3wIxiu8PD0h$gVF zAePyHjI#1C%1i}17vv~ALmRO1C}If5*#J$qaxO|uEXgkl$#WBQ#w{)tfeomr5TYmZM=QGWOiA%0XO;WT2bxbUwk{M>|GIbitF^1P zzsjEW{{4GX`5dL@)Q?7^FQ?Y z<3~XO0R@YPCdpHDWx03l*bxyBARs5FcYO2a&2#6@eeojW;ohA)Po6!y_4eBb+S=MX z4xVNX4-5+ ztgNho2hW`GsjB+*`0?WN&z1M=*kNI3Cs$u2xM|%!&`H64^u~ixy870I;Mcuk_BO@zo*8KVIBAYHx6+OAo z?ewu@Zd0-|G87E!9`M|HS#s_Ab?4*ZIWv=8Tv*mFk_oyP;pfMv^l<;4JwALTuin0m zjgDTuZe3c&gA2!w9SchJ_vg=_rLV8ge*fORWtE%c<>j4U6j(^`S*-3jeEYVv)08Pw zo;-be^3+A1C}sJ(K?YJ6IsduwpHO)0rq%Lrp=H6J4B4jvtxqd!YgOgum@^sh-cVd> zcMurZ&CSi5H*HEwPp_QZ8G2A&H*vxF=4R$oi%zz*uza|1A>h*{Rui@k#jQXUyJPfz zWF}-}WMpS|D?W~ojJ$dOKK}`IsdukmYwPN=nwNFb%OX`E>wuOT30psLCm`%-dFsG|JqhJ5@v?w(-Rc7MLCr?wyHko5PcUqdYOwnO5pR|`HT_NN0-MeQu?%X-EKu^`8XwvK^hZ(vL7nxseE}p2- zck!gwqw@k9t$C6hD_=^o$OxYFSO^S>HQK?df_IfvjnhjN)5C65Jj(y$$9U2xK*QWe z;GRIBUU5pmrj;V>ebMRZ>e-jCT{|YwvDB!b#-m}*k`!h^&a{GeJlC6xE>_(FKMFkAASuS|hQj=|{7M o=(e`(j!Rw3A3e%p`&;pk30U|pSh|(R98_g^y85}Sb4q9e08^j1xBvhE literal 0 HcmV?d00001 diff --git a/python/tests/test_ort.py b/python/tests/test_ort.py index 0814da3..9697075 100644 --- a/python/tests/test_ort.py +++ b/python/tests/test_ort.py @@ -13,6 +13,7 @@ root_dir = Path(__file__).resolve().parent.parent sys.path.append(str(root_dir)) from rapidocr_onnxruntime import LoadImageError, RapidOCR +from tests.base_module import download_file engine = RapidOCR() tests_dir = root_dir / "tests" / "test_files" @@ -20,6 +21,16 @@ img_path = tests_dir / "ch_en_num.jpg" package_name = "rapidocr_onnxruntime" +def test_long_img(): + img_url = "https://github.com/RapidAI/RapidOCR/releases/download/v1.1.0/long.jpeg" + img_path = tests_dir / "long.jpeg" + download_file(img_url, save_path=img_path) + result, _ = engine(img_path) + assert result is not None + assert len(result) == 55 + img_path.unlink() + + def test_ort_cuda_warning(caplog): engine = RapidOCR(det_use_cuda=True) caplog.set_level(logging.WARNING) @@ -111,6 +122,7 @@ def test_cls_rec(): img_path = tests_dir / "text_cls.jpg" result, _ = engine(img_path, use_det=False, use_cls=True, use_rec=True) + assert result is not None assert len(result) == 1 assert result[0][0] == "韩国小馆" @@ -119,6 +131,7 @@ def test_det_cls_rec(): img = cv2.imread(str(img_path)) result, _ = engine(img) + assert result is not None assert result[0][1] == "正品促销" assert len(result) == 18 @@ -193,6 +206,7 @@ def test_input_three_ndim_two_channel(): image_array = np.load(str(img_npy)) result, _ = engine(image_array) + assert result is not None assert len(result) == 1 assert result[0][1] == "TREND PLOT REPORT" @@ -205,5 +219,6 @@ def test_input_three_ndim_one_channel(): result, _ = engine(img) + assert result is not None assert result[0][1] == "正品促销" assert len(result) == 17