refactor(rapidocr_onnxruntime): Add typing annotations and remove invalid code

python/rapidocr_onnxruntime/ch_ppocr_v2_cls/text_cls.py

@@ -15,18 +15,17 @@ import argparse
 import copy
 import math
 import time
-from typing import List
+from typing import Any, Dict, List, Tuple, Union
 
 import cv2
 import numpy as np
-
 from rapidocr_onnxruntime.utils import OrtInferSession, read_yaml
 
 from .utils import ClsPostProcess
 
 
 class TextClassifier:
-    def __init__(self, config):
+    def __init__(self, config: Dict[str, Any]):
         self.cls_image_shape = config["cls_image_shape"]
         self.cls_batch_num = config["cls_batch_num"]
         self.cls_thresh = config["cls_thresh"]
@@ -34,7 +33,9 @@ class TextClassifier:
 
         self.infer = OrtInferSession(config)
 
-    def __call__(self, img_list: List[np.ndarray]):
+    def __call__(
+        self, img_list: Union[np.ndarray, List[np.ndarray]]
+    ) -> Tuple[List[np.ndarray], List[List[Union[str, float]]], float]:
         if isinstance(img_list, np.ndarray):
             img_list = [img_list]
 
@@ -65,8 +66,7 @@ class TextClassifier:
             cls_result = self.postprocess_op(prob_out)
             elapse += time.time() - starttime
 
-            for rno in range(len(cls_result)):
-                label, score = cls_result[rno]
+            for rno, (label, score) in enumerate(cls_result):
                 cls_res[indices[beg_img_no + rno]] = [label, score]
                 if "180" in label and score > self.cls_thresh:
                     img_list[indices[beg_img_no + rno]] = cv2.rotate(
@@ -74,7 +74,7 @@ class TextClassifier:
                     )
         return img_list, cls_res, elapse
 
-    def resize_norm_img(self, img):
+    def resize_norm_img(self, img: np.ndarray) -> np.ndarray:
         img_c, img_h, img_w = self.cls_image_shape
         h, w = img.shape[:2]
         ratio = w / float(h)

python/rapidocr_onnxruntime/ch_ppocr_v2_cls/utils.py

@@ -11,20 +11,18 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-class ClsPostProcess:
-    """Convert between text-label and text-index"""
+from typing import List, Tuple
 
-    def __init__(self, label_list):
-        super(ClsPostProcess, self).__init__()
+import numpy as np
+
+
+class ClsPostProcess:
+    def __init__(self, label_list: List[str]):
         self.label_list = label_list
 
-    def __call__(self, preds, label=None):
+    def __call__(self, preds: np.ndarray) -> List[Tuple[str, float]]:
         pred_idxs = preds.argmax(axis=1)
         decode_out = [
             (self.label_list[idx], preds[i, idx]) for i, idx in enumerate(pred_idxs)
         ]
-        if label is None:
-            return decode_out
-
-        label = [(self.label_list[idx], 1.0) for idx in label]
-        return decode_out, label
+        return decode_out

python/rapidocr_onnxruntime/ch_ppocr_v3_det/text_detect.py

@@ -14,19 +14,17 @@
 # -*- encoding: utf-8 -*-
 # @Author: SWHL
 # @Contact: liekkaskono@163.com
-import argparse
 import time
+from typing import Any, Dict, Optional, Tuple
 
-import cv2
 import numpy as np
-
-from rapidocr_onnxruntime.utils import OrtInferSession, read_yaml
+from rapidocr_onnxruntime.utils import OrtInferSession
 
 from .utils import DBPostProcess, create_operators, transform
 
 
 class TextDetector:
