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

@@ -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

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]],

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
 
 

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)

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

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
+onnxruntime>=1.7.0
+tqdm

python/tests/__init__.py

@@ -0,0 +1,3 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com

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

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