-    def __init__(self, config):
+    def __init__(self, config: Dict[str, Any]):
         pre_process_list = {
             "DetResizeForTest": {
                 "limit_side_len": config.get("limit_side_len", 736),
@@ -55,31 +53,50 @@ class TextDetector:
 
         self.infer = OrtInferSession(config)
 
-    def __call__(self, img):
+    def __call__(self, img: np.ndarray) -> Tuple[Optional[np.ndarray], float]:
         if img is None:
             raise ValueError("img is None")
 
-        ori_im_shape = img.shape[:2]
+        ori_im_shape = img.shape[0], img.shape[1]
 
-        data = {"image": img}
-        data = transform(data, self.preprocess_op)
-        img, shape_list = data
+        data = transform({"image": img}, self.preprocess_op)
+        if data is None:
+            return None, 0
+
+        img, shape = data
         if img is None:
             return None, 0
 
         img = np.expand_dims(img, axis=0).astype(np.float32)
-        shape_list = np.expand_dims(shape_list, axis=0)
+        shape = np.expand_dims(shape, axis=0)
 
         starttime = time.time()
         preds = self.infer(img)[0]
-        post_result = self.postprocess_op(preds, shape_list)
+        post_result = self.postprocess_op(preds, shape)
 
         dt_boxes = post_result[0]["points"]
         dt_boxes = self.filter_tag_det_res(dt_boxes, ori_im_shape)
         elapse = time.time() - starttime
         return dt_boxes, elapse
 
-    def order_points_clockwise(self, pts):
+    def filter_tag_det_res(
+        self, dt_boxes: np.ndarray, image_shape: Tuple[int, int]
+    ) -> np.ndarray:
+        img_height, img_width = image_shape
+        dt_boxes_new = []
+        for box in dt_boxes:
+            box = self.order_points_clockwise(box)
+            box = self.clip_det_res(box, img_height, img_width)
+
+            rect_width = int(np.linalg.norm(box[0] - box[1]))
+            rect_height = int(np.linalg.norm(box[0] - box[3]))
+            if rect_width <= 3 or rect_height <= 3:
+                continue
+
+            dt_boxes_new.append(box)
+        return np.array(dt_boxes_new)
+
+    def order_points_clockwise(self, pts: np.ndarray) -> np.ndarray:
         """
         reference from:
         https://github.com/jrosebr1/imutils/blob/master/imutils/perspective.py
@@ -104,42 +121,10 @@ class TextDetector:
         rect = np.array([tl, tr, br, bl], dtype="float32")
         return rect
 
-    def clip_det_res(self, points, img_height, img_width):
+    def clip_det_res(
+        self, points: np.ndarray, img_height: int, img_width: int
+    ) -> np.ndarray:
         for pno in range(points.shape[0]):
             points[pno, 0] = int(min(max(points[pno, 0], 0), img_width - 1))
             points[pno, 1] = int(min(max(points[pno, 1], 0), img_height - 1))
         return points
-
-    def filter_tag_det_res(self, dt_boxes, image_shape):
-        img_height, img_width = image_shape[:2]
-        dt_boxes_new = []
-        for box in dt_boxes:
-            box = self.order_points_clockwise(box)
-            box = self.clip_det_res(box, img_height, img_width)
-            rect_width = int(np.linalg.norm(box[0] - box[1]))
-            rect_height = int(np.linalg.norm(box[0] - box[3]))
-            if rect_width <= 3 or rect_height <= 3:
-                continue
-            dt_boxes_new.append(box)
-        dt_boxes = np.array(dt_boxes_new)
-        return dt_boxes
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--config_path", type=str, default="config.yaml")
-    parser.add_argument("--image_path", type=str, default=None)
-    args = parser.parse_args()
-
-    config = read_yaml(args.config_path)
-
-    text_detector = TextDetector(config)
-
-    img = cv2.imread(args.image_path)
-    dt_boxes, elapse = text_detector(img)
-
-    from utils import draw_text_det_res
-
-    src_im = draw_text_det_res(dt_boxes, args.image_path)
-    cv2.imwrite("det_results.jpg", src_im)
-    print("The det_results.jpg has been saved in the current directory.")

python/rapidocr_onnxruntime/ch_ppocr_v3_det/utils.py

@@ -1,4 +1,3 @@
-"""
 # Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -12,11 +11,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""
 # -*- encoding: utf-8 -*-
 # @Author: SWHL
 # @Contact: liekkaskono@163.com
-import sys
+import importlib
+from typing import Any, Dict, List, Optional, Tuple
 
 import cv2
 import numpy as np
@@ -24,48 +23,30 @@ import pyclipper
 from shapely.geometry import Polygon
 
 
-class NormalizeImage:
-    """normalize image such as substract mean, divide std"""
+def create_operators(op_param_dict: Dict[str, Optional[Dict[str, Any]]]) -> List[Any]:
+    ops = []
+    cur_module = importlib.import_module("rapidocr_onnxruntime.ch_ppocr_v3_det.utils")
+    for op_name, param in op_param_dict.items():
+        if param is None:
+            param = {}
 
-    def __init__(self, scale=None, mean=None, std=None, order="chw"):
-        if isinstance(scale, str):
-            scale = eval(scale)
-
-        self.scale = np.float32(scale if scale is not None else 1.0 / 255.0)
-        mean = mean if mean is not None else [0.485, 0.456, 0.406]
-        std = std if std is not None else [0.229, 0.224, 0.225]
-
-        shape = (3, 1, 1) if order == "chw" else (1, 1, 3)
-        self.mean = np.array(mean).reshape(shape).astype("float32")
-        self.std = np.array(std).reshape(shape).astype("float32")
-
-    def __call__(self, data):
-        img = np.array(data["image"]).astype(np.float32)
-        data["image"] = (img * self.scale - self.mean) / self.std
-        return data
+        process_class = getattr(cur_module, op_name)
+        op = process_class(**param)
+        ops.append(op)
+    return ops
 
 
-class ToCHWImage:
-    """convert hwc image to chw image"""
+def transform(
+    data: Dict[str, np.ndarray], ops: Optional[List[Any]] = None
+) -> Optional[Tuple[np.ndarray, np.ndarray]]:
+    if ops is None:
+        ops = []
 
-    def __init__(self):
-        pass
-
-    def __call__(self, data):
-        img = np.array(data["image"])
-        data["image"] = img.transpose((2, 0, 1))
-        return data
-
-
-class KeepKeys:
-    def __init__(self, keep_keys):
-        self.keep_keys = keep_keys
-
-    def __call__(self, data):
-        data_list = []
-        for key in self.keep_keys:
-            data_list.append(data[key])
-        return data_list
+    for op in ops:
+        data = op(data)
+        if data is None:
+            return None
+    return data
 
 
 class DetResizeForTest:
@@ -85,36 +66,28 @@ class DetResizeForTest:
             self.limit_side_len = kwargs.get("limit_side_len", 736)
             self.limit_type = kwargs.get("limit_type", "min")
 
-    def __call__(self, data):
+    def __call__(self, data: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
         img = data["image"]
         src_h, src_w = img.shape[:2]
 
         if self.resize_type == 0:
-            img, [ratio_h, ratio_w] = self.resize_image_type0(img)
+            img = self.resize_image_type0(img)
         elif self.resize_type == 2:
-            img, [ratio_h, ratio_w] = self.resize_image_type2(img)
+            img = self.resize_image_type2(img)
         else:
-            img, [ratio_h, ratio_w] = self.resize_image_type1(img)
+            img = self.resize_image_type1(img)
 
-        data["image"] = img
-        data["shape"] = np.array([src_h, src_w, ratio_h, ratio_w])
-        return data
+        new_data = {"image": img, "shape": np.array([src_h, src_w])}
+        return new_data
 
-    def resize_image_type1(self, img):
+    def resize_image_type1(self, img: np.ndarray) -> np.ndarray:
         resize_h, resize_w = self.image_shape
-        ori_h, ori_w = img.shape[:2]  # (h, w, c)
-        ratio_h = float(resize_h) / ori_h
-        ratio_w = float(resize_w) / ori_w
         img = cv2.resize(img, (int(resize_w), int(resize_h)))
-        return img, [ratio_h, ratio_w]
+        return img
 
-    def resize_image_type0(self, img):
+    def resize_image_type0(self, img: np.ndarray) -> Optional[np.ndarray]:
         """
         resize image to a size multiple of 32 which is required by the network
-        args:
-            img(array): array with shape [h, w, c]
-        return(tuple):
-            img, (ratio_h, ratio_w)
         """
         limit_side_len = self.limit_side_len
         h, w = img.shape[:2]
@@ -145,21 +118,16 @@ class DetResizeForTest:
 
         try:
             if int(resize_w) <= 0 or int(resize_h) <= 0:
-                return None, (None, None)
+                return None
             img = cv2.resize(img, (int(resize_w), int(resize_h)))
-        except:
-            print(img.shape, resize_w, resize_h)
-            sys.exit(0)
+        except Exception as exc:
+            raise ResizeImgError from exc
 
-        ratio_h = resize_h / float(h)
-        ratio_w = resize_w / float(w)
-        return img, [ratio_h, ratio_w]
+        return img
 
-    def resize_image_type2(self, img):
+    def resize_image_type2(self, img: np.ndarray) -> np.ndarray:
         h, w = img.shape[:2]
-
-        resize_w = w
-        resize_h = h
+        resize_h, resize_w = h, w
 
         # Fix the longer side
         if resize_h > resize_w:
@@ -174,43 +142,61 @@ class DetResizeForTest:
         resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
         resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
         img = cv2.resize(img, (int(resize_w), int(resize_h)))
-        ratio_h = resize_h / float(h)
-        ratio_w = resize_w / float(w)
-
-        return img, [ratio_h, ratio_w]
+        return img
 
 
-def transform(data, ops=None):
-    """transform"""
-    if ops is None:
-        ops = []
+class NormalizeImage:
+    """normalize image such as substract mean, divide std"""
 
-    for op in ops:
-        data = op(data)
-        if data is None:
-            return None
-    return data
+    def __init__(
+        self,
+        scale: Optional[str] = None,
+        mean: Optional[List[float]] = None,
+        std: Optional[List[float]] = None,
+        order: str = "chw",
+    ):
+        if isinstance(scale, str):
+            scale = eval(scale)
+
+        self.scale = np.float32(scale if scale is not None else 1.0 / 255.0)
+        mean = mean if mean is not None else [0.485, 0.456, 0.406]
+        std = std if std is not None else [0.229, 0.224, 0.225]
+
+        shape = (3, 1, 1) if order == "chw" else (1, 1, 3)
+        self.mean = np.array(mean).reshape(shape).astype("float32")
+        self.std = np.array(std).reshape(shape).astype("float32")
+
+    def __call__(self, data: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+        img = np.array(data["image"]).astype(np.float32)
+        data["image"] = (img * self.scale - self.mean) / self.std
+        return data
 
 
-def create_operators(op_param_dict):
-    """
-    create operators based on the config
-    """
-    ops = []
-    for op_name, param in op_param_dict.items():
-        if param is None:
-            param = {}
-        op = eval(op_name)(**param)
-        ops.append(op)
-    return ops
+class ToCHWImage:
+    """convert hwc image to chw image"""
+
+    def __init__(self):
+        pass
+
+    def __call__(self, data: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+        img = np.array(data["image"])
+        data["image"] = img.transpose((2, 0, 1))
+        return data
 
 
-def draw_text_det_res(dt_boxes, img_path):
-    src_im = cv2.imread(img_path)
-    for box in dt_boxes:
-        box = np.array(box).astype(np.int32).reshape(-1, 2)
-        cv2.polylines(src_im, [box], True, color=(255, 255, 0), thickness=2)
-    return src_im
+class KeepKeys:
+    def __init__(self, keep_keys: List[str]):
+        self.keep_keys = keep_keys
+
+    def __call__(self, data: Dict[str, np.ndarray]) -> List[np.ndarray]:
+        data_list = []
+        for key in self.keep_keys:
+            data_list.append(data[key])
+        return data_list
+
+
+class ResizeImgError(Exception):
+    pass
 
 
 class DBPostProcess:
@@ -218,12 +204,12 @@ class DBPostProcess:
 
     def __init__(
         self,
-        thresh=0.3,
-        box_thresh=0.7,
-        max_candidates=1000,
-        unclip_ratio=2.0,
-        score_mode="fast",
-        use_dilation=False,
+        thresh: float = 0.3,
+        box_thresh: float = 0.7,
+        max_candidates: int = 1000,
+        unclip_ratio: float = 2.0,
+        score_mode: str = "fast",
+        use_dilation: bool = False,
     ):
         self.thresh = thresh
         self.box_thresh = box_thresh
@@ -232,18 +218,42 @@ class DBPostProcess:
         self.min_size = 3
         self.score_mode = score_mode
 
+        self.dilation_kernel = None
         if use_dilation:
             self.dilation_kernel = np.array([[1, 1], [1, 1]])
-        else:
-            self.dilation_kernel = None
 
-    def boxes_from_bitmap(self, pred, _bitmap, dest_width, dest_height):
+    def __call__(
+        self, pred: np.ndarray, shape: np.ndarray
+    ) -> List[Dict[str, np.ndarray]]:
+        pred = pred[:, 0, :, :]
+        segmentation = pred > self.thresh
+
+        boxes_batch = []
+        for batch_index in range(pred.shape[0]):
+            src_h, src_w = shape[batch_index]
+            if self.dilation_kernel is None:
+                mask = segmentation[batch_index]
+            else:
+                mask = cv2.dilate(
+                    np.array(segmentation[batch_index]).astype(np.uint8),
+                    self.dilation_kernel,
+                )
+
+            boxes, scores = self.boxes_from_bitmap(
+                pred[batch_index], mask, src_w, src_h
+            )
+
+            boxes_batch.append({"points": boxes})
+        return boxes_batch
+
+    def boxes_from_bitmap(
+        self, pred: np.ndarray, bitmap: np.ndarray, dest_width: int, dest_height: int
+    ) -> Tuple[np.ndarray, List[float]]:
         """
-        _bitmap: single map with shape (1, H, W),
+        bitmap: single map with shape (1, H, W),
                 whose values are binarized as {0, 1}
         """
 
-        bitmap = _bitmap
         height, width = bitmap.shape
 
         outs = cv2.findContours(
@@ -256,45 +266,35 @@ class DBPostProcess:
 
         num_contours = min(len(contours), self.max_candidates)
 
-        boxes = []
-        scores = []
+        boxes, scores = [], []
         for index in range(num_contours):
             contour = contours[index]
             points, sside = self.get_mini_boxes(contour)
             if sside < self.min_size:
                 continue
-            points = np.array(points)
+
             if self.score_mode == "fast":
                 score = self.box_score_fast(pred, points.reshape(-1, 2))
             else:
                 score = self.box_score_slow(pred, contour)
+
             if self.box_thresh > score:
                 continue
 
-            box = self.unclip(points).reshape(-1, 1, 2)
+            box = self.unclip(points)
             box, sside = self.get_mini_boxes(box)
             if sside < self.min_size + 2:
                 continue
-            box = np.array(box)
 
             box[:, 0] = np.clip(np.round(box[:, 0] / width * dest_width), 0, dest_width)
             box[:, 1] = np.clip(
                 np.round(box[:, 1] / height * dest_height), 0, dest_height
             )
-            boxes.append(box.astype(np.int16))
+            boxes.append(box.astype(np.int32))
             scores.append(score)
-        return np.array(boxes, dtype=np.int16), scores
+        return np.array(boxes, dtype=np.int32), scores
 
-    def unclip(self, box):
-        unclip_ratio = self.unclip_ratio
-        poly = Polygon(box)
-        distance = poly.area * unclip_ratio / poly.length
-        offset = pyclipper.PyclipperOffset()
-        offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
-        expanded = np.array(offset.Execute(distance))
-        return expanded
-
-    def get_mini_boxes(self, contour):
+    def get_mini_boxes(self, contour: np.ndarray) -> Tuple[np.ndarray, float]:
         bounding_box = cv2.minAreaRect(contour)
         points = sorted(list(cv2.boxPoints(bounding_box)), key=lambda x: x[0])
 
@@ -305,6 +305,7 @@ class DBPostProcess:
         else:
             index_1 = 1
             index_4 = 0
+
         if points[3][1] > points[2][1]:
             index_2 = 2
             index_3 = 3
@@ -312,10 +313,13 @@ class DBPostProcess:
             index_2 = 3
             index_3 = 2
 
-        box = [points[index_1], points[index_2], points[index_3], points[index_4]]
+        box = np.array(
+            [points[index_1], points[index_2], points[index_3], points[index_4]]
+        )
         return box, min(bounding_box[1])
 
-    def box_score_fast(self, bitmap, _box):
+    @staticmethod
+    def box_score_fast(bitmap: np.ndarray, _box: np.ndarray) -> float:
         h, w = bitmap.shape[:2]
         box = _box.copy()
         xmin = np.clip(np.floor(box[:, 0].min()).astype(np.int32), 0, w - 1)
@@ -329,10 +333,8 @@ class DBPostProcess:
         cv2.fillPoly(mask, box.reshape(1, -1, 2).astype(np.int32), 1)
         return cv2.mean(bitmap[ymin : ymax + 1, xmin : xmax + 1], mask)[0]
 
-    def box_score_slow(self, bitmap, contour):
-        """
-        box_score_slow: use polyon mean score as the mean score
-        """
+    def box_score_slow(self, bitmap: np.ndarray, contour: np.ndarray) -> float:
+        """use polyon mean score as the mean score"""
         h, w = bitmap.shape[:2]
         contour = contour.copy()
         contour = np.reshape(contour, (-1, 2))
@@ -350,23 +352,11 @@ class DBPostProcess:
         cv2.fillPoly(mask, contour.reshape(1, -1, 2).astype(np.int32), 1)
         return cv2.mean(bitmap[ymin : ymax + 1, xmin : xmax + 1], mask)[0]
 
-    def __call__(self, pred, shape_list):
-        pred = pred[:, 0, :, :]
-        segmentation = pred > self.thresh
-
-        boxes_batch = []
-        for batch_index in range(pred.shape[0]):
-            src_h, src_w, ratio_h, ratio_w = shape_list[batch_index]
-            if self.dilation_kernel is not None:
-                mask = cv2.dilate(
-                    np.array(segmentation[batch_index]).astype(np.uint8),
-                    self.dilation_kernel,
-                )
-            else:
-                mask = segmentation[batch_index]
-            boxes, scores = self.boxes_from_bitmap(
-                pred[batch_index], mask, src_w, src_h
-            )
-
-            boxes_batch.append({"points": boxes})
-        return boxes_batch
+    def unclip(self, box: np.ndarray) -> np.ndarray:
+        unclip_ratio = self.unclip_ratio
+        poly = Polygon(box)
+        distance = poly.area * unclip_ratio / poly.length
+        offset = pyclipper.PyclipperOffset()
+        offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
+        expanded = np.array(offset.Execute(distance)).reshape((-1, 1, 2))
+        return expanded

python/rapidocr_onnxruntime/ch_ppocr_v3_rec/text_recognize.py

@@ -14,30 +14,34 @@
 import argparse
 import math
 import time
-from typing import List
+from typing import Any, Dict, List, Tuple, Union
 
 import cv2
 import numpy as np
-
 from rapidocr_onnxruntime.utils import OrtInferSession, read_yaml
 
 from .utils import CTCLabelDecode
 
 
 class TextRecognizer:
-    def __init__(self, config):
+    def __init__(self, config: Dict[str, Any]):
         self.session = OrtInferSession(config)
 
+        character = None
         if self.session.have_key():
-            character_dict_path = self.session.get_character_list()
-        else:
-            character_dict_path = config.get("rec_keys_path", None)
-        self.postprocess_op = CTCLabelDecode(character_dict_path)
+            character = self.session.get_character_list()
+
+        character_path = config.get("rec_keys_path", None)
+        self.postprocess_op = CTCLabelDecode(
+            character=character, character_path=character_path
+        )
 
         self.rec_batch_num = config["rec_batch_num"]
         self.rec_image_shape = config["rec_img_shape"]
 
-    def __call__(self, img_list: List[np.ndarray]):
+    def __call__(
+        self, img_list: Union[np.ndarray, List[np.ndarray]]
+    ) -> Tuple[List[Tuple[str, float]], float]:
         if isinstance(img_list, np.ndarray):
             img_list = [img_list]
 
@@ -48,7 +52,7 @@ class TextRecognizer:
         indices = np.argsort(np.array(width_list))
 
         img_num = len(img_list)
-        rec_res = [["", 0.0]] * img_num
+        rec_res = [("", 0.0)] * img_num
 
         batch_num = self.rec_batch_num
         elapse = 0
@@ -75,7 +79,7 @@ class TextRecognizer:
             elapse += time.time() - starttime
         return rec_res, elapse
 
-    def resize_norm_img(self, img, max_wh_ratio):
+    def resize_norm_img(self, img: np.ndarray, max_wh_ratio: float) -> np.ndarray:
         img_channel, img_height, img_width = self.rec_image_shape
         assert img_channel == img.shape[2]
 

python/rapidocr_onnxruntime/ch_ppocr_v3_rec/utils.py

@@ -1,73 +1,101 @@
 # -*- encoding: utf-8 -*-
 # @Author: SWHL
 # @Contact: liekkaskono@163.com
+from pathlib import Path
+from typing import List, Optional, Tuple, Union
+
 import numpy as np
 
 
 class CTCLabelDecode:
-    """Convert between text-label and text-index"""
+    def __init__(
+        self,
+        character: Optional[List[str]] = None,
+        character_path: Union[str, Path, None] = None,
+    ):
+        self.character = self.get_character(character, character_path)
+        self.dict = {char: i for i, char in enumerate(self.character)}
 
-    def __init__(self, character_dict_path):
-        self.character_str = []
-        assert character_dict_path is not None, "character_dict_path should not be None"
-
-        if isinstance(character_dict_path, str):
-            with open(character_dict_path, "rb") as fin:
-                lines = fin.readlines()
-                for line in lines:
-                    line = line.decode("utf-8").strip("\n").strip("\r\n")
-                    self.character_str.append(line)
-        else:
-            self.character_str = character_dict_path
-        self.character_str.append(" ")
-
-        dict_character = self.add_special_char(self.character_str)
-        self.character = dict_character
-
-        self.dict = {}
-        for i, char in enumerate(dict_character):
-            self.dict[char] = i
-
-    def __call__(self, preds, label=None):
+    def __call__(self, preds: np.ndarray) -> List[Tuple[str, float]]:
         preds_idx = preds.argmax(axis=2)
         preds_prob = preds.max(axis=2)
         text = self.decode(preds_idx, preds_prob, is_remove_duplicate=True)
-        if label is None:
-            return text
-        label = self.decode(label)
-        return text, label
+        return text
 
-    def add_special_char(self, dict_character):
-        dict_character = ["blank"] + dict_character
-        return dict_character
+    def get_character(
+        self,
+        character: Optional[List[str]] = None,
+        character_path: Union[str, Path, None] = None,
+    ) -> List[str]:
+        if character is None and character_path is None:
+            raise ValueError("character must not be None")
 
-    def get_ignored_tokens(self):
-        return [0]  # for ctc blank
+        character_list = None
+        if character:
+            character_list = character
 
-    def decode(self, text_index, text_prob=None, is_remove_duplicate=False):
+        if character_path:
+            character_list = self.read_character_file(character_path)
+
+        if character_list is None:
+            raise ValueError("character must not be None")
+
+        character_list = self.insert_special_char(
+            character_list, " ", len(character_list)
+        )
+        character_list = self.insert_special_char(character_list, "blank", 0)
+        return character_list
+
+    @staticmethod
+    def read_character_file(character_path: Union[str, Path]) -> List[str]:
+        character_list = []
+        with open(character_path, "rb") as f:
+            lines = f.readlines()
+            for line in lines:
+                line = line.decode("utf-8").strip("\n").strip("\r\n")
+                character_list.append(line)
+        return character_list
+
+    @staticmethod
+    def insert_special_char(
+        character_list: List[str], special_char: str, loc: int = -1
+    ) -> List[str]:
+        character_list.insert(loc, special_char)
+        return character_list
+
+    def decode(
+        self,
+        text_index: np.ndarray,
+        text_prob: Optional[np.ndarray] = None,
+        is_remove_duplicate: bool = False,
+    ) -> List[Tuple[str, float]]:
         """convert text-index into text-label."""
-
         result_list = []
         ignored_tokens = self.get_ignored_tokens()
         batch_size = len(text_index)
         for batch_idx in range(batch_size):
-            char_list = []
-            conf_list = []
-            for idx in range(len(text_index[batch_idx])):
-                if text_index[batch_idx][idx] in ignored_tokens:
+            char_list, conf_list = [], []
+            cur_pred_ids = text_index[batch_idx]
+            for idx, cur_idx in enumerate(cur_pred_ids):
+                if cur_idx in ignored_tokens:
                     continue
+
                 if is_remove_duplicate:
                     # only for predict
-                    if (
-                        idx > 0
-                        and text_index[batch_idx][idx - 1] == text_index[batch_idx][idx]
-                    ):
+                    if idx > 0 and cur_pred_ids[idx - 1] == cur_idx:
                         continue
-                char_list.append(self.character[int(text_index[batch_idx][idx])])
-                if text_prob is not None:
-                    conf_list.append(text_prob[batch_idx][idx])
-                else:
+
+                char_list.append(self.character[int(cur_idx)])
+
+                if text_prob is None:
                     conf_list.append(1)
+                else:
+                    conf_list.append(text_prob[batch_idx][idx])
+
             text = "".join(char_list)
             result_list.append((text, np.mean(conf_list if any(conf_list) else [0])))
         return result_list
+
+    @staticmethod
+    def get_ignored_tokens() -> List[int]:
+        return [0]  # for ctc blank

python/rapidocr_onnxruntime/main.py

@@ -3,7 +3,7 @@
 # @Contact: liekkaskono@163.com
 import copy
 from pathlib import Path
-from typing import List, Optional, Tuple, Union
+from typing import Any, List, Optional, Tuple, Union
 
 import cv2
 import numpy as np
@@ -124,7 +124,7 @@ class RapidOCR:
     def auto_text_det(
         self,
         img: np.ndarray,
-    ) -> Tuple[Optional[np.ndarray], float, Optional[List[np.ndarray]]]:
+    ) -> Tuple[Optional[List[np.ndarray]], float]:
         dt_boxes, det_elapse = self.text_det(img)
         if dt_boxes is None or len(dt_boxes) < 1:
             return None, 0.0
@@ -132,8 +132,10 @@ class RapidOCR:
         dt_boxes = self.sorted_boxes(dt_boxes)
         return dt_boxes, det_elapse
 
-    def get_crop_img_list(self, img, dt_boxes):
-        def get_rotate_crop_image(img, points):
+    def get_crop_img_list(
+        self, img: np.ndarray, dt_boxes: List[np.ndarray]
+    ) -> List[np.ndarray]:
+        def get_rotate_crop_image(img: np.ndarray, points: np.ndarray) -> np.ndarray:
             img_crop_width = int(
                 max(
                     np.linalg.norm(points[0] - points[1]),
@@ -146,14 +148,14 @@ class RapidOCR:
                     np.linalg.norm(points[1] - points[2]),
                 )
             )
-            pts_std = np.float32(
+            pts_std = np.array(
                 [
                     [0, 0],
                     [img_crop_width, 0],
                     [img_crop_width, img_crop_height],
                     [0, img_crop_height],
                 ]
-            )
+            ).astype(np.float32)
             M = cv2.getPerspectiveTransform(points, pts_std)
             dst_img = cv2.warpPerspective(
                 img,
@@ -175,7 +177,7 @@ class RapidOCR:
         return img_crop_list
 
     @staticmethod
-    def sorted_boxes(dt_boxes):
+    def sorted_boxes(dt_boxes: np.ndarray) -> List[np.ndarray]:
         """
         Sort text boxes in order from top to bottom, left to right
         args:
@@ -201,8 +203,14 @@ class RapidOCR:
         return _boxes
 
     def get_final_res(
-        self, dt_boxes, cls_res, rec_res, det_elapse, cls_elapse, rec_elapse
-    ):
+        self,
+        dt_boxes: Optional[List[np.ndarray]],
+        cls_res: Optional[List[List[Union[str, float]]]],
+        rec_res: Optional[List[Tuple[str, float]]],
+        det_elapse: float,
+        cls_elapse: float,
+        rec_elapse: float,
+    ) -> Tuple[Optional[List[List[Union[Any, str]]]], Optional[List[float]]]:
         if dt_boxes is None and rec_res is None and cls_res is not None:
             return cls_res, [cls_elapse]
 
@@ -216,7 +224,7 @@ class RapidOCR:
             return [box.tolist() for box in dt_boxes], [det_elapse]
 
         dt_boxes, rec_res = self.filter_result(dt_boxes, rec_res)
-        if len(dt_boxes) <= 0:
+        if not dt_boxes or not rec_res or len(dt_boxes) <= 0:
             return None, None
 
         ocr_res = [
@@ -224,7 +232,14 @@ class RapidOCR:
         ], [det_elapse, cls_elapse, rec_elapse]
         return ocr_res
 
-    def filter_result(self, dt_boxes, rec_res):
+    def filter_result(
+        self,
+        dt_boxes: Optional[List[np.ndarray]],
+        rec_res: Optional[List[Tuple[str, float]]],
+    ) -> Tuple[Optional[List[np.ndarray]], Optional[List[Tuple[str, float]]]]:
+        if dt_boxes is None or rec_res is None:
+            return None, None
+
         filter_boxes, filter_rec_res = [], []
         for box, rec_reuslt in zip(dt_boxes, rec_res):
             text, score = rec_reuslt

python/rapidocr_onnxruntime/utils/infer_engine.py

@@ -6,7 +6,7 @@ import platform
 import traceback
 from enum import Enum
 from pathlib import Path
-from typing import List, Tuple, Union
+from typing import Any, Dict, List, Tuple, Union
 
 import numpy as np
 from onnxruntime import (
@@ -27,7 +27,7 @@ class EP(Enum):
 
 
 class OrtInferSession:
-    def __init__(self, config):
+    def __init__(self, config: Dict[str, Any]):
         self.logger = get_logger("OrtInferSession")
 
         model_path = config.get("model_path", None)
@@ -48,7 +48,7 @@ class OrtInferSession:
         self._verify_providers()
 
     @staticmethod
-    def _init_sess_opts(config):
+    def _init_sess_opts(config: Dict[str, Any]) -> SessionOptions:
         sess_opt = SessionOptions()
         sess_opt.log_severity_level = 4
         sess_opt.enable_cpu_mem_arena = False
@@ -65,7 +65,7 @@ class OrtInferSession:
 
         return sess_opt
 
-    def _get_ep_list(self) -> List[Tuple[str, str]]:
+    def _get_ep_list(self) -> List[Tuple[str, Dict[str, Any]]]:
         cpu_provider_opts = {
             "arena_extend_strategy": "kSameAsRequested",
         }
@@ -172,7 +172,7 @@ class OrtInferSession:
         )
         return False
 
-    def _verify_providers(self) -> None:
+    def _verify_providers(self):
         session_providers = self.session.get_providers()
         first_provider = session_providers[0]
 
@@ -198,13 +198,13 @@ class OrtInferSession:
             error_info = traceback.format_exc()
             raise ONNXRuntimeError(error_info) from e
 
-    def get_input_names(self):
+    def get_input_names(self) -> List[str]:
         return [v.name for v in self.session.get_inputs()]
 
-    def get_output_names(self):
+    def get_output_names(self) -> List[str]:
         return [v.name for v in self.session.get_outputs()]
 
-    def get_character_list(self, key: str = "character"):
+    def get_character_list(self, key: str = "character") -> List[str]:
         meta_dict = self.session.get_modelmeta().custom_metadata_map
         return meta_dict[key].splitlines()
 

python/rapidocr_onnxruntime/utils/load_image.py

@@ -3,7 +3,7 @@
 # @Contact: liekkaskono@163.com
 from io import BytesIO
 from pathlib import Path
-from typing import Union
+from typing import Any, Union
 
 import cv2
 import numpy as np
@@ -55,7 +55,7 @@ class LoadImage:
             return np.array(img)
         return np.array(img)
 
-    def convert_img(self, img: np.ndarray, origin_img_type):
+    def convert_img(self, img: np.ndarray, origin_img_type: Any) -> np.ndarray:
         if img.ndim == 2:
             return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
 

python/rapidocr_onnxruntime/utils/logger.py

@@ -6,7 +6,7 @@ from functools import lru_cache
 
 
 @lru_cache(maxsize=32)
-def get_logger(name: str):
+def get_logger(name: str) -> logging.Logger:
     logger = logging.getLogger(name)
     logger.setLevel(logging.DEBUG)
 

python/rapidocr_onnxruntime/utils/parse_parameters.py

@@ -3,7 +3,7 @@
 # @Contact: liekkaskono@163.com
 import argparse
 from pathlib import Path
-from typing import Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 import numpy as np
 from PIL import Image
@@ -12,7 +12,7 @@ root_dir = Path(__file__).resolve().parent.parent
 InputType = Union[str, np.ndarray, bytes, Path, Image.Image]
 
 
-def update_model_path(config):
+def update_model_path(config: Dict[str, Any]) -> Dict[str, Any]:
     key = "model_path"
     config["Det"][key] = str(root_dir / config["Det"][key])
     config["Rec"][key] = str(root_dir / config["Rec"][key])

python/rapidocr_onnxruntime/utils/vis_res.py

@@ -29,7 +29,7 @@ class VisRes:
         scores: Optional[Tuple[float]] = None,
         font_path: Optional[str] = None,
     ) -> np.ndarray:
-        if txts is None and scores is None:
+        if txts 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)
 
@@ -52,7 +52,7 @@ class VisRes:
         self,
         img_content: InputType,
         dt_boxes: np.ndarray,
-        txts: Optional[Union[List[str], Tuple[str]]] = None,
+        txts: Union[List[str], Tuple[str]],
         scores: Optional[Tuple[float]] = None,
         font_path: Optional[str] = None,
     ) -> np.ndarray: