suport torch inference (#320)

* feat: add rapidocr_torch inference

* ci: fix push workflow

python/rapidocr_torch/__init__.py

@@ -0,0 +1,5 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+from .main import RapidOCR
+from .utils import LoadImageError, VisRes

python/rapidocr_torch/arch_config.yaml

@@ -0,0 +1,98 @@
+ch_ptocr_mobile_v2.0_cls_infer:
+  model_type: cls
+  algorithm: CLS
+  Transform:
+  Backbone:
+    name: MobileNetV3
+    scale: 0.35
+    model_name: small
+  Neck:
+  Head:
+    name: ClsHead
+    class_dim: 2
+
+ch_PP-OCRv4_det_infer:
+  model_type: det
+  algorithm: DB
+  Transform: null
+  Backbone:
+    name: PPLCNetV3
+    scale: 0.75
+    det: True
+  Neck:
+    name: RSEFPN
+    out_channels: 96
+    shortcut: True
+  Head:
+    name: DBHead
+    k: 50
+
+
+ch_PP-OCRv4_det_server_infer:
+  model_type: det
+  algorithm: DB
+  Transform: null
+  Backbone:
+    name: PPHGNet_small
+    det: True
+  Neck:
+    name: LKPAN
+    out_channels: 256
+    intracl: true
+  Head:
+    name: PFHeadLocal
+    k: 50
+    mode: "large"
+
+
+ch_PP-OCRv4_rec_infer:
+  model_type: rec
+  algorithm: SVTR_LCNet
+  Transform:
+  Backbone:
+    name: PPLCNetV3
+    scale: 0.95
+  Head:
+    name: MultiHead
+    out_channels_list:
+      CTCLabelDecode: 6625 #'blank' + ...(6623) + ' '
+    head_list:
+      - CTCHead:
+          Neck:
+            name: svtr
+            dims: 120
+            depth: 2
+            hidden_dims: 120
+            kernel_size: [ 1, 3 ]
+            use_guide: True
+          Head:
+            fc_decay: 0.00001
+      - NRTRHead:
+          nrtr_dim: 384
+          max_text_length: 25
+
+
+ch_PP-OCRv4_rec_server_infer:
+  model_type: rec
+  algorithm: SVTR_HGNet
+  Transform:
+  Backbone:
+    name: PPHGNet_small
+  Head:
+    name: MultiHead
+    out_channels_list:
+      CTCLabelDecode: 6625 #'blank' + ...(6623) + ' '
+    head_list:
+      - CTCHead:
+          Neck:
+            name: svtr
+            dims: 120
+            depth: 2
+            hidden_dims: 120
+            kernel_size: [ 1, 3 ]
+            use_guide: True
+          Head:
+            fc_decay: 0.00001
+      - NRTRHead:
+          nrtr_dim: 384
+          max_text_length: 25

python/rapidocr_torch/cal_rec_boxes/__init__.py

@@ -0,0 +1,4 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+from .main import CalRecBoxes

python/rapidocr_torch/cal_rec_boxes/main.py

@@ -0,0 +1,269 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL / Joker1212
+# @Contact: liekkaskono@163.com
+import copy
+import math
+from typing import Any, List, Optional, Tuple
+
+import cv2
+import numpy as np
+
+
+class CalRecBoxes:
+    """计算识别文字的汉字单字和英文单词的坐标框。代码借鉴自PaddlePaddle/PaddleOCR和fanqie03/char-detection"""
+
+    def __init__(self):
+        pass
+
+    def __call__(
+        self,
+        imgs: Optional[List[np.ndarray]],
+        dt_boxes: Optional[List[np.ndarray]],
+        rec_res: Optional[List[Any]],
+    ):
+        res = []
+        for img, box, rec_res in zip(imgs, dt_boxes, rec_res):
+            direction = self.get_box_direction(box)
+
+            rec_txt, rec_conf, rec_word_info = rec_res[0], rec_res[1], rec_res[2]
+            h, w = img.shape[:2]
+            img_box = np.array([[0, 0], [w, 0], [w, h], [0, h]])
+            word_box_content_list, word_box_list = self.cal_ocr_word_box(
+                rec_txt, img_box, rec_word_info
+            )
+            word_box_list = self.adjust_box_overlap(copy.deepcopy(word_box_list))
+            word_box_list = self.reverse_rotate_crop_image(
+                copy.deepcopy(box), word_box_list, direction
+            )
+            res.append([rec_txt, rec_conf, word_box_list, word_box_content_list])
+        return res
+
+    @staticmethod
+    def get_box_direction(box: np.ndarray) -> str:
+        direction = "w"
+        img_crop_width = int(
+            max(
+                np.linalg.norm(box[0] - box[1]),
+                np.linalg.norm(box[2] - box[3]),
+            )
+        )
+        img_crop_height = int(
+            max(
+                np.linalg.norm(box[0] - box[3]),
+                np.linalg.norm(box[1] - box[2]),
+            )
+        )
+        if img_crop_height * 1.0 / img_crop_width >= 1.5:
+            direction = "h"
+        return direction
+
+    @staticmethod
+    def cal_ocr_word_box(
+        rec_txt: str, box: np.ndarray, rec_word_info: List[Tuple[str, List[int]]]
+    ) -> Tuple[List[str], List[List[int]]]:
+        """Calculate the detection frame for each word based on the results of recognition and detection of ocr
+        汉字坐标是单字的
+        英语坐标是单词级别的
+        """
+
+        col_num, word_list, word_col_list, state_list = rec_word_info
+        box = box.tolist()
+        bbox_x_start = box[0][0]
+        bbox_x_end = box[1][0]
+        bbox_y_start = box[0][1]
+        bbox_y_end = box[2][1]
+
+        cell_width = (bbox_x_end - bbox_x_start) / col_num
+        word_box_list = []
+        word_box_content_list = []
+        cn_width_list = []
+        en_width_list = []
+        cn_col_list = []
+        en_col_list = []
+
+        def cal_char_width(width_list, word_col_):
+            if len(word_col_) == 1:
+                return
+            char_total_length = (word_col_[-1] - word_col_[0] + 1) * cell_width
+            char_width = char_total_length / (len(word_col_) - 1)
+            width_list.append(char_width)
+
+        def cal_box(col_list, width_list, word_box_list_):
+            if len(col_list) == 0:
+                return
+            if len(width_list) != 0:
+                avg_char_width = np.mean(width_list)
+            else:
+                avg_char_width = (bbox_x_end - bbox_x_start) / len(rec_txt)
+
+            for center_idx in col_list:
+                center_x = (center_idx + 0.5) * cell_width
+                cell_x_start = max(int(center_x - avg_char_width / 2), 0) + bbox_x_start
+                cell_x_end = (
+                    min(int(center_x + avg_char_width / 2), bbox_x_end - bbox_x_start)
+                    + bbox_x_start
+                )
+                cell = [
+                    [cell_x_start, bbox_y_start],
+                    [cell_x_end, bbox_y_start],
+                    [cell_x_end, bbox_y_end],
+                    [cell_x_start, bbox_y_end],
+                ]
+                word_box_list_.append(cell)
+
+        for word, word_col, state in zip(word_list, word_col_list, state_list):
+            if state == "cn":
+                cal_char_width(cn_width_list, word_col)
+                cn_col_list += word_col
+                word_box_content_list += word
+            else:
+                cal_char_width(en_width_list, word_col)
+                en_col_list += word_col
+                word_box_content_list += word
+
+        cal_box(cn_col_list, cn_width_list, word_box_list)
+        cal_box(en_col_list, en_width_list, word_box_list)
+        sorted_word_box_list = sorted(word_box_list, key=lambda box: box[0][0])
+        return word_box_content_list, sorted_word_box_list
+
+    @staticmethod
+    def adjust_box_overlap(
+        word_box_list: List[List[List[int]]],
+    ) -> List[List[List[int]]]:
+        # 调整bbox有重叠的地方
+        for i in range(len(word_box_list) - 1):
+            cur, nxt = word_box_list[i], word_box_list[i + 1]
+            if cur[1][0] > nxt[0][0]:  # 有交集
+                distance = abs(cur[1][0] - nxt[0][0])
+                cur[1][0] -= distance / 2
+                cur[2][0] -= distance / 2
+                nxt[0][0] += distance / 2
+                nxt[3][0] += distance / 2
+        return word_box_list
+
+    def reverse_rotate_crop_image(
+        self,
+        bbox_points: np.ndarray,
+        word_points_list: List[List[List[int]]],
+        direction: str = "w",
+    ) -> List[List[List[int]]]:
+        """
+        get_rotate_crop_image的逆操作
+        img为原图
+        part_img为crop后的图
+        bbox_points为part_img中对应在原图的bbox, 四个点，左上，右上，右下，左下
+        part_points为在part_img中的点[(x, y), (x, y)]
+        """
+        bbox_points = np.float32(bbox_points)
+
+        left = int(np.min(bbox_points[:, 0]))
+        top = int(np.min(bbox_points[:, 1]))
+        bbox_points[:, 0] = bbox_points[:, 0] - left
+        bbox_points[:, 1] = bbox_points[:, 1] - top
+
+        img_crop_width = int(np.linalg.norm(bbox_points[0] - bbox_points[1]))
+        img_crop_height = int(np.linalg.norm(bbox_points[0] - bbox_points[3]))
+
+        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(bbox_points, pts_std)
+        _, IM = cv2.invert(M)
+
+        new_word_points_list = []
+        for word_points in word_points_list:
+            new_word_points = []
+            for point in word_points:
+                new_point = point
+                if direction == "h":
+                    new_point = self.s_rotate(
+                        math.radians(-90), new_point[0], new_point[1], 0, 0
+                    )
+                    new_point[0] = new_point[0] + img_crop_width
+
+                p = np.float32(new_point + [1])
+                x, y, z = np.dot(IM, p)
+                new_point = [x / z, y / z]
+
+                new_point = [int(new_point[0] + left), int(new_point[1] + top)]
+                new_word_points.append(new_point)
+            new_word_points = self.order_points(new_word_points)
+            new_word_points_list.append(new_word_points)
+        return new_word_points_list
+
+    @staticmethod
+    def s_rotate(angle, valuex, valuey, pointx, pointy):
+        """绕pointx,pointy顺时针旋转
+        https://blog.csdn.net/qq_38826019/article/details/84233397
+        """
+        valuex = np.array(valuex)
+        valuey = np.array(valuey)
+        sRotatex = (
+            (valuex - pointx) * math.cos(angle)
+            + (valuey - pointy) * math.sin(angle)
+            + pointx
+        )
+        sRotatey = (
+            (valuey - pointy) * math.cos(angle)
+            - (valuex - pointx) * math.sin(angle)
+            + pointy
+        )
+        return [sRotatex, sRotatey]
+
+    @staticmethod
+    def order_points(box: List[List[int]]) -> List[List[int]]:
+        """矩形框顺序排列"""
+        box = np.array(box).reshape((-1, 2))
+        center_x, center_y = np.mean(box[:, 0]), np.mean(box[:, 1])
+        if np.any(box[:, 0] == center_x) and np.any(
+            box[:, 1] == center_y
+        ):  # 有两点横坐标相等，有两点纵坐标相等，菱形
+            p1 = box[np.where(box[:, 0] == np.min(box[:, 0]))]
+            p2 = box[np.where(box[:, 1] == np.min(box[:, 1]))]
+            p3 = box[np.where(box[:, 0] == np.max(box[:, 0]))]
+            p4 = box[np.where(box[:, 1] == np.max(box[:, 1]))]
+        elif np.all(box[:, 0] == center_x):  # 四个点的横坐标都相同
+            y_sort = np.argsort(box[:, 1])
+            p1 = box[y_sort[0]]
+            p2 = box[y_sort[1]]
+            p3 = box[y_sort[2]]
+            p4 = box[y_sort[3]]
+        elif np.any(box[:, 0] == center_x) and np.all(
+            box[:, 1] != center_y
+        ):  # 只有两点横坐标相等，先上下再左右
+            p12, p34 = (
+                box[np.where(box[:, 1] < center_y)],
+                box[np.where(box[:, 1] > center_y)],
+            )
+            p1, p2 = (
+                p12[np.where(p12[:, 0] == np.min(p12[:, 0]))],
+                p12[np.where(p12[:, 0] == np.max(p12[:, 0]))],
+            )
+            p3, p4 = (
+                p34[np.where(p34[:, 0] == np.max(p34[:, 0]))],
+                p34[np.where(p34[:, 0] == np.min(p34[:, 0]))],
+            )
+        else:  # 只有两点纵坐标相等，或者是没有相等的，先左右再上下
+            p14, p23 = (
+                box[np.where(box[:, 0] < center_x)],
+                box[np.where(box[:, 0] > center_x)],
+            )
+            p1, p4 = (
+                p14[np.where(p14[:, 1] == np.min(p14[:, 1]))],
+                p14[np.where(p14[:, 1] == np.max(p14[:, 1]))],
+            )
+            p2, p3 = (
+                p23[np.where(p23[:, 1] == np.min(p23[:, 1]))],
+                p23[np.where(p23[:, 1] == np.max(p23[:, 1]))],
+            )
+        # 解决单字矩形框重叠导致多个相同框的情况
+        p1 = p1[:1, :]
+        p2 = p2[:1, :]
+        p3 = p3[:1, :]
+        p4 = p4[:1, :]
+        return np.array([p1, p2, p3, p4]).reshape((-1, 2)).tolist()

python/rapidocr_torch/ch_ppocr_cls/__init__.py

@@ -0,0 +1,4 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+from .text_cls import TextClassifier

python/rapidocr_torch/ch_ppocr_cls/text_cls.py

@@ -0,0 +1,114 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+import argparse
+import copy
+import math
+import time
+from typing import Any, Dict, List, Tuple, Union
+
+import cv2
+import numpy as np
+from rapidocr_torch.utils import TorchInferSession, read_yaml
+from .utils import ClsPostProcess
+
+
+class TextClassifier:
+    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"]
+        self.postprocess_op = ClsPostProcess(config["label_list"])
+
+        self.infer = TorchInferSession(config)
+
+    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]
+
+        img_list = copy.deepcopy(img_list)
+
+        # Calculate the aspect ratio of all text bars
+        width_list = [img.shape[1] / float(img.shape[0]) for img in img_list]
+
+        # Sorting can speed up the cls process
+        indices = np.argsort(np.array(width_list))
+
+        img_num = len(img_list)
+        cls_res = [["", 0.0]] * img_num
+        batch_num = self.cls_batch_num
+        elapse = 0
+        for beg_img_no in range(0, img_num, batch_num):
+            end_img_no = min(img_num, beg_img_no + batch_num)
+
+            norm_img_batch = []
+            for ino in range(beg_img_no, end_img_no):
+                norm_img = self.resize_norm_img(img_list[indices[ino]])
+                norm_img = norm_img[np.newaxis, :]
+                norm_img_batch.append(norm_img)
+            norm_img_batch = np.concatenate(norm_img_batch).astype(np.float32)
+
+            starttime = time.time()
+            # 输出直接为(norm_img_batch,2)
+            prob_out = self.infer(norm_img_batch)[0]
+            cls_result = self.postprocess_op(prob_out)
+            elapse += time.time() - starttime
+
+            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(
+                        img_list[indices[beg_img_no + rno]], 1
+                    )
+        return img_list, cls_res, elapse
+
+    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)
+        if math.ceil(img_h * ratio) > img_w:
+            resized_w = img_w
+        else:
+            resized_w = int(math.ceil(img_h * ratio))
+
+        resized_image = cv2.resize(img, (resized_w, img_h))
+        resized_image = resized_image.astype("float32")
+        if img_c == 1:
+            resized_image = resized_image / 255
+            resized_image = resized_image[np.newaxis, :]
+        else:
+            resized_image = resized_image.transpose((2, 0, 1)) / 255
+
+        resized_image -= 0.5
+        resized_image /= 0.5
+        padding_im = np.zeros((img_c, img_h, img_w), dtype=np.float32)
+        padding_im[:, :, :resized_w] = resized_image
+        return padding_im
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--image_path", type=str, help="image_dir|image_path")
+    parser.add_argument("--config_path", type=str, default="config.yaml")
+    args = parser.parse_args()
+
+    config = read_yaml(args.config_path)
+
+    text_classifier = TextClassifier(config)
+
+    img = cv2.imread(args.image_path)
+    img_list, cls_res, predict_time = text_classifier(img)
+    for ino in range(len(img_list)):
+        print(f"cls result:{cls_res[ino]}")

python/rapidocr_torch/ch_ppocr_cls/utils.py

@@ -0,0 +1,28 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+from typing import List, Tuple
+
+import numpy as np
+
+
+class ClsPostProcess:
+    def __init__(self, label_list: List[str]):
+        self.label_list = label_list
+
+    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)
+        ]
+        return decode_out

python/rapidocr_torch/ch_ppocr_det/__init__.py

@@ -0,0 +1,4 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+from .text_detect import TextDetector

python/rapidocr_torch/ch_ppocr_det/text_detect.py

@@ -0,0 +1,108 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 time
+from typing import Any, Dict, Optional, Tuple
+
+import numpy as np
+from .utils import DBPostProcess, DetPreProcess
+from rapidocr_torch.utils.infer_engine import TorchInferSession
+
+
+class TextDetector:
+    def __init__(self, config: Dict[str, Any]):
+        limit_side_len = config.get("limit_side_len", 736)
+        limit_type = config.get("limit_type", "min")
+        self.preprocess_op = DetPreProcess(limit_side_len, limit_type)
+
+        post_process = {
+            "thresh": config.get("thresh", 0.3),
+            "box_thresh": config.get("box_thresh", 0.5),
+            "max_candidates": config.get("max_candidates", 1000),
+            "unclip_ratio": config.get("unclip_ratio", 1.6),
+            "use_dilation": config.get("use_dilation", True),
+            "score_mode": config.get("score_mode", "fast"),
+        }
+        self.postprocess_op = DBPostProcess(**post_process)
+        self.infer = TorchInferSession(config)
+
+
+    def __call__(self, img: np.ndarray) -> Tuple[Optional[np.ndarray], float]:
+        start_time = time.perf_counter()
+
+        if img is None:
+            raise ValueError("img is None")
+
+        ori_img_shape = img.shape[0], img.shape[1]
+        prepro_img = self.preprocess_op(img)
+        if prepro_img is None:
+            return None, 0
+
+        preds = self.infer(prepro_img)[0]
+        dt_boxes, dt_boxes_scores = self.postprocess_op(preds, ori_img_shape)
+        dt_boxes = self.filter_tag_det_res(dt_boxes, ori_img_shape)
+        elapse = time.perf_counter() - start_time
+        return dt_boxes, elapse
+
+    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
+        sort the points based on their x-coordinates
+        """
+        xSorted = pts[np.argsort(pts[:, 0]), :]
+
+        # grab the left-most and right-most points from the sorted
+        # x-roodinate points
+        leftMost = xSorted[:2, :]
+        rightMost = xSorted[2:, :]
+
+        # now, sort the left-most coordinates according to their
+        # y-coordinates so we can grab the top-left and bottom-left
+        # points, respectively
+        leftMost = leftMost[np.argsort(leftMost[:, 1]), :]
+        (tl, bl) = leftMost
+
+        rightMost = rightMost[np.argsort(rightMost[:, 1]), :]
+        (tr, br) = rightMost
+
+        rect = np.array([tl, tr, br, bl], dtype="float32")
+        return rect
+
+    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

python/rapidocr_torch/ch_ppocr_det/utils.py

@@ -0,0 +1,229 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+from typing import List, Optional, Tuple
+
+import cv2
+import numpy as np
+import pyclipper
+from shapely.geometry import Polygon
+
+
+class DetPreProcess:
+    def __init__(self, limit_side_len: int = 736, limit_type: str = "min"):
+        self.mean = np.array([0.5, 0.5, 0.5])
+        self.std = np.array([0.5, 0.5, 0.5])
+        self.scale = 1 / 255.0
+
+        self.limit_side_len = limit_side_len
+        self.limit_type = limit_type
+
+    def __call__(self, img: np.ndarray) -> Optional[np.ndarray]:
+        resized_img = self.resize(img)
+        if resized_img is None:
+            return None
+
+        img = self.normalize(resized_img)
+        img = self.permute(img)
+        img = np.expand_dims(img, axis=0).astype(np.float32)
+        return img
+
+    def normalize(self, img: np.ndarray) -> np.ndarray:
+        return (img.astype("float32") * self.scale - self.mean) / self.std
+
+    def permute(self, img: np.ndarray) -> np.ndarray:
+        return img.transpose((2, 0, 1))
+
+    def resize(self, img: np.ndarray) -> Optional[np.ndarray]:
+        """resize image to a size multiple of 32 which is required by the network"""
+        h, w = img.shape[:2]
+
+        if self.limit_type == "max":
+            if max(h, w) > self.limit_side_len:
+                if h > w:
+                    ratio = float(self.limit_side_len) / h
+                else:
+                    ratio = float(self.limit_side_len) / w
+            else:
+                ratio = 1.0
+        else:
+            if min(h, w) < self.limit_side_len:
+                if h < w:
+                    ratio = float(self.limit_side_len) / h
+                else:
+                    ratio = float(self.limit_side_len) / w
+            else:
+                ratio = 1.0
+
+        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:
+                return None
+            img = cv2.resize(img, (int(resize_w), int(resize_h)))
+        except Exception as exc:
+            raise ResizeImgError from exc
+
+        return img
+
+
+class ResizeImgError(Exception):
+    pass
+
+
+class DBPostProcess:
+    """The post process for Differentiable Binarization (DB)."""
+
+    def __init__(
+        self,
+        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
+        self.max_candidates = max_candidates
+        self.unclip_ratio = unclip_ratio
+        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]])
+
+    def __call__(
+        self, pred: np.ndarray, ori_shape: Tuple[int, int]
+    ) -> Tuple[np.ndarray, List[float]]:
+        src_h, src_w = ori_shape
+        pred = pred[:, 0, :, :]
+        segmentation = pred > self.thresh
+
+        mask = segmentation[0]
+        if self.dilation_kernel is not None:
+            mask = cv2.dilate(
+                np.array(segmentation[0]).astype(np.uint8), self.dilation_kernel
+            )
+        boxes, scores = self.boxes_from_bitmap(pred[0], mask, src_w, src_h)
+        return boxes, scores
+
+    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),
+                whose values are binarized as {0, 1}
+        """
+
+        height, width = bitmap.shape
+
+        outs = cv2.findContours(
+            (bitmap * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE
+        )
+        if len(outs) == 3:
+            img, contours, _ = outs[0], outs[1], outs[2]
+        elif len(outs) == 2:
+            contours, _ = outs[0], outs[1]
+
+        num_contours = min(len(contours), self.max_candidates)
+
+        boxes, scores = [], []
+        for index in range(num_contours):
+            contour = contours[index]
+            points, sside = self.get_mini_boxes(contour)
+            if sside < self.min_size:
+                continue
+
+            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)
+            box, sside = self.get_mini_boxes(box)
+            if sside < self.min_size + 2:
+                continue
+
+            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.int32))
+            scores.append(score)
+        return np.array(boxes, dtype=np.int32), scores
+
+    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])
+
+        index_1, index_2, index_3, index_4 = 0, 1, 2, 3
+        if points[1][1] > points[0][1]:
+            index_1 = 0
+            index_4 = 1
+        else:
+            index_1 = 1
+            index_4 = 0
+
+        if points[3][1] > points[2][1]:
+            index_2 = 2
+            index_3 = 3
+        else:
+            index_2 = 3
+            index_3 = 2
+
+        box = np.array(
+            [points[index_1], points[index_2], points[index_3], points[index_4]]
+        )
+        return box, min(bounding_box[1])
+
+    @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)
+        xmax = np.clip(np.ceil(box[:, 0].max()).astype(np.int32), 0, w - 1)
+        ymin = np.clip(np.floor(box[:, 1].min()).astype(np.int32), 0, h - 1)
+        ymax = np.clip(np.ceil(box[:, 1].max()).astype(np.int32), 0, h - 1)
+
+        mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
+        box[:, 0] = box[:, 0] - xmin
+        box[:, 1] = box[:, 1] - ymin
+        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: 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))
+
+        xmin = np.clip(np.min(contour[:, 0]), 0, w - 1)
+        xmax = np.clip(np.max(contour[:, 0]), 0, w - 1)
+        ymin = np.clip(np.min(contour[:, 1]), 0, h - 1)
+        ymax = np.clip(np.max(contour[:, 1]), 0, h - 1)
+
+        mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
+
+        contour[:, 0] = contour[:, 0] - xmin
+        contour[:, 1] = contour[:, 1] - ymin
+
+        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 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_torch/ch_ppocr_rec/__init__.py

@@ -0,0 +1,4 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+from .text_recognize import TextRecognizer

python/rapidocr_torch/ch_ppocr_rec/text_recognize.py

@@ -0,0 +1,126 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+import argparse
+import math
+import time
+from pathlib import Path
+from typing import Any, Dict, List, Tuple, Union
+
+import cv2
+import numpy as np
+
+from rapidocr_torch.utils import TorchInferSession, read_yaml
+from .utils import CTCLabelDecode
+
+DEFAULT_DICT_PATH = str(Path(__file__).parent / "ppocr_keys_v1.txt")
+class TextRecognizer:
+    def __init__(self, config: Dict[str, Any]):
+        dict_path = config.get("rec_keys_path", None)
+        self.character_dict_path = dict_path if dict_path else DEFAULT_DICT_PATH
+        self.postprocess_op = CTCLabelDecode(
+           character_path=self.character_dict_path
+        )
+
+        self.rec_batch_num = config["rec_batch_num"]
+        self.rec_image_shape = config["rec_img_shape"]
+        self.infer = TorchInferSession(config)
+
+    def __call__(
+        self,
+        img_list: Union[np.ndarray, List[np.ndarray]],
+        return_word_box: bool = False,
+    ) -> Tuple[List[Tuple[str, float]], float]:
+        if isinstance(img_list, np.ndarray):
+            img_list = [img_list]
+
+        # Calculate the aspect ratio of all text bars
+        width_list = [img.shape[1] / float(img.shape[0]) for img in img_list]
+
+        # Sorting can speed up the recognition process
+        indices = np.argsort(np.array(width_list))
+
+        img_num = len(img_list)
+        rec_res = [("", 0.0)] * img_num
+
+        batch_num = self.rec_batch_num
+        elapse = 0
+        for beg_img_no in range(0, img_num, batch_num):
+            end_img_no = min(img_num, beg_img_no + batch_num)
+
+            # Parameter Alignment for PaddleOCR
+            imgC, imgH, imgW = self.rec_image_shape[:3]
+            max_wh_ratio = imgW / imgH
+            wh_ratio_list = []
+            for ino in range(beg_img_no, end_img_no):
+                h, w = img_list[indices[ino]].shape[0:2]
+                wh_ratio = w * 1.0 / h
+                max_wh_ratio = max(max_wh_ratio, wh_ratio)
+                wh_ratio_list.append(wh_ratio)
+
+            norm_img_batch = []
+            for ino in range(beg_img_no, end_img_no):
+                norm_img = self.resize_norm_img(img_list[indices[ino]], max_wh_ratio)
+                norm_img_batch.append(norm_img[np.newaxis, :])
+            norm_img_batch = np.concatenate(norm_img_batch).astype(np.float32)
+
+            starttime = time.time()
+            preds = self.infer(norm_img_batch)[0]
+            rec_result = self.postprocess_op(
+                preds,
+                return_word_box,
+                wh_ratio_list=wh_ratio_list,
+                max_wh_ratio=max_wh_ratio,
+            )
+
+            for rno, one_res in enumerate(rec_result):
+                rec_res[indices[beg_img_no + rno]] = one_res
+            elapse += time.time() - starttime
+        return rec_res, elapse
+
+    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]
+
+        img_width = int(img_height * max_wh_ratio)
+
+        h, w = img.shape[:2]
+        ratio = w / float(h)
+        if math.ceil(img_height * ratio) > img_width:
+            resized_w = img_width
+        else:
+            resized_w = int(math.ceil(img_height * ratio))
+
+        resized_image = cv2.resize(img, (resized_w, img_height))
+        resized_image = resized_image.astype("float32")
+        resized_image = resized_image.transpose((2, 0, 1)) / 255
+        resized_image -= 0.5
+        resized_image /= 0.5
+
+        padding_im = np.zeros((img_channel, img_height, img_width), dtype=np.float32)
+        padding_im[:, :, 0:resized_w] = resized_image
+        return padding_im
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--image_path", type=str, help="image_dir|image_path")
+    parser.add_argument("--config_path", type=str, default="config.yaml")
+    args = parser.parse_args()
+
+    config = read_yaml(args.config_path)
+    text_recognizer = TextRecognizer(config)
+
+    img = cv2.imread(args.image_path)
+    rec_res, predict_time = text_recognizer(img)
+    print(f"rec result: {rec_res}\t cost: {predict_time}s")

python/rapidocr_torch/ch_ppocr_rec/utils.py

@@ -0,0 +1,182 @@
+# -*- 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:
+    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 __call__(
+        self, preds: np.ndarray, return_word_box: bool = False, **kwargs
+    ) -> List[Tuple[str, float]]:
+        preds_idx = preds.argmax(axis=2)
+        preds_prob = preds.max(axis=2)
+        text = self.decode(
+            preds_idx, preds_prob, return_word_box, is_remove_duplicate=True
+        )
+        if return_word_box:
+            for rec_idx, rec in enumerate(text):
+                wh_ratio = kwargs["wh_ratio_list"][rec_idx]
+                max_wh_ratio = kwargs["max_wh_ratio"]
+                rec[2][0] = rec[2][0] * (wh_ratio / max_wh_ratio)
+        return text
+
+    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")
+
+        character_list = None
+        if character:
+            character_list = character
+
+        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,
+        return_word_box: bool = False,
+        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):
+            selection = np.ones(len(text_index[batch_idx]), dtype=bool)
+            if is_remove_duplicate:
+                selection[1:] = text_index[batch_idx][1:] != text_index[batch_idx][:-1]
+
+            for ignored_token in ignored_tokens:
+                selection &= text_index[batch_idx] != ignored_token
+
+            if text_prob is not None:
+                conf_list = text_prob[batch_idx][selection]
+            else:
+                conf_list = [1] * len(selection)
+
+            if len(conf_list) == 0:
+                conf_list = [0]
+
+            char_list = [
+                self.character[text_id] for text_id in text_index[batch_idx][selection]
+            ]
+            text = "".join(char_list)
+            if return_word_box:
+                word_list, word_col_list, state_list = self.get_word_info(
+                    text, selection
+                )
+                result_list.append(
+                    (
+                        text,
+                        np.mean(conf_list).tolist(),
+                        [
+                            len(text_index[batch_idx]),
+                            word_list,
+                            word_col_list,
+                            state_list,
+                        ],
+                    )
+                )
+            else:
+                result_list.append((text, np.mean(conf_list).tolist()))
+        return result_list
+
+    @staticmethod
+    def get_word_info(
+        text: str, selection: np.ndarray
+    ) -> Tuple[List[List[str]], List[List[int]], List[str]]:
+        """
+        Group the decoded characters and record the corresponding decoded positions.
+        from https://github.com/PaddlePaddle/PaddleOCR/blob/fbba2178d7093f1dffca65a5b963ec277f1a6125/ppocr/postprocess/rec_postprocess.py#L70
+
+        Args:
+            text: the decoded text
+            selection: the bool array that identifies which columns of features are decoded as non-separated characters
+        Returns:
+            word_list: list of the grouped words
+            word_col_list: list of decoding positions corresponding to each character in the grouped word
+            state_list: list of marker to identify the type of grouping words, including two types of grouping words:
+                        - 'cn': continous chinese characters (e.g., 你好啊)
+                        - 'en&num': continous english characters (e.g., hello), number (e.g., 123, 1.123), or mixed of them connected by '-' (e.g., VGG-16)
+        """
+        state = None
+        word_content = []
+        word_col_content = []
+        word_list = []
+        word_col_list = []
+        state_list = []
+        valid_col = np.where(selection == True)[0]
+
+        for c_i, char in enumerate(text):
+            if "\u4e00" <= char <= "\u9fff":
+                c_state = "cn"
+            else:
+                c_state = "en&num"
+
+            if state == None:
+                state = c_state
+
+            if state != c_state:
+                if len(word_content) != 0:
+                    word_list.append(word_content)
+                    word_col_list.append(word_col_content)
+                    state_list.append(state)
+                    word_content = []
+                    word_col_content = []
+                state = c_state
+
+            word_content.append(char)
+            word_col_content.append(int(valid_col[c_i]))
+
+        if len(word_content) != 0:
+            word_list.append(word_content)
+            word_col_list.append(word_col_content)
+            state_list.append(state)
+
+        return word_list, word_col_list, state_list
+
+    @staticmethod
+    def get_ignored_tokens() -> List[int]:
+        return [0]  # for ctc blank

python/rapidocr_torch/config.yaml

@@ -0,0 +1,59 @@
+Global:
+    text_score: 0.5
+    use_det: true
+    use_cls: true
+    use_rec: true
+    print_verbose: false
+    min_height: 30
+    width_height_ratio: 8
+    max_side_len: 2000
+    min_side_len: 30
+    return_word_box: false
+
+    intra_op_num_threads: &intra_nums -1
+    inter_op_num_threads: &inter_nums -1
+
+Det:
+    intra_op_num_threads: *intra_nums
+    inter_op_num_threads: *inter_nums
+
+    use_cuda: false
+    use_dml: false
+
+    model_path: models/ch_PP-OCRv4_det_infer.pth
+
+    limit_side_len: 736
+    limit_type: min
+
+    thresh: 0.3
+    box_thresh: 0.5
+    max_candidates: 1000
+    unclip_ratio: 1.5
+    use_dilation: true
+    score_mode: fast
+
+Cls:
+    intra_op_num_threads: *intra_nums
+    inter_op_num_threads: *inter_nums
+
+    use_cuda: false
+    use_dml: false
+
+    model_path: models/ch_ptocr_mobile_v2.0_cls_infer.pth
+
+    cls_image_shape: [3, 48, 192]
+    cls_batch_num: 6
+    cls_thresh: 0.9
+    label_list: ['0', '180']
+
+Rec:
+    intra_op_num_threads: *intra_nums
+    inter_op_num_threads: *inter_nums
+
+    use_cuda: false
+    use_dml: false
+
+    model_path: models/ch_PP-OCRv4_rec_infer.pth
+
+    rec_img_shape: [3, 48, 320]
+    rec_batch_num: 6

python/rapidocr_torch/main.py

@@ -0,0 +1,368 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+import copy
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import cv2
+import numpy as np
+
+from .cal_rec_boxes import CalRecBoxes
+from .ch_ppocr_cls import TextClassifier
+from .ch_ppocr_det import TextDetector
+from .ch_ppocr_rec import TextRecognizer
+from .utils import (
+    LoadImage,
+    UpdateParameters,
+    VisRes,
+    add_round_letterbox,
+    get_logger,
+    increase_min_side,
+    init_args,
+    read_yaml,
+    reduce_max_side,
+    update_model_path,
+)
+
+root_dir = Path(__file__).resolve().parent
+DEFAULT_CFG_PATH = root_dir / "config.yaml"
+logger = get_logger("RapidOCR")
+
+
+class RapidOCR:
+    def __init__(self, config_path: Optional[str] = None, **kwargs):
+        if config_path is not None and Path(config_path).exists():
+            config = read_yaml(config_path)
+        else:
+            config = read_yaml(DEFAULT_CFG_PATH)
+        config = update_model_path(config)
+
+        if kwargs:
+            updater = UpdateParameters()
+            config = updater(config, **kwargs)
+
+        global_config = config["Global"]
+        self.print_verbose = global_config["print_verbose"]
+        self.text_score = global_config["text_score"]
+        self.min_height = global_config["min_height"]
+        self.width_height_ratio = global_config["width_height_ratio"]
+
+        self.use_det = global_config["use_det"]
+        self.text_det = TextDetector(config["Det"])
+
+        self.use_cls = global_config["use_cls"]
+        self.text_cls = TextClassifier(config["Cls"])
+
+        self.use_rec = global_config["use_rec"]
+        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"]
+
+        self.cal_rec_boxes = CalRecBoxes()
+
+    def __call__(
+        self,
+        img_content: Union[str, np.ndarray, bytes, Path],
+        use_det: Optional[bool] = None,
+        use_cls: Optional[bool] = None,
+        use_rec: Optional[bool] = None,
+        **kwargs,
+    ) -> Tuple[Optional[List[List[Union[Any, str]]]], Optional[List[float]]]:
+        use_det = self.use_det if use_det is None else use_det
+        use_cls = self.use_cls if use_cls is None else use_cls
+        use_rec = self.use_rec if use_rec is None else use_rec
+        return_word_box = False
+        if kwargs:
+            box_thresh = kwargs.get("box_thresh", 0.5)
+            unclip_ratio = kwargs.get("unclip_ratio", 1.6)
+            text_score = kwargs.get("text_score", 0.5)
+            return_word_box = kwargs.get("return_word_box", False)
+            self.text_det.postprocess_op.box_thresh = box_thresh
+            self.text_det.postprocess_op.unclip_ratio = unclip_ratio
+            self.text_score = text_score
+
+        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, 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
+
+            img = self.get_crop_img_list(img, dt_boxes)
+
+        if use_cls:
+            img, cls_res, cls_elapse = self.text_cls(img)
+
+        if use_rec:
+            rec_res, rec_elapse = self.text_rec(img, return_word_box)
+
+        if dt_boxes is not None and rec_res is not None and return_word_box:
+            rec_res = self.cal_rec_boxes(img, dt_boxes, rec_res)
+            for rec_res_i in rec_res:
+                if rec_res_i[2]:
+                    rec_res_i[2] = (
+                        self._get_origin_points(rec_res_i[2], op_record, raw_h, raw_w)
+                        .astype(np.int32)
+                        .tolist()
+                    )
+
+        if dt_boxes 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 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:
+            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 = 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
+    ) -> 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
+
+        dt_boxes = self.sorted_boxes(dt_boxes)
+        return dt_boxes, det_elapse
+
+    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]),
+                    np.linalg.norm(points[2] - points[3]),
+                )
+            )
+            img_crop_height = int(
+                max(
+                    np.linalg.norm(points[0] - points[3]),
+                    np.linalg.norm(points[1] - points[2]),
+                )
+            )
+            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,
+                M,
+                (img_crop_width, img_crop_height),
+                borderMode=cv2.BORDER_REPLICATE,
+                flags=cv2.INTER_CUBIC,
+            )
+            dst_img_height, dst_img_width = dst_img.shape[0:2]
+            if dst_img_height * 1.0 / dst_img_width >= 1.5:
+                dst_img = np.rot90(dst_img)
+            return dst_img
+
+        img_crop_list = []
+        for box in dt_boxes:
+            tmp_box = copy.deepcopy(box)
+            img_crop = get_rotate_crop_image(img, tmp_box)
+            img_crop_list.append(img_crop)
+        return img_crop_list
+
+    @staticmethod
+    def sorted_boxes(dt_boxes: np.ndarray) -> List[np.ndarray]:
+        """
+        Sort text boxes in order from top to bottom, left to right
+        args:
+            dt_boxes(array):detected text boxes with shape [4, 2]
+        return:
+            sorted boxes(array) with shape [4, 2]
+        """
+        num_boxes = dt_boxes.shape[0]
+        sorted_boxes = sorted(dt_boxes, key=lambda x: (x[0][1], x[0][0]))
+        _boxes = list(sorted_boxes)
+
+        for i in range(num_boxes - 1):
+            for j in range(i, -1, -1):
+                if (
+                    abs(_boxes[j + 1][0][1] - _boxes[j][0][1]) < 10
+                    and _boxes[j + 1][0][0] < _boxes[j][0][0]
+                ):
+                    tmp = _boxes[j]
+                    _boxes[j] = _boxes[j + 1]
+                    _boxes[j + 1] = tmp
+                else:
+                    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).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
+
+    def get_final_res(
+        self,
+        dt_boxes: Optional[List[np.ndarray]],
+        cls_res: Optional[List[List[Union[str, float]]]],
+        rec_res: Optional[List[Tuple[str, float, List[Union[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]
+
+        if dt_boxes is None and rec_res is None:
+            return None, None
+
+        if dt_boxes is None and rec_res is not None:
+            return [[res[0], res[1]] for res in rec_res], [rec_elapse]
+
+        if dt_boxes is not None and rec_res is None:
+            return [box.tolist() for box in dt_boxes], [det_elapse]
+
+        dt_boxes, rec_res = self.filter_result(dt_boxes, rec_res)
+        if not dt_boxes or not rec_res or len(dt_boxes) <= 0:
+            return None, None
+
+        ocr_res = [[box.tolist(), *res] for box, res in zip(dt_boxes, rec_res)], [
+            det_elapse,
+            cls_elapse,
+            rec_elapse,
+        ]
+        return ocr_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[0], rec_reuslt[1]
+            if float(score) >= self.text_score:
+                filter_boxes.append(box)
+                filter_rec_res.append(rec_reuslt)
+
+        return filter_boxes, filter_rec_res
+
+
+def main():
+    args = init_args()
+    ocr_engine = RapidOCR(**vars(args))
+
+    use_det = not args.no_det
+    use_cls = not args.no_cls
+    use_rec = not args.no_rec
+    args.img_path = "D:\pythonProjects\RapidOCR\python\\tests\\test_files\\text_det.jpg"
+    args.vis_save_path = "../output/"
+    result, elapse_list = ocr_engine(
+        args.img_path,
+        use_det=use_det,
+        use_cls=use_cls,
+        use_rec=use_rec,
+    )
+    logger.info(result)
+
+    if args.print_cost:
+        logger.info(elapse_list)
+
+    if args.vis_res:
+        vis = VisRes()
+        Path(args.vis_save_path).mkdir(parents=True, exist_ok=True)
+        save_path = Path(args.vis_save_path) / f"{Path(args.img_path).stem}_vis.png"
+
+        if use_det and not use_cls and not use_rec:
+            boxes, *_ = list(zip(*result))
+            vis_img = vis(args.img_path, boxes)
+            cv2.imwrite(str(save_path), vis_img)
+            logger.info("The vis result has saved in %s", save_path)
+
+        elif use_det and use_rec:
+            font_path = Path(args.vis_font_path)
+            if not font_path.exists():
+                raise FileExistsError(f"{font_path} does not exist!")
+
+            boxes, txts, scores = list(zip(*result))
+            vis_img = vis(args.img_path, boxes, txts, scores, font_path=font_path)
+            cv2.imwrite(str(save_path), vis_img)
+            logger.info("The vis result has saved in %s", save_path)
+
+
+if __name__ == "__main__":
+    main()

python/rapidocr_torch/modeling/architectures/__init__.py

@@ -0,0 +1,25 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import copy
+
+__all__ = ['build_model']
+
+
+def build_model(config, **kwargs):
+    from .base_model import BaseModel
+
+    config = copy.deepcopy(config)
+    module_class = BaseModel(config, **kwargs)
+    return module_class

python/rapidocr_torch/modeling/architectures/base_model.py

@@ -0,0 +1,105 @@
+import torch.nn as nn
+from rapidocr_torch.modeling.backbones import build_backbone
+from rapidocr_torch.modeling.heads import build_head
+from rapidocr_torch.modeling.necks import build_neck
+
+
+class BaseModel(nn.Module):
+    def __init__(self, config, **kwargs):
+        """
+        the module for OCR.
+        args:
+            config (dict): the super parameters for module.
+        """
+        super(BaseModel, self).__init__()
+
+        in_channels = config.get('in_channels', 3)
+        model_type = config['model_type']
+        # build backbone, backbone is need for del, rec and cls
+        if 'Backbone' not in config or config['Backbone'] is None:
+            self.use_backbone = False
+        else:
+            self.use_backbone = True
+            config["Backbone"]['in_channels'] = in_channels
+            self.backbone = build_backbone(config["Backbone"], model_type)
+            in_channels = self.backbone.out_channels
+
+        # build neck
+        # for rec, neck can be cnn,rnn or reshape(None)
+        # for det, neck can be FPN, BIFPN and so on.
+        # for cls, neck should be none
+        if 'Neck' not in config or config['Neck'] is None:
+            self.use_neck = False
+        else:
+            self.use_neck = True
+            config['Neck']['in_channels'] = in_channels
+            self.neck = build_neck(config['Neck'])
+            in_channels = self.neck.out_channels
+
+        # # build head, head is need for det, rec and cls
+        if 'Head' not in config or config['Head'] is None:
+            self.use_head = False
+        else:
+            self.use_head = True
+            config["Head"]['in_channels'] = in_channels
+            self.head = build_head(config["Head"], **kwargs)
+
+        self.return_all_feats = config.get("return_all_feats", False)
+
+        self._initialize_weights()
+
+    def _initialize_weights(self):
+        # weight initialization
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out')
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.ones_(m.weight)
+                nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.ConvTranspose2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out')
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+
+    def forward(self, x):
+        y = dict()
+        if self.use_backbone:
+            x = self.backbone(x)
+        if isinstance(x, dict):
+            y.update(x)
+        else:
+            y["backbone_out"] = x
+        final_name = "backbone_out"
+        if self.use_neck:
+            x = self.neck(x)
+            if isinstance(x, dict):
+                y.update(x)
+            else:
+                y["neck_out"] = x
+            final_name = "neck_out"
+        if self.use_head:
+            x = self.head(x)
+        # for multi head, save ctc neck out for udml
+        if isinstance(x, dict) and 'ctc_nect' in x.keys():
+            y['neck_out'] = x['ctc_neck']
+            y['head_out'] = x
+        elif isinstance(x, dict):
+            y.update(x)
+        else:
+            y["head_out"] = x
+        if self.return_all_feats:
+            if self.training:
+                return y
+            elif isinstance(x, dict):
+                return x
+            else:
+                return {final_name: x}
+        else:
+            return x

python/rapidocr_torch/modeling/backbones/__init__.py

@@ -0,0 +1,39 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+__all__ = ['build_backbone']
+
+
+def build_backbone(config, model_type):
+    if model_type == 'det':
+        from .det_mobilenet_v3 import MobileNetV3
+        from .rec_lcnetv3 import PPLCNetV3
+        from .rec_hgnet import PPHGNet_small
+        support_dict = ['MobileNetV3', 'ResNet', 'ResNet_vd', 'ResNet_SAST', 'PPLCNetV3', 'PPHGNet_small']
+    elif model_type == 'rec' or model_type == 'cls':
+        from .rec_mobilenet_v3 import MobileNetV3
+        from .rec_svtrnet import SVTRNet
+        from .rec_lcnetv3 import PPLCNetV3
+        from .rec_hgnet import PPHGNet_small
+        support_dict = ['MobileNetV1Enhance', 'MobileNetV3', 'ResNet', 'ResNetFPN', 'MTB',
+                        'ResNet31', 'SVTRNet', 'ViTSTR', 'DenseNet', 'PPLCNetV3', 'PPHGNet_small']
+    else:
+        raise NotImplementedError
+
+    module_name = config.pop('name')
+    assert module_name in support_dict, Exception(
+        'when model typs is {}, backbone only support {}'.format(model_type,
+                                                                 support_dict))
+    module_class = eval(module_name)(**config)
+    return module_class

python/rapidocr_torch/modeling/backbones/det_mobilenet_v3.py

@@ -0,0 +1,253 @@
+import torch.nn as nn
+from rapidocr_torch.modeling.common import Activation
+
+def make_divisible(v, divisor=8, min_value=None):
+    if min_value is None:
+        min_value = divisor
+    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+    if new_v < 0.9 * v:
+        new_v += divisor
+    return new_v
+
+
+class ConvBNLayer(nn.Module):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride,
+                 padding,
+                 groups=1,
+                 if_act=True,
+                 act=None,
+                 name=None):
+        super(ConvBNLayer, self).__init__()
+        self.if_act = if_act
+        self.conv = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            groups=groups,
+            bias=False)
+
+        self.bn = nn.BatchNorm2d(
+            out_channels,
+            )
+        if self.if_act:
+            self.act = Activation(act_type=act, inplace=True)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        if self.if_act:
+            x = self.act(x)
+        return x
+
+
+class SEModule(nn.Module):
+    def __init__(self, in_channels, reduction=4, name=""):
+        super(SEModule, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.conv1 = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=in_channels // reduction,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=True)
+        self.relu1 = Activation(act_type='relu', inplace=True)
+        self.conv2 = nn.Conv2d(
+            in_channels=in_channels // reduction,
+            out_channels=in_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=True)
+        self.hard_sigmoid = Activation(act_type='hard_sigmoid', inplace=True)
+
+    def forward(self, inputs):
+        outputs = self.avg_pool(inputs)
+        outputs = self.conv1(outputs)
+        outputs = self.relu1(outputs)
+        outputs = self.conv2(outputs)
+        outputs = self.hard_sigmoid(outputs)
+        outputs = inputs * outputs
+        return outputs
+
+
+class ResidualUnit(nn.Module):
+    def __init__(self,
+                 in_channels,
+                 mid_channels,
+                 out_channels,
+                 kernel_size,
+                 stride,
+                 use_se,
+                 act=None,
+                 name=''):
+        super(ResidualUnit, self).__init__()
+        self.if_shortcut = stride == 1 and in_channels == out_channels
+        self.if_se = use_se
+
+        self.expand_conv = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=mid_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            if_act=True,
+            act=act,
+            name=name + "_expand")
+        self.bottleneck_conv = ConvBNLayer(
+            in_channels=mid_channels,
+            out_channels=mid_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=int((kernel_size - 1) // 2),
+            groups=mid_channels,
+            if_act=True,
+            act=act,
+            name=name + "_depthwise")
+        if self.if_se:
+            self.mid_se = SEModule(mid_channels, name=name + "_se")
+        self.linear_conv = ConvBNLayer(
+            in_channels=mid_channels,
+            out_channels=out_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            if_act=False,
+            act=None,
+            name=name + "_linear")
+
+    def forward(self, inputs):
+        x = self.expand_conv(inputs)
+        x = self.bottleneck_conv(x)
+        if self.if_se:
+            x = self.mid_se(x)
+        x = self.linear_conv(x)
+        if self.if_shortcut:
+            x = inputs + x
+        return x
+
+
+class MobileNetV3(nn.Module):
+    def __init__(self,
+                 in_channels=3,
+                 model_name='large',
+                 scale=0.5,
+                 disable_se=False,
+                 **kwargs):
+        """
+        the MobilenetV3 backbone network for detection module.
+        Args:
+            params(dict): the super parameters for build network
+        """
+        super(MobileNetV3, self).__init__()
+
+        self.disable_se = disable_se
+
+        if model_name == "large":
+            cfg = [
+                # k, exp, c,  se,     nl,  s,
+                [3, 16, 16, False, 'relu', 1],
+                [3, 64, 24, False, 'relu', 2],
+                [3, 72, 24, False, 'relu', 1],
+                [5, 72, 40, True, 'relu', 2],
+                [5, 120, 40, True, 'relu', 1],
+                [5, 120, 40, True, 'relu', 1],
+                [3, 240, 80, False, 'hard_swish', 2],
+                [3, 200, 80, False, 'hard_swish', 1],
+                [3, 184, 80, False, 'hard_swish', 1],
+                [3, 184, 80, False, 'hard_swish', 1],
+                [3, 480, 112, True, 'hard_swish', 1],
+                [3, 672, 112, True, 'hard_swish', 1],
+                [5, 672, 160, True, 'hard_swish', 2],
+                [5, 960, 160, True, 'hard_swish', 1],
+                [5, 960, 160, True, 'hard_swish', 1],
+            ]
+            cls_ch_squeeze = 960
+        elif model_name == "small":
+            cfg = [
+                # k, exp, c,  se,     nl,  s,
+                [3, 16, 16, True, 'relu', 2],
+                [3, 72, 24, False, 'relu', 2],
+                [3, 88, 24, False, 'relu', 1],
+                [5, 96, 40, True, 'hard_swish', 2],
+                [5, 240, 40, True, 'hard_swish', 1],
+                [5, 240, 40, True, 'hard_swish', 1],
+                [5, 120, 48, True, 'hard_swish', 1],
+                [5, 144, 48, True, 'hard_swish', 1],
+                [5, 288, 96, True, 'hard_swish', 2],
+                [5, 576, 96, True, 'hard_swish', 1],
+                [5, 576, 96, True, 'hard_swish', 1],
+            ]
+            cls_ch_squeeze = 576
+        else:
+            raise NotImplementedError("mode[" + model_name +
+                                      "_model] is not implemented!")
+
+        supported_scale = [0.35, 0.5, 0.75, 1.0, 1.25]
+        assert scale in supported_scale, \
+            "supported scale are {} but input scale is {}".format(supported_scale, scale)
+        inplanes = 16
+        # conv1
+        self.conv = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=make_divisible(inplanes * scale),
+            kernel_size=3,
+            stride=2,
+            padding=1,
+            groups=1,
+            if_act=True,
+            act='hard_swish',
+            name='conv1')
+
+        self.stages = nn.ModuleList()
+        self.out_channels = []
+        block_list = []
+        i = 0
+        inplanes = make_divisible(inplanes * scale)
+        for (k, exp, c, se, nl, s) in cfg:
+            se = se and not self.disable_se
+            if s == 2 and i > 2:
+                self.out_channels.append(inplanes)
+                self.stages.append(nn.Sequential(*block_list))
+                block_list = []
+            block_list.append(
+                ResidualUnit(
+                    in_channels=inplanes,
+                    mid_channels=make_divisible(scale * exp),
+                    out_channels=make_divisible(scale * c),
+                    kernel_size=k,
+                    stride=s,
+                    use_se=se,
+                    act=nl,
+                    name="conv" + str(i + 2)))
+            inplanes = make_divisible(scale * c)
+            i += 1
+        block_list.append(
+            ConvBNLayer(
+                in_channels=inplanes,
+                out_channels=make_divisible(scale * cls_ch_squeeze),
+                kernel_size=1,
+                stride=1,
+                padding=0,
+                groups=1,
+                if_act=True,
+                act='hard_swish',
+                name='conv_last'))
+        self.stages.append(nn.Sequential(*block_list))
+        self.out_channels.append(make_divisible(scale * cls_ch_squeeze))
+        # for i, stage in enumerate(self.stages):
+        #     self.add_sublayer(sublayer=stage, name="stage{}".format(i))
+
+    def forward(self, x):
+        x = self.conv(x)
+        out_list = []
+        for stage in self.stages:
+            x = stage(x)
+            out_list.append(x)
+        return out_list

python/rapidocr_torch/modeling/backbones/rec_hgnet.py

@@ -0,0 +1,272 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ConvBNAct(nn.Module):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride,
+                 groups=1,
+                 use_act=True):
+        super().__init__()
+        self.use_act = use_act
+        self.conv = nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            padding=(kernel_size - 1) // 2,
+            groups=groups,
+            bias=False)
+        self.bn = nn.BatchNorm2d(out_channels)
+        if self.use_act:
+            self.act = nn.ReLU()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        if self.use_act:
+            x = self.act(x)
+        return x
+
+
+class ESEModule(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.conv = nn.Conv2d(
+            in_channels=channels,
+            out_channels=channels,
+            kernel_size=1,
+            stride=1,
+            padding=0)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        identity = x
+        x = self.avg_pool(x)
+        x = self.conv(x)
+        x = self.sigmoid(x)
+        return x * identity
+
+
+class HG_Block(nn.Module):
+    def __init__(
+            self,
+            in_channels,
+            mid_channels,
+            out_channels,
+            layer_num,
+            identity=False, ):
+        super().__init__()
+        self.identity = identity
+
+        self.layers = nn.ModuleList()
+        self.layers.append(
+            ConvBNAct(
+                in_channels=in_channels,
+                out_channels=mid_channels,
+                kernel_size=3,
+                stride=1))
+        for _ in range(layer_num - 1):
+            self.layers.append(
+                ConvBNAct(
+                    in_channels=mid_channels,
+                    out_channels=mid_channels,
+                    kernel_size=3,
+                    stride=1))
+
+        # feature aggregation
+        total_channels = in_channels + layer_num * mid_channels
+        self.aggregation_conv = ConvBNAct(
+            in_channels=total_channels,
+            out_channels=out_channels,
+            kernel_size=1,
+            stride=1)
+        self.att = ESEModule(out_channels)
+
+    def forward(self, x):
+        identity = x
+        output = []
+        output.append(x)
+        for layer in self.layers:
+            x = layer(x)
+            output.append(x)
+        x = torch.cat(output, dim=1)
+        x = self.aggregation_conv(x)
+        x = self.att(x)
+        if self.identity:
+            x += identity
+        return x
+
+
+class HG_Stage(nn.Module):
+    def __init__(self,
+                 in_channels,
+                 mid_channels,
+                 out_channels,
+                 block_num,
+                 layer_num,
+                 downsample=True,
+                 stride=[2, 1]):
+        super().__init__()
+        self.downsample = downsample
+        if downsample:
+            self.downsample = ConvBNAct(
+                in_channels=in_channels,
+                out_channels=in_channels,
+                kernel_size=3,
+                stride=stride,
+                groups=in_channels,
+                use_act=False)
+
+        blocks_list = []
+        blocks_list.append(
+            HG_Block(
+                in_channels,
+                mid_channels,
+                out_channels,
+                layer_num,
+                identity=False))
+        for _ in range(block_num - 1):
+            blocks_list.append(
+                HG_Block(
+                    out_channels,
+                    mid_channels,
+                    out_channels,
+                    layer_num,
+                    identity=True))
+        self.blocks = nn.Sequential(*blocks_list)
+
+    def forward(self, x):
+        if self.downsample:
+            x = self.downsample(x)
+        x = self.blocks(x)
+        return x
+
+
+class PPHGNet(nn.Module):
+    """
+    PPHGNet
+    Args:
+        stem_channels: list. Stem channel list of PPHGNet.
+        stage_config: dict. The configuration of each stage of PPHGNet. such as the number of channels, stride, etc.
+        layer_num: int. Number of layers of HG_Block.
+        use_last_conv: boolean. Whether to use a 1x1 convolutional layer before the classification layer.
+        class_expand: int=2048. Number of channels for the last 1x1 convolutional layer.
+        dropout_prob: float. Parameters of dropout, 0.0 means dropout is not used.
+        class_num: int=1000. The number of classes.
+    Returns:
+        model: nn.Layer. Specific PPHGNet model depends on args.
+    """
+
+    def __init__(
+            self,
+            stem_channels,
+            stage_config,
+            layer_num,
+            in_channels=3,
+            det=False,
+            out_indices=None):
+        super().__init__()
+        self.det = det
+        self.out_indices = out_indices if out_indices is not None else [
+            0, 1, 2, 3
+        ]
+
+        # stem
+        stem_channels.insert(0, in_channels)
+        self.stem = nn.Sequential(* [
+            ConvBNAct(
+                in_channels=stem_channels[i],
+                out_channels=stem_channels[i + 1],
+                kernel_size=3,
+                stride=2 if i == 0 else 1) for i in range(
+                    len(stem_channels) - 1)
+        ])
+
+        if self.det:
+            self.pool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        # stages
+        self.stages = nn.ModuleList()
+        self.out_channels = []
+        for block_id, k in enumerate(stage_config):
+            in_channels, mid_channels, out_channels, block_num, downsample, stride = stage_config[
+                k]
+            self.stages.append(
+                HG_Stage(in_channels, mid_channels, out_channels, block_num,
+                         layer_num, downsample, stride))
+            if block_id in self.out_indices:
+                self.out_channels.append(out_channels)
+
+        if not self.det:
+            self.out_channels = stage_config["stage4"][2]
+
+        self._init_weights()
+
+    def _init_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.ones_(m.weight)
+                nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.Linear):
+                nn.init.zeros_(m.bias)
+
+    def forward(self, x):
+        x = self.stem(x)
+        if self.det:
+            x = self.pool(x)
+
+        out = []
+        for i, stage in enumerate(self.stages):
+            x = stage(x)
+            if self.det and i in self.out_indices:
+                out.append(x)
+        if self.det:
+            return out
+
+        if self.training:
+            x = F.adaptive_avg_pool2d(x, [1, 40])
+        else:
+            x = F.avg_pool2d(x, [3, 2])
+        return x
+
+def PPHGNet_small(pretrained=False, use_ssld=False, det=False, **kwargs):
+    """
+    PPHGNet_small
+    Args:
+        pretrained: bool=False or str. If `True` load pretrained parameters, `False` otherwise.
+                    If str, means the path of the pretrained model.
+        use_ssld: bool=False. Whether using distillation pretrained model when pretrained=True.
+    Returns:
+        model: nn.Layer. Specific `PPHGNet_small` model depends on args.
+    """
+    stage_config_det = {
+        # in_channels, mid_channels, out_channels, blocks, downsample
+        "stage1": [128, 128, 256, 1, False, 2],
+        "stage2": [256, 160, 512, 1, True, 2],
+        "stage3": [512, 192, 768, 2, True, 2],
+        "stage4": [768, 224, 1024, 1, True, 2],
+    }
+
+    stage_config_rec = {
+        # in_channels, mid_channels, out_channels, blocks, downsample
+        "stage1": [128, 128, 256, 1, True, [2, 1]],
+        "stage2": [256, 160, 512, 1, True, [1, 2]],
+        "stage3": [512, 192, 768, 2, True, [2, 1]],
+        "stage4": [768, 224, 1024, 1, True, [2, 1]],
+    }
+
+    model = PPHGNet(
+        stem_channels=[64, 64, 128],
+        stage_config=stage_config_det if det else stage_config_rec,
+        layer_num=6,
+        det=det,
+        **kwargs)
+    return model
+

python/rapidocr_torch/modeling/backbones/rec_lcnetv3.py

@@ -0,0 +1,474 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from rapidocr_torch.modeling.common import Activation
+
+# from paddle.nn.initializer import Constant, KaimingNormal
+# from paddle.nn import AdaptiveAvgPool2D, BatchNorm2D, Conv2D, Dropout, Hardsigmoid, Hardswish, Identity, Linear, ReLU
+# from paddle.regularizer import L2Decay
+
+NET_CONFIG_det = {
+    "blocks2":
+    #k, in_c, out_c, s, use_se
+    [[3, 16, 32, 1, False]],
+    "blocks3": [[3, 32, 64, 2, False], [3, 64, 64, 1, False]],
+    "blocks4": [[3, 64, 128, 2, False], [3, 128, 128, 1, False]],
+    "blocks5":
+    [[3, 128, 256, 2, False], [5, 256, 256, 1, False], [5, 256, 256, 1, False],
+     [5, 256, 256, 1, False], [5, 256, 256, 1, False]],
+    "blocks6": [[5, 256, 512, 2, True], [5, 512, 512, 1, True],
+                [5, 512, 512, 1, False], [5, 512, 512, 1, False]]
+}
+
+NET_CONFIG_rec = {
+    "blocks2":
+    #k, in_c, out_c, s, use_se
+    [[3, 16, 32, 1, False]],
+    "blocks3": [[3, 32, 64, 1, False], [3, 64, 64, 1, False]],
+    "blocks4": [[3, 64, 128, (2, 1), False], [3, 128, 128, 1, False]],
+    "blocks5":
+    [[3, 128, 256, (1, 2), False], [5, 256, 256, 1, False],
+     [5, 256, 256, 1, False], [5, 256, 256, 1, False], [5, 256, 256, 1, False]],
+    "blocks6": [[5, 256, 512, (2, 1), True], [5, 512, 512, 1, True],
+                [5, 512, 512, (2, 1), False], [5, 512, 512, 1, False]]
+}
+
+
+def make_divisible(v, divisor=16, min_value=None):
+    if min_value is None:
+        min_value = divisor
+    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+    if new_v < 0.9 * v:
+        new_v += divisor
+    return new_v
+
+
+class LearnableAffineBlock(nn.Module):
+    def __init__(self, scale_value=1.0, bias_value=0.0, lr_mult=1.0,
+                 lab_lr=0.1):
+        super().__init__()
+        self.scale = nn.Parameter(torch.Tensor([scale_value]))
+        self.bias = nn.Parameter(torch.Tensor([bias_value]))
+
+    def forward(self, x):
+        return self.scale * x + self.bias
+
+
+class ConvBNLayer(nn.Module):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride,
+                 groups=1,
+                 lr_mult=1.0):
+        super().__init__()
+        self.conv = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=(kernel_size - 1) // 2,
+            groups=groups,
+            bias=False)
+
+        self.bn = nn.BatchNorm2d(
+            out_channels,
+        )
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        return x
+
+
+class Act(nn.Module):
+    def __init__(self, act="hswish", lr_mult=1.0, lab_lr=0.1):
+        super().__init__()
+        if act == "hswish":
+            self.act = nn.Hardswish(inplace=True)
+        else:
+            assert act == "relu"
+            self.act = Activation(act)
+        self.lab = LearnableAffineBlock(lr_mult=lr_mult, lab_lr=lab_lr)
+
+    def forward(self, x):
+        return self.lab(self.act(x))
+
+
+class LearnableRepLayer(nn.Module):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 groups=1,
+                 num_conv_branches=1,
+                 lr_mult=1.0,
+                 lab_lr=0.1):
+        super().__init__()
+        self.is_repped = False
+        self.groups = groups
+        self.stride = stride
+        self.kernel_size = kernel_size
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.num_conv_branches = num_conv_branches
+        self.padding = (kernel_size - 1) // 2
+
+        self.identity = nn.BatchNorm2d(
+            num_features=in_channels,
+        ) if out_channels == in_channels and stride == 1 else None
+
+        self.conv_kxk = nn.ModuleList([
+            ConvBNLayer(
+                in_channels,
+                out_channels,
+                kernel_size,
+                stride,
+                groups=groups,
+                lr_mult=lr_mult) for _ in range(self.num_conv_branches)
+        ])
+
+        self.conv_1x1 = ConvBNLayer(
+            in_channels,
+            out_channels,
+            1,
+            stride,
+            groups=groups,
+            lr_mult=lr_mult) if kernel_size > 1 else None
+
+        self.lab = LearnableAffineBlock(lr_mult=lr_mult, lab_lr=lab_lr)
+        self.act = Act(lr_mult=lr_mult, lab_lr=lab_lr)
+
+    def forward(self, x):
+        # for export
+        if self.is_repped:
+            out = self.lab(self.reparam_conv(x))
+            if self.stride != 2:
+                out = self.act(out)
+            return out
+
+        out = 0
+        if self.identity is not None:
+            out += self.identity(x)
+
+        if self.conv_1x1 is not None:
+            out += self.conv_1x1(x)
+
+        for conv in self.conv_kxk:
+            out += conv(x)
+
+        out = self.lab(out)
+        if self.stride != 2:
+            out = self.act(out)
+        return out
+
+    def rep(self):
+        if self.is_repped:
+            return
+        kernel, bias = self._get_kernel_bias()
+        self.reparam_conv = nn.Conv2d(
+            in_channels=self.in_channels,
+            out_channels=self.out_channels,
+            kernel_size=self.kernel_size,
+            stride=self.stride,
+            padding=self.padding,
+            groups=self.groups)
+        self.reparam_conv.weight.data = kernel
+        self.reparam_conv.bias.data = bias
+        self.is_repped = True
+
+    def _pad_kernel_1x1_to_kxk(self, kernel1x1, pad):
+        if not isinstance(kernel1x1, torch.Tensor):
+            return 0
+        else:
+            return nn.functional.pad(kernel1x1, [pad, pad, pad, pad])
+
+    def _get_kernel_bias(self):
+        kernel_conv_1x1, bias_conv_1x1 = self._fuse_bn_tensor(self.conv_1x1)
+        kernel_conv_1x1 = self._pad_kernel_1x1_to_kxk(kernel_conv_1x1,
+                                                      self.kernel_size // 2)
+
+        kernel_identity, bias_identity = self._fuse_bn_tensor(self.identity)
+
+        kernel_conv_kxk = 0
+        bias_conv_kxk = 0
+        for conv in self.conv_kxk:
+            kernel, bias = self._fuse_bn_tensor(conv)
+            kernel_conv_kxk += kernel
+            bias_conv_kxk += bias
+
+        kernel_reparam = kernel_conv_kxk + kernel_conv_1x1 + kernel_identity
+        bias_reparam = bias_conv_kxk + bias_conv_1x1 + bias_identity
+        return kernel_reparam, bias_reparam
+
+    def _fuse_bn_tensor(self, branch):
+        if not branch:
+            return 0, 0
+        elif isinstance(branch, ConvBNLayer):
+            kernel = branch.conv.weight
+            running_mean = branch.bn._mean
+            running_var = branch.bn._variance
+            gamma = branch.bn.weight
+            beta = branch.bn.bias
+            eps = branch.bn._epsilon
+        else:
+            assert isinstance(branch, nn.BatchNorm2d)
+            if not hasattr(self, 'id_tensor'):
+                input_dim = self.in_channels // self.groups
+                kernel_value = torch.zeros(
+                    (self.in_channels, input_dim, self.kernel_size,
+                     self.kernel_size),
+                    dtype=branch.weight.dtype)
+                for i in range(self.in_channels):
+                    kernel_value[i, i % input_dim, self.kernel_size // 2,
+                                 self.kernel_size // 2] = 1
+                self.id_tensor = kernel_value
+            kernel = self.id_tensor
+            running_mean = branch._mean
+            running_var = branch._variance
+            gamma = branch.weight
+            beta = branch.bias
+            eps = branch._epsilon
+        std = (running_var + eps).sqrt()
+        t = (gamma / std).reshape((-1, 1, 1, 1))
+        return kernel * t, beta - running_mean * gamma / std
+
+
+class SELayer(nn.Module):
+    def __init__(self, channel, reduction=4, lr_mult=1.0):
+        super().__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.conv1 = nn.Conv2d(
+            in_channels=channel,
+            out_channels=channel // reduction,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        )
+        self.relu = nn.ReLU()
+        self.conv2 = nn.Conv2d(
+            in_channels=channel // reduction,
+            out_channels=channel,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        )
+        self.hardsigmoid = nn.Hardsigmoid(inplace=True)
+
+    def forward(self, x):
+        identity = x
+        x = self.avg_pool(x)
+        x = self.conv1(x)
+        x = self.relu(x)
+        x = self.conv2(x)
+        x = self.hardsigmoid(x)
+        x = identity * x
+        return x
+
+
+class LCNetV3Block(nn.Module):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 stride,
+                 dw_size,
+                 use_se=False,
+                 conv_kxk_num=4,
+                 lr_mult=1.0,
+                 lab_lr=0.1):
+        super().__init__()
+        self.use_se = use_se
+        self.dw_conv = LearnableRepLayer(
+            in_channels=in_channels,
+            out_channels=in_channels,
+            kernel_size=dw_size,
+            stride=stride,
+            groups=in_channels,
+            num_conv_branches=conv_kxk_num,
+            lr_mult=lr_mult,
+            lab_lr=lab_lr)
+        if use_se:
+            self.se = SELayer(in_channels, lr_mult=lr_mult)
+        self.pw_conv = LearnableRepLayer(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=1,
+            stride=1,
+            num_conv_branches=conv_kxk_num,
+            lr_mult=lr_mult,
+            lab_lr=lab_lr)
+
+    def forward(self, x):
+        x = self.dw_conv(x)
+        if self.use_se:
+            x = self.se(x)
+        x = self.pw_conv(x)
+        return x
+
+
+class PPLCNetV3(nn.Module):
+    def __init__(self,
+                 scale=1.0,
+                 conv_kxk_num=4,
+                 lr_mult_list=[1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
+                 lab_lr=0.1,
+                 det=False,
+                 **kwargs):
+        super().__init__()
+        self.scale = scale
+        self.lr_mult_list = lr_mult_list
+        self.det = det
+
+        self.net_config = NET_CONFIG_det if self.det else NET_CONFIG_rec
+
+        assert isinstance(self.lr_mult_list, (
+            list, tuple
+        )), "lr_mult_list should be in (list, tuple) but got {}".format(
+            type(self.lr_mult_list))
+        assert len(self.lr_mult_list
+                   ) == 6, "lr_mult_list length should be 6 but got {}".format(
+                       len(self.lr_mult_list))
+
+        self.conv1 = ConvBNLayer(
+            in_channels=3,
+            out_channels=make_divisible(16 * scale),
+            kernel_size=3,
+            stride=2,
+            lr_mult=self.lr_mult_list[0])
+
+        self.blocks2 = nn.Sequential(*[
+            LCNetV3Block(
+                in_channels=make_divisible(in_c * scale),
+                out_channels=make_divisible(out_c * scale),
+                dw_size=k,
+                stride=s,
+                use_se=se,
+                conv_kxk_num=conv_kxk_num,
+                lr_mult=self.lr_mult_list[1],
+                lab_lr=lab_lr)
+            for i, (k, in_c, out_c, s, se
+                    ) in enumerate(self.net_config["blocks2"])
+        ])
+
+        self.blocks3 = nn.Sequential(*[
+            LCNetV3Block(
+                in_channels=make_divisible(in_c * scale),
+                out_channels=make_divisible(out_c * scale),
+                dw_size=k,
+                stride=s,
+                use_se=se,
+                conv_kxk_num=conv_kxk_num,
+                lr_mult=self.lr_mult_list[2],
+                lab_lr=lab_lr)
+            for i, (k, in_c, out_c, s, se
+                    ) in enumerate(self.net_config["blocks3"])
+        ])
+
+        self.blocks4 = nn.Sequential(*[
+            LCNetV3Block(
+                in_channels=make_divisible(in_c * scale),
+                out_channels=make_divisible(out_c * scale),
+                dw_size=k,
+                stride=s,
+                use_se=se,
+                conv_kxk_num=conv_kxk_num,
+                lr_mult=self.lr_mult_list[3],
+                lab_lr=lab_lr)
+            for i, (k, in_c, out_c, s, se
+                    ) in enumerate(self.net_config["blocks4"])
+        ])
+
+        self.blocks5 = nn.Sequential(*[
+            LCNetV3Block(
+                in_channels=make_divisible(in_c * scale),
+                out_channels=make_divisible(out_c * scale),
+                dw_size=k,
+                stride=s,
+                use_se=se,
+                conv_kxk_num=conv_kxk_num,
+                lr_mult=self.lr_mult_list[4],
+                lab_lr=lab_lr)
+            for i, (k, in_c, out_c, s, se
+                    ) in enumerate(self.net_config["blocks5"])
+        ])
+
+        self.blocks6 = nn.Sequential(*[
+            LCNetV3Block(
+                in_channels=make_divisible(in_c * scale),
+                out_channels=make_divisible(out_c * scale),
+                dw_size=k,
+                stride=s,
+                use_se=se,
+                conv_kxk_num=conv_kxk_num,
+                lr_mult=self.lr_mult_list[5],
+                lab_lr=lab_lr)
+            for i, (k, in_c, out_c, s, se
+                    ) in enumerate(self.net_config["blocks6"])
+        ])
+        self.out_channels = make_divisible(512 * scale)
+
+        if self.det:
+            mv_c = [16, 24, 56, 480]
+            self.out_channels = [
+                make_divisible(self.net_config["blocks3"][-1][2] * scale),
+                make_divisible(self.net_config["blocks4"][-1][2] * scale),
+                make_divisible(self.net_config["blocks5"][-1][2] * scale),
+                make_divisible(self.net_config["blocks6"][-1][2] * scale),
+            ]
+
+            self.layer_list = nn.ModuleList([
+                nn.Conv2d(self.out_channels[0], int(mv_c[0] * scale), 1, 1, 0),
+                nn.Conv2d(self.out_channels[1], int(mv_c[1] * scale), 1, 1, 0),
+                nn.Conv2d(self.out_channels[2], int(mv_c[2] * scale), 1, 1, 0),
+                nn.Conv2d(self.out_channels[3], int(mv_c[3] * scale), 1, 1, 0)
+            ])
+            self.out_channels = [
+                int(mv_c[0] * scale), int(mv_c[1] * scale),
+                int(mv_c[2] * scale), int(mv_c[3] * scale)
+            ]
+
+    def forward(self, x):
+        out_list = []
+        x = self.conv1(x)
+        x = self.blocks2(x)
+        x = self.blocks3(x)
+        out_list.append(x)
+        x = self.blocks4(x)
+        out_list.append(x)
+        x = self.blocks5(x)
+        out_list.append(x)
+        x = self.blocks6(x)
+        out_list.append(x)
+
+        if self.det:
+            out_list[0] = self.layer_list[0](out_list[0])
+            out_list[1] = self.layer_list[1](out_list[1])
+            out_list[2] = self.layer_list[2](out_list[2])
+            out_list[3] = self.layer_list[3](out_list[3])
+            return out_list
+
+        if self.training:
+            x = F.adaptive_avg_pool2d(x, [1, 40])
+        else:
+            x = F.avg_pool2d(x, [3, 2])
+        return x

python/rapidocr_torch/modeling/backbones/rec_mobilenet_v3.py

@@ -0,0 +1,125 @@
+import os, sys
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from rapidocr_torch.modeling.common import Activation
+
+from .det_mobilenet_v3 import ResidualUnit, ConvBNLayer, make_divisible
+
+class MobileNetV3(nn.Module):
+    def __init__(self,
+                 in_channels=3,
+                 model_name='small',
+                 scale=0.5,
+                 large_stride=None,
+                 small_stride=None,
+                 **kwargs):
+        super(MobileNetV3, self).__init__()
+        if small_stride is None:
+            small_stride = [2, 2, 2, 2]
+        if large_stride is None:
+            large_stride = [1, 2, 2, 2]
+
+        assert isinstance(large_stride, list), "large_stride type must " \
+                                               "be list but got {}".format(type(large_stride))
+        assert isinstance(small_stride, list), "small_stride type must " \
+                                               "be list but got {}".format(type(small_stride))
+        assert len(large_stride) == 4, "large_stride length must be " \
+                                       "4 but got {}".format(len(large_stride))
+        assert len(small_stride) == 4, "small_stride length must be " \
+                                       "4 but got {}".format(len(small_stride))
+
+        if model_name == "large":
+            cfg = [
+                # k, exp, c,  se,     nl,  s,
+                [3, 16, 16, False, 'relu', large_stride[0]],
+                [3, 64, 24, False, 'relu', (large_stride[1], 1)],
+                [3, 72, 24, False, 'relu', 1],
+                [5, 72, 40, True, 'relu', (large_stride[2], 1)],
+                [5, 120, 40, True, 'relu', 1],
+                [5, 120, 40, True, 'relu', 1],
+                [3, 240, 80, False, 'hard_swish', 1],
+                [3, 200, 80, False, 'hard_swish', 1],
+                [3, 184, 80, False, 'hard_swish', 1],
+                [3, 184, 80, False, 'hard_swish', 1],
+                [3, 480, 112, True, 'hard_swish', 1],
+                [3, 672, 112, True, 'hard_swish', 1],
+                [5, 672, 160, True, 'hard_swish', (large_stride[3], 1)],
+                [5, 960, 160, True, 'hard_swish', 1],
+                [5, 960, 160, True, 'hard_swish', 1],
+            ]
+            cls_ch_squeeze = 960
+        elif model_name == "small":
+            cfg = [
+                # k, exp, c,  se,     nl,  s,
+                [3, 16, 16, True, 'relu', (small_stride[0], 1)],
+                [3, 72, 24, False, 'relu', (small_stride[1], 1)],
+                [3, 88, 24, False, 'relu', 1],
+                [5, 96, 40, True, 'hard_swish', (small_stride[2], 1)],
+                [5, 240, 40, True, 'hard_swish', 1],
+                [5, 240, 40, True, 'hard_swish', 1],
+                [5, 120, 48, True, 'hard_swish', 1],
+                [5, 144, 48, True, 'hard_swish', 1],
+                [5, 288, 96, True, 'hard_swish', (small_stride[3], 1)],
+                [5, 576, 96, True, 'hard_swish', 1],
+                [5, 576, 96, True, 'hard_swish', 1],
+            ]
+            cls_ch_squeeze = 576
+        else:
+            raise NotImplementedError("mode[" + model_name +
+                                      "_model] is not implemented!")
+
+        supported_scale = [0.35, 0.5, 0.75, 1.0, 1.25]
+        assert scale in supported_scale, \
+            "supported scales are {} but input scale is {}".format(supported_scale, scale)
+
+        inplanes = 16
+        # conv1
+        self.conv1 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=make_divisible(inplanes * scale),
+            kernel_size=3,
+            stride=2,
+            padding=1,
+            groups=1,
+            if_act=True,
+            act='hard_swish',
+            name='conv1')
+        i = 0
+        block_list = []
+        inplanes = make_divisible(inplanes * scale)
+        for (k, exp, c, se, nl, s) in cfg:
+            block_list.append(
+                ResidualUnit(
+                    in_channels=inplanes,
+                    mid_channels=make_divisible(scale * exp),
+                    out_channels=make_divisible(scale * c),
+                    kernel_size=k,
+                    stride=s,
+                    use_se=se,
+                    act=nl,
+                    name='conv' + str(i + 2)))
+            inplanes = make_divisible(scale * c)
+            i += 1
+        self.blocks = nn.Sequential(*block_list)
+
+        self.conv2 = ConvBNLayer(
+            in_channels=inplanes,
+            out_channels=make_divisible(scale * cls_ch_squeeze),
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            groups=1,
+            if_act=True,
+            act='hard_swish',
+            name='conv_last')
+
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
+        self.out_channels = make_divisible(scale * cls_ch_squeeze)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.blocks(x)
+        x = self.conv2(x)
+        x = self.pool(x)
+        return x

python/rapidocr_torch/modeling/backbones/rec_svtrnet.py

@@ -0,0 +1,582 @@
+import torch
+import torch.nn as nn
+from rapidocr_torch.modeling.common import Activation
+import numpy as np
+
+def drop_path(x, drop_prob=0., training=False):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ...
+    """
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = torch.as_tensor(1 - drop_prob)
+    shape = (x.shape[0], ) + (1, ) * (x.ndim - 1)
+    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype)
+    random_tensor = torch.floor(random_tensor)  # binarize
+    output = x.divide(keep_prob) * random_tensor
+    return output
+
+
+class ConvBNLayer(nn.Module):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size=3,
+                 stride=1,
+                 padding=0,
+                 bias_attr=False,
+                 groups=1,
+                 act='gelu'):
+        super().__init__()
+        self.conv = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            groups=groups,
+            bias=bias_attr)
+        self.norm = nn.BatchNorm2d(out_channels)
+        self.act = Activation(act_type=act, inplace=True)
+
+    def forward(self, inputs):
+        out = self.conv(inputs)
+        out = self.norm(out)
+        out = self.act(out)
+        return out
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+
+class Identity(nn.Module):
+    def __init__(self):
+        super(Identity, self).__init__()
+
+    def forward(self, input):
+        return input
+
+
+class Mlp(nn.Module):
+    def __init__(self,
+                 in_features,
+                 hidden_features=None,
+                 out_features=None,
+                 act_layer='gelu',
+                 drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = Activation(act_type=act_layer, inplace=True)
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class ConvMixer(nn.Module):
+    def __init__(
+            self,
+            dim,
+            num_heads=8,
+            HW=[8, 25],
+            local_k=[3, 3], ):
+        super().__init__()
+        self.HW = HW
+        self.dim = dim
+        self.local_mixer = nn.Conv2d(
+            dim,
+            dim,
+            local_k,
+            1, [local_k[0] // 2, local_k[1] // 2],
+            groups=num_heads,
+            )
+
+    def forward(self, x):
+        h = self.HW[0]
+        w = self.HW[1]
+        x = x.transpose([0, 2, 1]).reshape([0, self.dim, h, w])
+        x = self.local_mixer(x)
+        x = x.flatten(2).permute(0, 2, 1)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(self,
+                 dim,
+                 num_heads=8,
+                 mixer='Global',
+                 HW=[8, 25],
+                 local_k=[7, 11],
+                 qkv_bias=False,
+                 qk_scale=None,
+                 attn_drop=0.,
+                 proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.HW = HW
+        if HW is not None:
+            H = HW[0]
+            W = HW[1]
+            self.N = H * W
+            self.C = dim
+        if mixer == 'Local' and HW is not None:
+            hk = local_k[0]
+            wk = local_k[1]
+            mask = torch.ones(H * W, H + hk - 1, W + wk - 1, dtype=torch.float32)
+            for h in range(0, H):
+                for w in range(0, W):
+                    mask[h * W + w, h:h + hk, w:w + wk] = 0.
+            mask_paddle = mask[:, hk // 2:H + hk // 2, wk // 2:W + wk //
+                               2].flatten(1)
+            mask_inf = torch.full([H * W, H * W], fill_value=float("-Inf"), dtype=torch.float32)
+            mask = torch.where(mask_paddle < 1, mask_paddle, mask_inf)
+            self.mask = mask.unsqueeze(0).unsqueeze(1)
+            # self.mask = mask[None, None, :]
+        self.mixer = mixer
+
+    def forward(self, x):
+        if self.HW is not None:
+            N = self.N
+            C = self.C
+        else:
+            _, N, C = x.shape
+        qkv = self.qkv(x)
+        qkv = qkv.reshape((-1, N, 3, self.num_heads, C // self.num_heads)).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
+
+        attn = (q.matmul(k.permute(0, 1, 3, 2)))
+        if self.mixer == 'Local':
+            attn += self.mask
+        attn = nn.functional.softmax(attn, dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn.matmul(v)).permute(0, 2, 1, 3).reshape((-1, N, C))
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(self,
+                 dim,
+                 num_heads,
+                 mixer='Global',
+                 local_mixer=[7, 11],
+                 HW=None,
+                 mlp_ratio=4.,
+                 qkv_bias=False,
+                 qk_scale=None,
+                 drop=0.,
+                 attn_drop=0.,
+                 drop_path=0.,
+                 act_layer='gelu',
+                 norm_layer='nn.LayerNorm',
+                 epsilon=1e-6,
+                 prenorm=True):
+        super().__init__()
+        if isinstance(norm_layer, str):
+            self.norm1 = eval(norm_layer)(dim, eps=epsilon)
+        else:
+            self.norm1 = norm_layer(dim)
+        if mixer == 'Global' or mixer == 'Local':
+            self.mixer = Attention(
+                dim,
+                num_heads=num_heads,
+                mixer=mixer,
+                HW=HW,
+                local_k=local_mixer,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                attn_drop=attn_drop,
+                proj_drop=drop)
+        elif mixer == 'Conv':
+            self.mixer = ConvMixer(
+                dim, num_heads=num_heads, HW=HW, local_k=local_mixer)
+        else:
+            raise TypeError("The mixer must be one of [Global, Local, Conv]")
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else Identity()
+        if isinstance(norm_layer, str):
+            self.norm2 = eval(norm_layer)(dim, eps=epsilon)
+        else:
+            self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp_ratio = mlp_ratio
+        self.mlp = Mlp(in_features=dim,
+                       hidden_features=mlp_hidden_dim,
+                       act_layer=act_layer,
+                       drop=drop)
+        self.prenorm = prenorm
+
+    def forward(self, x):
+        if self.prenorm:
+            x = self.norm1(x + self.drop_path(self.mixer(x)))
+            x = self.norm2(x + self.drop_path(self.mlp(x)))
+        else:
+            x = x + self.drop_path(self.mixer(self.norm1(x)))
+            x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+
+    def __init__(self,
+                 img_size=[32, 100],
+                 in_channels=3,
+                 embed_dim=768,
+                 sub_num=2,
+                 patch_size=[4, 4],
+                 mode='pope',
+                 ):
+        super().__init__()
+        num_patches = (img_size[1] // (2 ** sub_num)) * \
+                      (img_size[0] // (2 ** sub_num))
+        self.img_size = img_size
+        self.num_patches = num_patches
+        self.embed_dim = embed_dim
+        self.norm = None
+        if mode == 'pope':
+            if sub_num == 2:
+                self.proj = nn.Sequential(
+                    ConvBNLayer(
+                        in_channels=in_channels,
+                        out_channels=embed_dim // 2,
+                        kernel_size=3,
+                        stride=2,
+                        padding=1,
+                        act='gelu',
+                        bias_attr=True),
+                    ConvBNLayer(
+                        in_channels=embed_dim // 2,
+                        out_channels=embed_dim,
+                        kernel_size=3,
+                        stride=2,
+                        padding=1,
+                        act='gelu',
+                        bias_attr=True))
+            if sub_num == 3:
+                self.proj = nn.Sequential(
+                    ConvBNLayer(
+                        in_channels=in_channels,
+                        out_channels=embed_dim // 4,
+                        kernel_size=3,
+                        stride=2,
+                        padding=1,
+                        act='gelu',
+                        bias_attr=True),
+                    ConvBNLayer(
+                        in_channels=embed_dim // 4,
+                        out_channels=embed_dim // 2,
+                        kernel_size=3,
+                        stride=2,
+                        padding=1,
+                        act='gelu',
+                        bias_attr=True),
+                    ConvBNLayer(
+                        in_channels=embed_dim // 2,
+                        out_channels=embed_dim,
+                        kernel_size=3,
+                        stride=2,
+                        padding=1,
+                        act='gelu',
+                        bias_attr=True))
+        elif mode == 'linear':
+            self.proj = nn.Conv2d(
+                1, embed_dim, kernel_size=patch_size, stride=patch_size)
+            self.num_patches = img_size[0] // patch_size[0] * img_size[
+                1] // patch_size[1]
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        assert H == self.img_size[0] and W == self.img_size[1], \
+            "Input image size ({}*{}) doesn't match model ({}*{}).".format(
+                H,W,self.img_size[0],self.img_size[1]
+            )
+        x = self.proj(x).flatten(2).permute(0, 2, 1)
+        return x
+
+
+class SubSample(nn.Module):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 types='Pool',
+                 stride=[2, 1],
+                 sub_norm='nn.LayerNorm',
+                 act=None):
+        super().__init__()
+        self.types = types
+        if types == 'Pool':
+            self.avgpool = nn.AvgPool2d(
+                kernel_size=[3, 5], stride=stride, padding=[1, 2])
+            self.maxpool = nn.MaxPool2d(
+                kernel_size=[3, 5], stride=stride, padding=[1, 2])
+            self.proj = nn.Linear(in_channels, out_channels)
+        else:
+            self.conv = nn.Conv2d(
+                in_channels,
+                out_channels,
+                kernel_size=3,
+                stride=stride,
+                padding=1,
+                )
+        self.norm = eval(sub_norm)(out_channels)
+        if act is not None:
+            self.act = act()
+        else:
+            self.act = None
+
+    def forward(self, x):
+
+        if self.types == 'Pool':
+            x1 = self.avgpool(x)
+            x2 = self.maxpool(x)
+            x = (x1 + x2) * 0.5
+            out = self.proj(x.flatten(2).permute(0, 2, 1))
+        else:
+            x = self.conv(x)
+            out = x.flatten(2).permute(0, 2, 1)
+        out = self.norm(out)
+        if self.act is not None:
+            out = self.act(out)
+
+        return out
+
+
+class SVTRNet(nn.Module):
+    def __init__(
+            self,
+            img_size=[32, 100],
+            in_channels=3,
+            embed_dim=[64, 128, 256],
+            depth=[3, 6, 3],
+            num_heads=[2, 4, 8],
+            mixer=['Local'] * 6 + ['Global'] *
+            6,  # Local atten, Global atten, Conv
+            local_mixer=[[7, 11], [7, 11], [7, 11]],
+            patch_merging='Conv',  # Conv, Pool, None
+            mlp_ratio=4,
+            qkv_bias=True,
+            qk_scale=None,
+            drop_rate=0.,
+            last_drop=0.0,
+            attn_drop_rate=0.,
+            drop_path_rate=0.1,
+            norm_layer='nn.LayerNorm',
+            sub_norm='nn.LayerNorm',
+            epsilon=1e-6,
+            out_channels=192,
+            out_char_num=25,
+            block_unit='Block',
+            act='gelu',
+            last_stage=True,
+            sub_num=2,
+            prenorm=True,
+            use_lenhead=False,
+            **kwargs):
+        super().__init__()
+        self.img_size = img_size
+        self.embed_dim = embed_dim
+        self.out_channels = out_channels
+        self.prenorm = prenorm
+        patch_merging = None if patch_merging != 'Conv' and patch_merging != 'Pool' else patch_merging
+        self.patch_embed = PatchEmbed(
+            img_size=img_size,
+            in_channels=in_channels,
+            embed_dim=embed_dim[0],
+            sub_num=sub_num)
+        num_patches = self.patch_embed.num_patches
+        self.HW = [img_size[0] // (2**sub_num), img_size[1] // (2**sub_num)]
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim[0]))
+        self.pos_drop = nn.Dropout(p=drop_rate)
+        Block_unit = eval(block_unit)
+
+        dpr = np.linspace(0, drop_path_rate, sum(depth))
+        self.blocks1 = nn.ModuleList([
+            Block_unit(
+                dim=embed_dim[0],
+                num_heads=num_heads[0],
+                mixer=mixer[0:depth[0]][i],
+                HW=self.HW,
+                local_mixer=local_mixer[0],
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop_rate,
+                act_layer=act,
+                attn_drop=attn_drop_rate,
+                drop_path=dpr[0:depth[0]][i],
+                norm_layer=norm_layer,
+                epsilon=epsilon,
+                prenorm=prenorm) for i in range(depth[0])
+        ])
+        if patch_merging is not None:
+            self.sub_sample1 = SubSample(
+                embed_dim[0],
+                embed_dim[1],
+                sub_norm=sub_norm,
+                stride=[2, 1],
+                types=patch_merging)
+            HW = [self.HW[0] // 2, self.HW[1]]
+        else:
+            HW = self.HW
+        self.patch_merging = patch_merging
+        self.blocks2 = nn.ModuleList([
+            Block_unit(
+                dim=embed_dim[1],
+                num_heads=num_heads[1],
+                mixer=mixer[depth[0]:depth[0] + depth[1]][i],
+                HW=HW,
+                local_mixer=local_mixer[1],
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop_rate,
+                act_layer=act,
+                attn_drop=attn_drop_rate,
+                drop_path=dpr[depth[0]:depth[0] + depth[1]][i],
+                norm_layer=norm_layer,
+                epsilon=epsilon,
+                prenorm=prenorm) for i in range(depth[1])
+        ])
+        if patch_merging is not None:
+            self.sub_sample2 = SubSample(
+                embed_dim[1],
+                embed_dim[2],
+                sub_norm=sub_norm,
+                stride=[2, 1],
+                types=patch_merging)
+            HW = [self.HW[0] // 4, self.HW[1]]
+        else:
+            HW = self.HW
+        self.blocks3 = nn.ModuleList([
+            Block_unit(
+                dim=embed_dim[2],
+                num_heads=num_heads[2],
+                mixer=mixer[depth[0] + depth[1]:][i],
+                HW=HW,
+                local_mixer=local_mixer[2],
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop_rate,
+                act_layer=act,
+                attn_drop=attn_drop_rate,
+                drop_path=dpr[depth[0] + depth[1]:][i],
+                norm_layer=norm_layer,
+                epsilon=epsilon,
+                prenorm=prenorm) for i in range(depth[2])
+        ])
+        self.last_stage = last_stage
+        if last_stage:
+            self.avg_pool = nn.AdaptiveAvgPool2d([1, out_char_num])
+            self.last_conv = nn.Conv2d(
+                in_channels=embed_dim[2],
+                out_channels=self.out_channels,
+                kernel_size=1,
+                stride=1,
+                padding=0,
+                bias=False)
+            self.hardswish = Activation('hard_swish', inplace=True) #nn.Hardswish()
+            # self.dropout = nn.Dropout(p=last_drop, mode="downscale_in_infer")
+            self.dropout = nn.Dropout(p=last_drop)
+        if not prenorm:
+            self.norm = eval(norm_layer)(embed_dim[-1], eps=epsilon)
+        self.use_lenhead = use_lenhead
+        if use_lenhead:
+            self.len_conv = nn.Linear(embed_dim[2], self.out_channels)
+            self.hardswish_len = Activation('hard_swish', inplace=True)# nn.Hardswish()
+            self.dropout_len = nn.Dropout(
+                p=last_drop)
+
+        torch.nn.init.xavier_normal_(self.pos_embed)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        # weight initialization
+        if isinstance(m, nn.Conv2d):
+            nn.init.kaiming_normal_(m.weight, mode='fan_out')
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.BatchNorm2d):
+            nn.init.ones_(m.weight)
+            nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.Linear):
+            nn.init.normal_(m.weight, 0, 0.01)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.ConvTranspose2d):
+            nn.init.kaiming_normal_(m.weight, mode='fan_out')
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.ones_(m.weight)
+            nn.init.zeros_(m.bias)
+
+    def forward_features(self, x):
+        x = self.patch_embed(x)
+        x = x + self.pos_embed
+        x = self.pos_drop(x)
+        for blk in self.blocks1:
+            x = blk(x)
+        if self.patch_merging is not None:
+            x = self.sub_sample1(
+                x.permute(0, 2, 1).reshape(
+                    [-1, self.embed_dim[0], self.HW[0], self.HW[1]]))
+        for blk in self.blocks2:
+            x = blk(x)
+        if self.patch_merging is not None:
+            x = self.sub_sample2(
+                x.permute(0, 2, 1).reshape(
+                    [-1, self.embed_dim[1], self.HW[0] // 2, self.HW[1]]))
+        for blk in self.blocks3:
+            x = blk(x)
+        if not self.prenorm:
+            x = self.norm(x)
+        return x
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        if self.use_lenhead:
+            len_x = self.len_conv(x.mean(1))
+            len_x = self.dropout_len(self.hardswish_len(len_x))
+        if self.last_stage:
+            if self.patch_merging is not None:
+                h = self.HW[0] // 4
+            else:
+                h = self.HW[0]
+            x = self.avg_pool(
+                x.permute(0, 2, 1).reshape(
+                    [-1, self.embed_dim[2], h, self.HW[1]]))
+            x = self.last_conv(x)
+            x = self.hardswish(x)
+            x = self.dropout(x)
+        if self.use_lenhead:
+            return x, len_x
+        return x

python/rapidocr_torch/modeling/common.py

@@ -0,0 +1,73 @@
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class Hswish(nn.Module):
+    def __init__(self, inplace=True):
+        super(Hswish, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        return x * F.relu6(x + 3., inplace=self.inplace) / 6.
+
+# out = max(0, min(1, slop*x+offset))
+# paddle.fluid.layers.hard_sigmoid(x, slope=0.2, offset=0.5, name=None)
+class Hsigmoid(nn.Module):
+    def __init__(self, inplace=True):
+        super(Hsigmoid, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        # torch: F.relu6(x + 3., inplace=self.inplace) / 6.
+        # paddle: F.relu6(1.2 * x + 3., inplace=self.inplace) / 6.
+        return F.relu6(1.2 * x + 3., inplace=self.inplace) / 6.
+
+class GELU(nn.Module):
+    def __init__(self, inplace=True):
+        super(GELU, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        return torch.nn.functional.gelu(x)
+
+
+class Swish(nn.Module):
+    def __init__(self, inplace=True):
+        super(Swish, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        if self.inplace:
+            x.mul_(torch.sigmoid(x))
+            return x
+        else:
+            return x*torch.sigmoid(x)
+
+
+class Activation(nn.Module):
+    def __init__(self, act_type, inplace=True):
+        super(Activation, self).__init__()
+        act_type = act_type.lower()
+        if act_type == 'relu':
+            self.act = nn.ReLU(inplace=inplace)
+        elif act_type == 'relu6':
+            self.act = nn.ReLU6(inplace=inplace)
+        elif act_type == 'sigmoid':
+            raise NotImplementedError
+        elif act_type == 'hard_sigmoid':
+            self.act = Hsigmoid(inplace)#nn.Hardsigmoid(inplace=inplace)#Hsigmoid(inplace)#
+        elif act_type == 'hard_swish' or act_type == 'hswish':
+            self.act = Hswish(inplace=inplace)
+        elif act_type == 'leakyrelu':
+            self.act = nn.LeakyReLU(inplace=inplace)
+        elif act_type == 'gelu':
+            self.act = GELU(inplace=inplace)
+        elif act_type == 'swish':
+            self.act = Swish(inplace=inplace)
+        else:
+            raise NotImplementedError
+
+    def forward(self, inputs):
+        return self.act(inputs)

python/rapidocr_torch/modeling/heads/__init__.py

@@ -0,0 +1,36 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+__all__ = ['build_head']
+
+
+def build_head(config, **kwargs):
+    # det head
+    from .det_db_head import DBHead, PFHeadLocal
+    # rec head
+    from .rec_ctc_head import CTCHead
+    from .rec_multi_head import MultiHead
+
+    # cls head
+    from .cls_head import ClsHead
+    support_dict = [
+        'DBHead',  'CTCHead', 'ClsHead', 'MultiHead', 'PFHeadLocal',
+    ]
+
+    module_name = config.pop('name')
+    char_num = config.pop('char_num', 6625)
+    assert module_name in support_dict, Exception('head only support {}'.format(
+        support_dict))
+    module_class = eval(module_name)(**config, **kwargs)
+    return module_class

python/rapidocr_torch/modeling/heads/cls_head.py

@@ -0,0 +1,26 @@
+import os, sys
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class ClsHead(nn.Module):
+    """
+    Class orientation
+    Args:
+        params(dict): super parameters for build Class network
+    """
+
+    def __init__(self, in_channels, class_dim, **kwargs):
+        super(ClsHead, self).__init__()
+        self.pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Linear(
+            in_channels,
+            class_dim,
+            bias=True)
+
+    def forward(self, x):
+        x = self.pool(x)
+        x = torch.reshape(x, shape=[x.shape[0], x.shape[1]])
+        x = self.fc(x)
+        x = F.softmax(x, dim=1)
+        return x

python/rapidocr_torch/modeling/heads/det_db_head.py

@@ -0,0 +1,122 @@
+import os, sys
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from rapidocr_torch.modeling.common import Activation
+from rapidocr_torch.modeling.backbones.det_mobilenet_v3 import ConvBNLayer
+
+class Head(nn.Module):
+    def __init__(self, in_channels, **kwargs):
+        super(Head, self).__init__()
+        self.conv1 = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=in_channels // 4,
+            kernel_size=3,
+            padding=1,
+            bias=False)
+        self.conv_bn1 = nn.BatchNorm2d(
+            in_channels // 4)
+        self.relu1 = Activation(act_type='relu')
+
+        self.conv2 = nn.ConvTranspose2d(
+            in_channels=in_channels // 4,
+            out_channels=in_channels // 4,
+            kernel_size=2,
+            stride=2)
+        self.conv_bn2 = nn.BatchNorm2d(
+            in_channels // 4)
+        self.relu2 = Activation(act_type='relu')
+
+        self.conv3 = nn.ConvTranspose2d(
+            in_channels=in_channels // 4,
+            out_channels=1,
+            kernel_size=2,
+            stride=2)
+
+    def forward(self, x, return_f=False):
+        x = self.conv1(x)
+        x = self.conv_bn1(x)
+        x = self.relu1(x)
+        x = self.conv2(x)
+        x = self.conv_bn2(x)
+        x = self.relu2(x)
+        if return_f is True:
+            f = x
+        x = self.conv3(x)
+        x = torch.sigmoid(x)
+        if return_f is True:
+            return x, f
+        return x
+
+
+class DBHead(nn.Module):
+    """
+    Differentiable Binarization (DB) for text detection:
+        see https://arxiv.org/abs/1911.08947
+    args:
+        params(dict): super parameters for build DB network
+    """
+
+    def __init__(self, in_channels, k=50, **kwargs):
+        super(DBHead, self).__init__()
+        self.k = k
+        binarize_name_list = [
+            'conv2d_56', 'batch_norm_47', 'conv2d_transpose_0', 'batch_norm_48',
+            'conv2d_transpose_1', 'binarize'
+        ]
+        thresh_name_list = [
+            'conv2d_57', 'batch_norm_49', 'conv2d_transpose_2', 'batch_norm_50',
+            'conv2d_transpose_3', 'thresh'
+        ]
+        self.binarize = Head(in_channels, **kwargs)# binarize_name_list)
+        self.thresh = Head(in_channels, **kwargs)#thresh_name_list)
+
+    def step_function(self, x, y):
+        return torch.reciprocal(1 + torch.exp(-self.k * (x - y)))
+
+    def forward(self, x):
+        # shrink_maps = self.binarize(x)
+        return self.binarize(x)
+        # if not self.training:
+        #     return {'maps': shrink_maps}
+
+        # threshold_maps = self.thresh(x)
+        # binary_maps = self.step_function(shrink_maps, threshold_maps)
+        # y = torch.cat([shrink_maps, threshold_maps, binary_maps], dim=1)
+        # return y
+
+class LocalModule(nn.Module):
+    def __init__(self, in_c, mid_c, use_distance=True):
+        super(self.__class__, self).__init__()
+        self.last_3 = ConvBNLayer(in_c + 1, mid_c, 3, 1, 1, act='relu')
+        self.last_1 = nn.Conv2d(mid_c, 1, 1, 1, 0)
+
+    def forward(self, x, init_map, distance_map):
+        outf = torch.cat([init_map, x], dim=1)
+        # last Conv
+        out = self.last_1(self.last_3(outf))
+        return out
+
+class PFHeadLocal(DBHead):
+    def __init__(self, in_channels, k=50, mode='small', **kwargs):
+        super(PFHeadLocal, self).__init__(in_channels, k, **kwargs)
+        self.mode = mode
+
+        self.up_conv = nn.Upsample(scale_factor=2, mode="nearest")
+        if self.mode == 'large':
+            self.cbn_layer = LocalModule(in_channels // 4, in_channels // 4)
+        elif self.mode == 'small':
+            self.cbn_layer = LocalModule(in_channels // 4, in_channels // 8)
+
+    def forward(self, x, targets=None):
+        shrink_maps, f = self.binarize(x, return_f=True)
+        base_maps = shrink_maps
+        cbn_maps = self.cbn_layer(self.up_conv(f), shrink_maps, None)
+        cbn_maps = F.sigmoid(cbn_maps)
+        return 0.5 * (base_maps + cbn_maps)
+            # return {'maps': 0.5 * (base_maps + cbn_maps), 'cbn_maps': cbn_maps}
+
+        # threshold_maps = self.thresh(x)
+        # binary_maps = self.step_function(shrink_maps, threshold_maps)
+        # y = torch.cat([cbn_maps, threshold_maps, binary_maps], dim=1)
+        # return y

python/rapidocr_torch/modeling/heads/rec_ctc_head.py

@@ -0,0 +1,53 @@
+import os, sys
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class CTCHead(nn.Module):
+    def __init__(self,
+                 in_channels,
+                 out_channels=6625,
+                 fc_decay=0.0004,
+                 mid_channels=None,
+                 return_feats=False,
+                 **kwargs):
+        super(CTCHead, self).__init__()
+        if mid_channels is None:
+            self.fc = nn.Linear(
+                in_channels,
+                out_channels,
+                bias=True,)
+        else:
+            self.fc1 = nn.Linear(
+                in_channels,
+                mid_channels,
+                bias=True,
+            )
+            self.fc2 = nn.Linear(
+                mid_channels,
+                out_channels,
+                bias=True,
+            )
+
+        self.out_channels = out_channels
+        self.mid_channels = mid_channels
+        self.return_feats = return_feats
+
+
+    def forward(self, x, labels=None):
+        if self.mid_channels is None:
+            predicts = self.fc(x)
+        else:
+            x = self.fc1(x)
+            predicts = self.fc2(x)
+
+        if self.return_feats:
+            result = (x, predicts)
+        else:
+            result = predicts
+
+        if not self.training:
+            predicts = F.softmax(predicts, dim=2)
+            result = predicts
+
+        return result

python/rapidocr_torch/modeling/heads/rec_multi_head.py

@@ -0,0 +1,55 @@
+import torch
+import torch.nn as nn
+
+from rapidocr_torch.modeling.necks.rnn import Im2Seq, SequenceEncoder
+from .rec_ctc_head import CTCHead
+
+class FCTranspose(nn.Module):
+    def __init__(self, in_channels, out_channels, only_transpose=False):
+        super().__init__()
+        self.only_transpose = only_transpose
+        if not self.only_transpose:
+            self.fc = nn.Linear(in_channels, out_channels, bias=False)
+
+    def forward(self, x):
+        if self.only_transpose:
+            return x.permute([0, 2, 1])
+        else:
+            return self.fc(x.permute([0, 2, 1]))
+
+
+class MultiHead(nn.Module):
+    def __init__(self, in_channels, out_channels_list, **kwargs):
+        super().__init__()
+        self.head_list = kwargs.pop('head_list')
+
+        self.gtc_head = 'sar'
+        assert len(self.head_list) >= 2
+        for idx, head_name in enumerate(self.head_list):
+            name = list(head_name)[0]
+            if name == 'SARHead':
+                pass
+
+            elif name == 'NRTRHead':
+                pass
+            elif name == 'CTCHead':
+                # ctc neck
+                self.encoder_reshape = Im2Seq(in_channels)
+                neck_args = self.head_list[idx][name]['Neck']
+                encoder_type = neck_args.pop('name')
+                self.ctc_encoder = SequenceEncoder(in_channels=in_channels, \
+                                                   encoder_type=encoder_type, **neck_args)
+                # ctc head
+                head_args = self.head_list[idx][name].get('Head', {})
+                if head_args is None:
+                    head_args = {}
+                self.ctc_head = eval(name)(in_channels=self.ctc_encoder.out_channels, \
+                                           out_channels=out_channels_list['CTCLabelDecode'], **head_args)
+            else:
+                raise NotImplementedError(
+                    '{} is not supported in MultiHead yet'.format(name))
+
+    def forward(self, x, data=None):
+        ctc_encoder = self.ctc_encoder(x)
+        return self.ctc_head(ctc_encoder)
+

python/rapidocr_torch/modeling/necks/__init__.py

@@ -0,0 +1,27 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+__all__ = ['build_neck']
+
+
+def build_neck(config):
+    from .db_fpn import DBFPN, RSEFPN, LKPAN
+    from .rnn import SequenceEncoder
+    support_dict = ['DBFPN', 'SequenceEncoder', 'RSEFPN', 'LKPAN']
+
+    module_name = config.pop('name')
+    assert module_name in support_dict, Exception('neck only support {}'.format(
+        support_dict))
+    module_class = eval(module_name)(**config)
+    return module_class

python/rapidocr_torch/modeling/necks/db_fpn.py

@@ -0,0 +1,416 @@
+import os, sys
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from rapidocr_torch.modeling.backbones.det_mobilenet_v3 import SEModule
+from rapidocr_torch.modeling.necks.intracl import IntraCLBlock
+
+
+def hard_swish(x, inplace=True):
+    return x * F.relu6(x + 3., inplace=inplace) / 6.
+
+
+class DSConv(nn.Module):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 padding,
+                 stride=1,
+                 groups=None,
+                 if_act=True,
+                 act="relu",
+                 **kwargs):
+        super(DSConv, self).__init__()
+        if groups == None:
+            groups = in_channels
+        self.if_act = if_act
+        self.act = act
+        self.conv1 = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=in_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            groups=groups,
+            bias=False)
+
+        self.bn1 = nn.BatchNorm2d(in_channels)
+
+        self.conv2 = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=int(in_channels * 4),
+            kernel_size=1,
+            stride=1,
+            bias=False)
+
+        self.bn2 = nn.BatchNorm2d(int(in_channels * 4))
+
+        self.conv3 = nn.Conv2d(
+            in_channels=int(in_channels * 4),
+            out_channels=out_channels,
+            kernel_size=1,
+            stride=1,
+            bias=False)
+        self._c = [in_channels, out_channels]
+        if in_channels != out_channels:
+            self.conv_end = nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=1,
+                stride=1,
+                bias=False)
+
+    def forward(self, inputs):
+
+        x = self.conv1(inputs)
+        x = self.bn1(x)
+
+        x = self.conv2(x)
+        x = self.bn2(x)
+        if self.if_act:
+            if self.act == "relu":
+                x = F.relu(x)
+            elif self.act == "hardswish":
+                x = hard_swish(x)
+            else:
+                print("The activation function({}) is selected incorrectly.".
+                      format(self.act))
+                exit()
+
+        x = self.conv3(x)
+        if self._c[0] != self._c[1]:
+            x = x + self.conv_end(inputs)
+        return x
+
+
+class DBFPN(nn.Module):
+    def __init__(self, in_channels, out_channels, use_asf=False, **kwargs):
+        super(DBFPN, self).__init__()
+        self.out_channels = out_channels
+        self.use_asf = use_asf
+
+        self.in2_conv = nn.Conv2d(
+            in_channels=in_channels[0],
+            out_channels=self.out_channels,
+            kernel_size=1,
+            bias=False)
+        self.in3_conv = nn.Conv2d(
+            in_channels=in_channels[1],
+            out_channels=self.out_channels,
+            kernel_size=1,
+            bias=False)
+        self.in4_conv = nn.Conv2d(
+            in_channels=in_channels[2],
+            out_channels=self.out_channels,
+            kernel_size=1,
+            bias=False)
+        self.in5_conv = nn.Conv2d(
+            in_channels=in_channels[3],
+            out_channels=self.out_channels,
+            kernel_size=1,
+            bias=False)
+        self.p5_conv = nn.Conv2d(
+            in_channels=self.out_channels,
+            out_channels=self.out_channels // 4,
+            kernel_size=3,
+            padding=1,
+            bias=False)
+        self.p4_conv = nn.Conv2d(
+            in_channels=self.out_channels,
+            out_channels=self.out_channels // 4,
+            kernel_size=3,
+            padding=1,
+            bias=False)
+        self.p3_conv = nn.Conv2d(
+            in_channels=self.out_channels,
+            out_channels=self.out_channels // 4,
+            kernel_size=3,
+            padding=1,
+            bias=False)
+        self.p2_conv = nn.Conv2d(
+            in_channels=self.out_channels,
+            out_channels=self.out_channels // 4,
+            kernel_size=3,
+            padding=1,
+            bias=False)
+
+        if self.use_asf is True:
+            self.asf = ASFBlock(self.out_channels, self.out_channels // 4)
+
+    def forward(self, x):
+        c2, c3, c4, c5 = x
+
+        in5 = self.in5_conv(c5)
+        in4 = self.in4_conv(c4)
+        in3 = self.in3_conv(c3)
+        in2 = self.in2_conv(c2)
+
+        out4 = in4 + F.interpolate(
+            in5, scale_factor=2, mode="nearest", )  # align_mode=1)  # 1/16
+        out3 = in3 + F.interpolate(
+            out4, scale_factor=2, mode="nearest", )  # align_mode=1)  # 1/8
+        out2 = in2 + F.interpolate(
+            out3, scale_factor=2, mode="nearest", )  # align_mode=1)  # 1/4
+
+        p5 = self.p5_conv(in5)
+        p4 = self.p4_conv(out4)
+        p3 = self.p3_conv(out3)
+        p2 = self.p2_conv(out2)
+        p5 = F.interpolate(p5, scale_factor=8, mode="nearest", )  # align_mode=1)
+        p4 = F.interpolate(p4, scale_factor=4, mode="nearest", )  # align_mode=1)
+        p3 = F.interpolate(p3, scale_factor=2, mode="nearest", )  # align_mode=1)
+
+        fuse = torch.cat([p5, p4, p3, p2], dim=1)
+
+        if self.use_asf is True:
+            fuse = self.asf(fuse, [p5, p4, p3, p2])
+
+        return fuse
+
+
+class RSELayer(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, shortcut=True):
+        super(RSELayer, self).__init__()
+        self.out_channels = out_channels
+        self.in_conv = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=self.out_channels,
+            kernel_size=kernel_size,
+            padding=int(kernel_size // 2),
+            bias=False)
+        self.se_block = SEModule(self.out_channels)
+        self.shortcut = shortcut
+
+    def forward(self, ins):
+        x = self.in_conv(ins)
+        if self.shortcut:
+            out = x + self.se_block(x)
+        else:
+            out = self.se_block(x)
+        return out
+
+
+class RSEFPN(nn.Module):
+    def __init__(self, in_channels, out_channels, shortcut=True, **kwargs):
+        super(RSEFPN, self).__init__()
+        self.out_channels = out_channels
+        self.ins_conv = nn.ModuleList()
+        self.inp_conv = nn.ModuleList()
+        self.intracl = False
+        if 'intracl' in kwargs.keys() and kwargs['intracl'] is True:
+            self.intracl = kwargs['intracl']
+            self.incl1 = IntraCLBlock(self.out_channels // 4, reduce_factor=2)
+            self.incl2 = IntraCLBlock(self.out_channels // 4, reduce_factor=2)
+            self.incl3 = IntraCLBlock(self.out_channels // 4, reduce_factor=2)
+            self.incl4 = IntraCLBlock(self.out_channels // 4, reduce_factor=2)
+
+        for i in range(len(in_channels)):
+            self.ins_conv.append(
+                RSELayer(
+                    in_channels[i],
+                    out_channels,
+                    kernel_size=1,
+                    shortcut=shortcut))
+            self.inp_conv.append(
+                RSELayer(
+                    out_channels,
+                    out_channels // 4,
+                    kernel_size=3,
+                    shortcut=shortcut))
+
+    def forward(self, x):
+        c2, c3, c4, c5 = x
+
+        in5 = self.ins_conv[3](c5)
+        in4 = self.ins_conv[2](c4)
+        in3 = self.ins_conv[1](c3)
+        in2 = self.ins_conv[0](c2)
+
+        out4 = in4 + F.interpolate(
+            in5, scale_factor=2, mode="nearest")  # 1/16
+        out3 = in3 + F.interpolate(
+            out4, scale_factor=2, mode="nearest")  # 1/8
+        out2 = in2 + F.interpolate(
+            out3, scale_factor=2, mode="nearest")  # 1/4
+
+        p5 = self.inp_conv[3](in5)
+        p4 = self.inp_conv[2](out4)
+        p3 = self.inp_conv[1](out3)
+        p2 = self.inp_conv[0](out2)
+
+        if self.intracl is True:
+            p5 = self.incl4(p5)
+            p4 = self.incl3(p4)
+            p3 = self.incl2(p3)
+            p2 = self.incl1(p2)
+
+        p5 = F.interpolate(p5, scale_factor=8, mode="nearest")
+        p4 = F.interpolate(p4, scale_factor=4, mode="nearest")
+        p3 = F.interpolate(p3, scale_factor=2, mode="nearest")
+
+        fuse = torch.cat([p5, p4, p3, p2], dim=1)
+        return fuse
+
+
+class LKPAN(nn.Module):
+    def __init__(self, in_channels, out_channels, mode='large', **kwargs):
+        super(LKPAN, self).__init__()
+        self.out_channels = out_channels
+
+        self.ins_conv = nn.ModuleList()
+        self.inp_conv = nn.ModuleList()
+        # pan head
+        self.pan_head_conv = nn.ModuleList()
+        self.pan_lat_conv = nn.ModuleList()
+
+        if mode.lower() == 'lite':
+            p_layer = DSConv
+        elif mode.lower() == 'large':
+            p_layer = nn.Conv2d
+        else:
+            raise ValueError(
+                "mode can only be one of ['lite', 'large'], but received {}".
+                    format(mode))
+
+        for i in range(len(in_channels)):
+            self.ins_conv.append(
+                nn.Conv2d(
+                    in_channels=in_channels[i],
+                    out_channels=self.out_channels,
+                    kernel_size=1,
+                    bias=False))
+
+            self.inp_conv.append(
+                p_layer(
+                    in_channels=self.out_channels,
+                    out_channels=self.out_channels // 4,
+                    kernel_size=9,
+                    padding=4,
+                    bias=False))
+
+            if i > 0:
+                self.pan_head_conv.append(
+                    nn.Conv2d(
+                        in_channels=self.out_channels // 4,
+                        out_channels=self.out_channels // 4,
+                        kernel_size=3,
+                        padding=1,
+                        stride=2,
+                        bias=False))
+            self.pan_lat_conv.append(
+                p_layer(
+                    in_channels=self.out_channels // 4,
+                    out_channels=self.out_channels // 4,
+                    kernel_size=9,
+                    padding=4,
+                    bias=False))
+            self.intracl = False
+            if 'intracl' in kwargs.keys() and kwargs['intracl'] is True:
+                self.intracl = kwargs['intracl']
+                self.incl1 = IntraCLBlock(self.out_channels // 4, reduce_factor=2)
+                self.incl2 = IntraCLBlock(self.out_channels // 4, reduce_factor=2)
+                self.incl3 = IntraCLBlock(self.out_channels // 4, reduce_factor=2)
+                self.incl4 = IntraCLBlock(self.out_channels // 4, reduce_factor=2)
+
+    def forward(self, x):
+        c2, c3, c4, c5 = x
+
+        in5 = self.ins_conv[3](c5)
+        in4 = self.ins_conv[2](c4)
+        in3 = self.ins_conv[1](c3)
+        in2 = self.ins_conv[0](c2)
+
+        out4 = in4 + F.interpolate(
+            in5, scale_factor=2, mode="nearest")  # 1/16
+        out3 = in3 + F.interpolate(
+            out4, scale_factor=2, mode="nearest")  # 1/8
+        out2 = in2 + F.interpolate(
+            out3, scale_factor=2, mode="nearest")  # 1/4
+
+        f5 = self.inp_conv[3](in5)
+        f4 = self.inp_conv[2](out4)
+        f3 = self.inp_conv[1](out3)
+        f2 = self.inp_conv[0](out2)
+
+        pan3 = f3 + self.pan_head_conv[0](f2)
+        pan4 = f4 + self.pan_head_conv[1](pan3)
+        pan5 = f5 + self.pan_head_conv[2](pan4)
+
+        p2 = self.pan_lat_conv[0](f2)
+        p3 = self.pan_lat_conv[1](pan3)
+        p4 = self.pan_lat_conv[2](pan4)
+        p5 = self.pan_lat_conv[3](pan5)
+
+        if self.intracl is True:
+            p5 = self.incl4(p5)
+            p4 = self.incl3(p4)
+            p3 = self.incl2(p3)
+            p2 = self.incl1(p2)
+
+        p5 = F.interpolate(p5, scale_factor=8, mode="nearest")
+        p4 = F.interpolate(p4, scale_factor=4, mode="nearest")
+        p3 = F.interpolate(p3, scale_factor=2, mode="nearest")
+
+        fuse = torch.cat([p5, p4, p3, p2], dim=1)
+        return fuse
+
+
+class ASFBlock(nn.Module):
+    """
+    This code is refered from:
+        https://github.com/MhLiao/DB/blob/master/decoders/feature_attention.py
+    """
+
+    def __init__(self, in_channels, inter_channels, out_features_num=4):
+        """
+        Adaptive Scale Fusion (ASF) block of DBNet++
+        Args:
+            in_channels: the number of channels in the input data
+            inter_channels: the number of middle channels
+            out_features_num: the number of fused stages
+        """
+        super(ASFBlock, self).__init__()
+        self.in_channels = in_channels
+        self.inter_channels = inter_channels
+        self.out_features_num = out_features_num
+        self.conv = nn.Conv2d(in_channels, inter_channels, 3, padding=1)
+
+        self.spatial_scale = nn.Sequential(
+            # Nx1xHxW
+            nn.Conv2d(
+                in_channels=1,
+                out_channels=1,
+                kernel_size=3,
+                bias=False,
+                padding=1,
+            ),
+            nn.ReLU(),
+            nn.Conv2d(
+                in_channels=1,
+                out_channels=1,
+                kernel_size=1,
+                bias=False,
+            ),
+            nn.Sigmoid())
+
+        self.channel_scale = nn.Sequential(
+            nn.Conv2d(
+                in_channels=inter_channels,
+                out_channels=out_features_num,
+                kernel_size=1,
+                bias=False,
+            ),
+            nn.Sigmoid())
+
+    def forward(self, fuse_features, features_list):
+        fuse_features = self.conv(fuse_features)
+        spatial_x = torch.mean(fuse_features, dim=1, keepdim=True)
+        attention_scores = self.spatial_scale(spatial_x) + fuse_features
+        attention_scores = self.channel_scale(attention_scores)
+        assert len(features_list) == self.out_features_num
+
+        out_list = []
+        for i in range(self.out_features_num):
+            out_list.append(attention_scores[:, i:i + 1] * features_list[i])
+        return torch.cat(out_list, dim=1)

python/rapidocr_torch/modeling/necks/intracl.py

@@ -0,0 +1,114 @@
+from torch import nn
+
+
+class IntraCLBlock(nn.Module):
+    def __init__(self, in_channels=96, reduce_factor=4):
+        super(IntraCLBlock, self).__init__()
+        self.channels = in_channels
+        self.rf = reduce_factor
+        self.conv1x1_reduce_channel = nn.Conv2d(
+            self.channels,
+            self.channels // self.rf,
+            kernel_size=1,
+            stride=1,
+            padding=0)
+        self.conv1x1_return_channel = nn.Conv2d(
+            self.channels // self.rf,
+            self.channels,
+            kernel_size=1,
+            stride=1,
+            padding=0)
+
+        self.v_layer_7x1 = nn.Conv2d(
+            self.channels // self.rf,
+            self.channels // self.rf,
+            kernel_size=(7, 1),
+            stride=(1, 1),
+            padding=(3, 0))
+        self.v_layer_5x1 = nn.Conv2d(
+            self.channels // self.rf,
+            self.channels // self.rf,
+            kernel_size=(5, 1),
+            stride=(1, 1),
+            padding=(2, 0))
+        self.v_layer_3x1 = nn.Conv2d(
+            self.channels // self.rf,
+            self.channels // self.rf,
+            kernel_size=(3, 1),
+            stride=(1, 1),
+            padding=(1, 0))
+
+        self.q_layer_1x7 = nn.Conv2d(
+            self.channels // self.rf,
+            self.channels // self.rf,
+            kernel_size=(1, 7),
+            stride=(1, 1),
+            padding=(0, 3))
+        self.q_layer_1x5 = nn.Conv2d(
+            self.channels // self.rf,
+            self.channels // self.rf,
+            kernel_size=(1, 5),
+            stride=(1, 1),
+            padding=(0, 2))
+        self.q_layer_1x3 = nn.Conv2d(
+            self.channels // self.rf,
+            self.channels // self.rf,
+            kernel_size=(1, 3),
+            stride=(1, 1),
+            padding=(0, 1))
+
+        # base
+        self.c_layer_7x7 = nn.Conv2d(
+            self.channels // self.rf,
+            self.channels // self.rf,
+            kernel_size=(7, 7),
+            stride=(1, 1),
+            padding=(3, 3))
+        self.c_layer_5x5 = nn.Conv2d(
+            self.channels // self.rf,
+            self.channels // self.rf,
+            kernel_size=(5, 5),
+            stride=(1, 1),
+            padding=(2, 2))
+        self.c_layer_3x3 = nn.Conv2d(
+            self.channels // self.rf,
+            self.channels // self.rf,
+            kernel_size=(3, 3),
+            stride=(1, 1),
+            padding=(1, 1))
+
+        self.bn = nn.BatchNorm2d(self.channels)
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        x_new = self.conv1x1_reduce_channel(x)
+
+        x_7_c = self.c_layer_7x7(x_new)
+        x_7_v = self.v_layer_7x1(x_new)
+        x_7_q = self.q_layer_1x7(x_new)
+        x_7 = x_7_c + x_7_v + x_7_q
+
+        x_5_c = self.c_layer_5x5(x_7)
+        x_5_v = self.v_layer_5x1(x_7)
+        x_5_q = self.q_layer_1x5(x_7)
+        x_5 = x_5_c + x_5_v + x_5_q
+
+        x_3_c = self.c_layer_3x3(x_5)
+        x_3_v = self.v_layer_3x1(x_5)
+        x_3_q = self.q_layer_1x3(x_5)
+        x_3 = x_3_c + x_3_v + x_3_q
+
+        x_relation = self.conv1x1_return_channel(x_3)
+
+        x_relation = self.bn(x_relation)
+        x_relation = self.relu(x_relation)
+
+        return x + x_relation
+
+
+def build_intraclblock_list(num_block):
+    IntraCLBlock_list = nn.ModuleList()
+    for i in range(num_block):
+        IntraCLBlock_list.append(IntraCLBlock())
+
+    return IntraCLBlock_list

python/rapidocr_torch/modeling/necks/rnn.py

@@ -0,0 +1,205 @@
+import os, sys
+import torch
+import torch.nn as nn
+from rapidocr_torch.modeling.backbones.rec_svtrnet import Block, ConvBNLayer
+
+class Im2Seq(nn.Module):
+    def __init__(self, in_channels, **kwargs):
+        super().__init__()
+        self.out_channels = in_channels
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        # assert H == 1
+        x = x.squeeze(dim=2)
+        # x = x.transpose([0, 2, 1])  # paddle (NTC)(batch, width, channels)
+        x = x.permute(0,2,1)
+        return x
+
+
+class EncoderWithRNN_(nn.Module):
+    def __init__(self, in_channels, hidden_size):
+        super(EncoderWithRNN_, self).__init__()
+        self.out_channels = hidden_size * 2
+        self.rnn1 = nn.LSTM(in_channels, hidden_size, bidirectional=False, batch_first=True, num_layers=2)
+        self.rnn2 = nn.LSTM(in_channels, hidden_size, bidirectional=False, batch_first=True, num_layers=2)
+
+    def forward(self, x):
+        self.rnn1.flatten_parameters()
+        self.rnn2.flatten_parameters()
+        out1, h1 = self.rnn1(x)
+        out2, h2 = self.rnn2(torch.flip(x, [1]))
+        return torch.cat([out1, torch.flip(out2, [1])], 2)
+
+
+class EncoderWithRNN(nn.Module):
+    def __init__(self, in_channels, hidden_size):
+        super(EncoderWithRNN, self).__init__()
+        self.out_channels = hidden_size * 2
+        self.lstm = nn.LSTM(
+            in_channels, hidden_size, num_layers=2, batch_first=True, bidirectional=True) # batch_first:=True
+
+    def forward(self, x):
+        x, _ = self.lstm(x)
+        return x
+
+
+class EncoderWithFC(nn.Module):
+    def __init__(self, in_channels, hidden_size):
+        super(EncoderWithFC, self).__init__()
+        self.out_channels = hidden_size
+        self.fc = nn.Linear(
+            in_channels,
+            hidden_size,
+            bias=True,
+            )
+
+    def forward(self, x):
+        x = self.fc(x)
+        return x
+
+
+class EncoderWithSVTR(nn.Module):
+    def __init__(
+            self,
+            in_channels,
+            dims=64,  # XS
+            depth=2,
+            hidden_dims=120,
+            use_guide=False,
+            num_heads=8,
+            qkv_bias=True,
+            mlp_ratio=2.0,
+            drop_rate=0.1,
+            kernel_size=[3,3],
+            attn_drop_rate=0.1,
+            drop_path=0.,
+            qk_scale=None):
+        super(EncoderWithSVTR, self).__init__()
+        self.depth = depth
+        self.use_guide = use_guide
+        self.conv1 = ConvBNLayer(
+            in_channels,
+            in_channels // 8,
+            kernel_size=kernel_size,
+            padding=[kernel_size[0] // 2, kernel_size[1] // 2],
+            act='swish')
+        self.conv2 = ConvBNLayer(
+            in_channels // 8, hidden_dims, kernel_size=1, act='swish')
+
+        self.svtr_block = nn.ModuleList([
+            Block(
+                dim=hidden_dims,
+                num_heads=num_heads,
+                mixer='Global',
+                HW=None,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop_rate,
+                act_layer='swish',
+                attn_drop=attn_drop_rate,
+                drop_path=drop_path,
+                norm_layer='nn.LayerNorm',
+                epsilon=1e-05,
+                prenorm=False) for i in range(depth)
+        ])
+        self.norm = nn.LayerNorm(hidden_dims, eps=1e-6)
+        self.conv3 = ConvBNLayer(
+            hidden_dims, in_channels, kernel_size=1, act='swish')
+        # last conv-nxn, the input is concat of input tensor and conv3 output tensor
+        self.conv4 = ConvBNLayer(
+            2 * in_channels, in_channels // 8, padding=1, act='swish')
+
+        self.conv1x1 = ConvBNLayer(
+            in_channels // 8, dims, kernel_size=1, act='swish')
+        self.out_channels = dims
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        # weight initialization
+        if isinstance(m, nn.Conv2d):
+            nn.init.kaiming_normal_(m.weight, mode='fan_out')
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.BatchNorm2d):
+            nn.init.ones_(m.weight)
+            nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.Linear):
+            nn.init.normal_(m.weight, 0, 0.01)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.ConvTranspose2d):
+            nn.init.kaiming_normal_(m.weight, mode='fan_out')
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.ones_(m.weight)
+            nn.init.zeros_(m.bias)
+
+    def forward(self, x):
+        # for use guide
+        if self.use_guide:
+            z = x.clone()
+            z.stop_gradient = True
+        else:
+            z = x
+        # for short cut
+        h = z
+        # reduce dim
+        z = self.conv1(z)
+        z = self.conv2(z)
+        # SVTR global block
+        B, C, H, W = z.shape
+        z = z.flatten(2).permute(0, 2, 1)
+
+        for blk in self.svtr_block:
+            z = blk(z)
+
+        z = self.norm(z)
+        # last stage
+        z = z.reshape([-1, H, W, C]).permute(0, 3, 1, 2)
+        z = self.conv3(z)
+        z = torch.cat((h, z), dim=1)
+        z = self.conv1x1(self.conv4(z))
+
+        return z
+
+
+class SequenceEncoder(nn.Module):
+    def __init__(self, in_channels, encoder_type, hidden_size=48, **kwargs):
+        super(SequenceEncoder, self).__init__()
+        self.encoder_reshape = Im2Seq(in_channels)
+        self.out_channels = self.encoder_reshape.out_channels
+        self.encoder_type = encoder_type
+        if encoder_type == 'reshape':
+            self.only_reshape = True
+        else:
+            support_encoder_dict = {
+                'reshape': Im2Seq,
+                'fc': EncoderWithFC,
+                'rnn': EncoderWithRNN,
+                'svtr': EncoderWithSVTR,
+            }
+            assert encoder_type in support_encoder_dict, '{} must in {}'.format(
+                encoder_type, support_encoder_dict.keys())
+
+            if encoder_type == "svtr":
+                self.encoder = support_encoder_dict[encoder_type](
+                    self.encoder_reshape.out_channels, **kwargs)
+            else:
+                self.encoder = support_encoder_dict[encoder_type](
+                    self.encoder_reshape.out_channels, hidden_size)
+            self.out_channels = self.encoder.out_channels
+            self.only_reshape = False
+
+    def forward(self, x):
+        if self.encoder_type != 'svtr':
+            x = self.encoder_reshape(x)
+            if not self.only_reshape:
+                x = self.encoder(x)
+            return x
+        else:
+            x = self.encoder(x)
+            x = self.encoder_reshape(x)
+            return x

python/rapidocr_torch/utils/__init__.py

@@ -0,0 +1,20 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+from pathlib import Path
+from typing import Dict, Union
+
+import yaml
+
+from .infer_engine import TorchInferSession
+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
+
+
+def read_yaml(yaml_path: Union[str, Path]) -> Dict[str, Dict]:
+    with open(yaml_path, "rb") as f:
+        data = yaml.load(f, Loader=yaml.Loader)
+    return data

python/rapidocr_torch/utils/infer_engine.py

@@ -0,0 +1,63 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+import os
+import platform
+from pathlib import Path
+from typing import Optional, Union, Dict
+
+import numpy as np
+import torch
+import yaml
+
+root_dir = Path(__file__).resolve().parent.parent
+DEFAULT_CFG_PATH = root_dir / "arch_config.yaml"
+
+def read_yaml(yaml_path: Union[str, Path]) -> Dict[str, Dict]:
+    with open(yaml_path, "rb") as f:
+        data = yaml.load(f, Loader=yaml.Loader)
+    return data
+
+from .logger import get_logger
+from rapidocr_torch.modeling.architectures.base_model import BaseModel
+
+
+class TorchInferSession:
+    def __init__(self, config, mode: Optional[str] = None) -> None:
+
+        all_arch_config = read_yaml(DEFAULT_CFG_PATH)
+
+        self.logger = get_logger("TorchInferSession")
+        self.mode = mode
+        model_path = Path(config["model_path"])
+        self._verify_model(model_path)
+        file_name = model_path.stem
+        if file_name not in all_arch_config:
+            raise ValueError(f"architecture {file_name} is not in config.yaml")
+        arch_config = all_arch_config[file_name]
+        self.predictor = BaseModel(arch_config)
+        self.predictor.load_state_dict(torch.load(model_path, weights_only=True))
+        self.predictor.eval()
+        self.use_gpu = False
+        if config["use_cuda"]:
+            self.predictor.cuda()
+            self.use_gpu = True
+    def __call__(self, img: np.ndarray):
+        with torch.no_grad():
+            inp = torch.from_numpy(img)
+            if self.use_gpu:
+                inp = inp.cuda()
+            # 适配跟onnx对齐取值逻辑
+            outputs = self.predictor(inp).unsqueeze(0)
+            return outputs.cpu().numpy()
+    @staticmethod
+    def _verify_model(model_path):
+        model_path = Path(model_path)
+        if not model_path.exists():
+            raise FileNotFoundError(f"{model_path} does not exists.")
+        if not model_path.is_file():
+            raise FileExistsError(f"{model_path} is not a file.")
+
+
+class TorchInferError(Exception):
+    pass

python/rapidocr_torch/utils/load_image.py

@@ -0,0 +1,123 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+from io import BytesIO
+from pathlib import Path
+from typing import Any, Union
+
+import cv2
+import numpy as np
+from PIL import Image, UnidentifiedImageError
+
+root_dir = Path(__file__).resolve().parent
+InputType = Union[str, np.ndarray, bytes, Path, Image.Image]
+
+
+class LoadImage:
+    def __init__(self):
+        pass
+
+    def __call__(self, img: InputType) -> np.ndarray:
+        if not isinstance(img, InputType.__args__):
+            raise LoadImageError(
+                f"The img type {type(img)} does not in {InputType.__args__}"
+            )
+
+        origin_img_type = type(img)
+        img = self.load_img(img)
+        img = self.convert_img(img, origin_img_type)
+        return img
+
+    def load_img(self, img: InputType) -> np.ndarray:
+        if isinstance(img, (str, Path)):
+            self.verify_exist(img)
+            try:
+                img = self.img_to_ndarray(Image.open(img))
+            except UnidentifiedImageError as e:
+                raise LoadImageError(f"cannot identify image file {img}") from e
+            return img
+
+        if isinstance(img, bytes):
+            img = self.img_to_ndarray(Image.open(BytesIO(img)))
+            return img
+
+        if isinstance(img, np.ndarray):
+            return img
+
+        if isinstance(img, Image.Image):
+            return self.img_to_ndarray(img)
+
+        raise LoadImageError(f"{type(img)} is not supported!")
+
+    def img_to_ndarray(self, img: Image.Image) -> np.ndarray:
+        if img.mode == "1":
+            img = img.convert("L")
+            return np.array(img)
+        return np.array(img)
+
+    def convert_img(self, img: np.ndarray, origin_img_type: Any) -> np.ndarray:
+        if img.ndim == 2:
+            return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+
+        if img.ndim == 3:
+            channel = img.shape[2]
+            if channel == 1:
+                return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+
+            if channel == 2:
+                return self.cvt_two_to_three(img)
+
+            if channel == 3:
+                if issubclass(origin_img_type, (str, Path, bytes, Image.Image)):
+                    return cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+                return img
+
+            if channel == 4:
+                return self.cvt_four_to_three(img)
+
+            raise LoadImageError(
+                f"The channel({channel}) of the img is not in [1, 2, 3, 4]"
+            )
+
+        raise LoadImageError(f"The ndim({img.ndim}) of the img is not in [2, 3]")
+
+    @staticmethod
+    def cvt_two_to_three(img: np.ndarray) -> np.ndarray:
+        """gray + alpha → BGR"""
+        img_gray = img[..., 0]
+        img_bgr = cv2.cvtColor(img_gray, cv2.COLOR_GRAY2BGR)
+
+        img_alpha = img[..., 1]
+        not_a = cv2.bitwise_not(img_alpha)
+        not_a = cv2.cvtColor(not_a, cv2.COLOR_GRAY2BGR)
+
+        new_img = cv2.bitwise_and(img_bgr, img_bgr, mask=img_alpha)
+        new_img = cv2.add(new_img, not_a)
+        return new_img
+
+    @staticmethod
+    def cvt_four_to_three(img: np.ndarray) -> np.ndarray:
+        """RGBA → BGR"""
+        r, g, b, a = cv2.split(img)
+        new_img = cv2.merge((b, g, r))
+
+        not_a = cv2.bitwise_not(a)
+        not_a = cv2.cvtColor(not_a, cv2.COLOR_GRAY2BGR)
+
+        new_img = cv2.bitwise_and(new_img, new_img, mask=a)
+
+        mean_color = np.mean(new_img)
+        if mean_color <= 0.0:
+            new_img = cv2.add(new_img, not_a)
+        else:
+            new_img = cv2.bitwise_not(new_img)
+        return new_img
+
+    @staticmethod
+    def verify_exist(file_path: Union[str, Path]):
+        if not Path(file_path).exists():
+            raise LoadImageError(f"{file_path} does not exist.")
+
+
+class LoadImageError(Exception):
+    pass

python/rapidocr_torch/utils/logger.py

@@ -0,0 +1,21 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+import logging
+from functools import lru_cache
+
+
+@lru_cache(maxsize=32)
+def get_logger(name: str) -> logging.Logger:
+    logger = logging.getLogger(name)
+    logger.setLevel(logging.DEBUG)
+
+    fmt = "%(asctime)s - %(name)s - %(levelname)s: %(message)s"
+    format_str = logging.Formatter(fmt)
+
+    sh = logging.StreamHandler()
+    sh.setLevel(logging.DEBUG)
+
+    logger.addHandler(sh)
+    sh.setFormatter(format_str)
+    return logger

python/rapidocr_torch/utils/parse_parameters.py

@@ -0,0 +1,196 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+import argparse
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Union
+
+import numpy as np
+from PIL import Image
+
+root_dir = Path(__file__).resolve().parent.parent
+InputType = Union[str, np.ndarray, bytes, Path, Image.Image]
+
+
+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])
+    config["Cls"][key] = str(root_dir / config["Cls"][key])
+    return config
+
+
+def init_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-img", "--img_path", type=str, default=None, required=True)
+    parser.add_argument("-p", "--print_cost", action="store_true", default=False)
+
+    global_group = parser.add_argument_group(title="Global")
+    global_group.add_argument("--text_score", type=float, default=0.5)
+
+    global_group.add_argument("--no_det", action="store_true", default=False)
+    global_group.add_argument("--no_cls", action="store_true", default=False)
+    global_group.add_argument("--no_rec", action="store_true", default=False)
+
+    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("--return_word_box", action="store_true", default=False)
+
+    global_group.add_argument("--intra_op_num_threads", type=int, default=-1)
+    global_group.add_argument("--inter_op_num_threads", type=int, default=-1)
+
+    det_group = parser.add_argument_group(title="Det")
+    det_group.add_argument("--det_use_cuda", action="store_true", default=False)
+    det_group.add_argument("--det_use_dml", action="store_true", default=False)
+    det_group.add_argument("--det_model_path", type=str, default=None)
+    det_group.add_argument("--det_limit_side_len", type=float, default=736)
+    det_group.add_argument(
+        "--det_limit_type", type=str, default="min", choices=["max", "min"]
+    )
+    det_group.add_argument("--det_thresh", type=float, default=0.3)
+    det_group.add_argument("--det_box_thresh", type=float, default=0.5)
+    det_group.add_argument("--det_unclip_ratio", type=float, default=1.6)
+    det_group.add_argument(
+        "--det_donot_use_dilation", action="store_true", default=False
+    )
+    det_group.add_argument(
+        "--det_score_mode", type=str, default="fast", choices=["slow", "fast"]
+    )
+
+    cls_group = parser.add_argument_group(title="Cls")
+    cls_group.add_argument("--cls_use_cuda", action="store_true", default=False)
+    cls_group.add_argument("--cls_use_dml", action="store_true", default=False)
+    cls_group.add_argument("--cls_model_path", type=str, default=None)
+    cls_group.add_argument("--cls_image_shape", type=list, default=[3, 48, 192])
+    cls_group.add_argument("--cls_label_list", type=list, default=["0", "180"])
+    cls_group.add_argument("--cls_batch_num", type=int, default=6)
+    cls_group.add_argument("--cls_thresh", type=float, default=0.9)
+
+    rec_group = parser.add_argument_group(title="Rec")
+    rec_group.add_argument("--rec_use_cuda", action="store_true", default=False)
+    rec_group.add_argument("--rec_use_dml", action="store_true", default=False)
+    rec_group.add_argument("--rec_model_path", type=str, default=None)
+    rec_group.add_argument("--rec_keys_path", type=str, default=None)
+    rec_group.add_argument("--rec_img_shape", type=list, default=[3, 48, 320])
+    rec_group.add_argument("--rec_batch_num", type=int, default=6)
+
+    vis_group = parser.add_argument_group(title="Visual Result")
+    vis_group.add_argument("-vis", "--vis_res", action="store_true", default=False)
+    vis_group.add_argument(
+        "--vis_font_path",
+        type=str,
+        default=None,
+        help="When -vis is True, the font_path must have value.",
+    )
+    vis_group.add_argument(
+        "--vis_save_path",
+        type=str,
+        default=".",
+        help="The directory of saving the vis image.",
+    )
+
+    args = parser.parse_args()
+    return args
+
+
+class UpdateParameters:
+    def __init__(self) -> None:
+        pass
+
+    def parse_kwargs(self, **kwargs):
+        global_dict, det_dict, cls_dict, rec_dict = {}, {}, {}, {}
+        for k, v in kwargs.items():
+            if k.startswith("det"):
+                k = k.split("det_")[1]
+                if k == "donot_use_dilation":
+                    k = "use_dilation"
+                    v = not v
+
+                det_dict[k] = v
+            elif k.startswith("cls"):
+                cls_dict[k] = v
+            elif k.startswith("rec"):
+                rec_dict[k] = v
+            else:
+                global_dict[k] = v
+        return global_dict, det_dict, cls_dict, rec_dict
+
+    def __call__(self, config, **kwargs):
+        global_dict, det_dict, cls_dict, rec_dict = self.parse_kwargs(**kwargs)
+        new_config = {
+            "Global": self.update_global_params(config["Global"], global_dict),
+            "Det": self.update_params(
+                config["Det"],
+                det_dict,
+                "det_",
+                ["det_model_path", "det_use_cuda", "det_use_dml"],
+            ),
+            "Cls": self.update_params(
+                config["Cls"],
+                cls_dict,
+                "cls_",
+                ["cls_label_list", "cls_model_path", "cls_use_cuda", "cls_use_dml"],
+            ),
+            "Rec": self.update_params(
+                config["Rec"],
+                rec_dict,
+                "rec_",
+                ["rec_model_path", "rec_use_cuda", "rec_use_dml"],
+            ),
+        }
+
+        update_params = ["intra_op_num_threads", "inter_op_num_threads"]
+        new_config = self.update_global_to_module(
+            config, update_params, src="Global", dsts=["Det", "Cls", "Rec"]
+        )
+        return new_config
+
+    def update_global_to_module(
+        self, config, params: List[str], src: str, dsts: List[str]
+    ):
+        for dst in dsts:
+            for param in params:
+                config[dst].update({param: config[src][param]})
+        return config
+
+    def update_global_params(self, config, global_dict):
+        if global_dict:
+            config.update(global_dict)
+        return config
+
+    def update_params(
+        self,
+        config,
+        param_dict: Dict[str, str],
+        prefix: str,
+        need_remove_prefix: Optional[List[str]] = None,
+    ):
+        if not param_dict:
+            return config
+
+        filter_dict = self.remove_prefix(param_dict, prefix, need_remove_prefix)
+        model_path = filter_dict.get("model_path", None)
+        if not model_path:
+            filter_dict["model_path"] = str(root_dir / config["model_path"])
+
+        config.update(filter_dict)
+        return config
+
+    @staticmethod
+    def remove_prefix(
+        config: Dict[str, str],
+        prefix: str,
+        need_remove_prefix: Optional[List[str]] = None,
+    ) -> Dict[str, str]:
+        if not need_remove_prefix:
+            return config
+
+        new_rec_dict = {}
+        for k, v in config.items():
+            if k in need_remove_prefix:
+                k = k.split(prefix)[1]
+            new_rec_dict[k] = v
+        return new_rec_dict

python/rapidocr_torch/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/rapidocr_torch/utils/vis_res.py

@@ -0,0 +1,143 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+import math
+import random
+from pathlib import Path
+from typing import List, Optional, Tuple, Union
+
+import cv2
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+
+from .load_image import LoadImage
+
+root_dir = Path(__file__).resolve().parent
+InputType = Union[str, np.ndarray, bytes, Path, Image.Image]
+
+
+class VisRes:
+    def __init__(self, text_score: float = 0.5):
+        self.text_score = text_score
+        self.load_img = LoadImage()
+
+    def __call__(
+        self,
+        img_content: InputType,
+        dt_boxes: np.ndarray,
+        txts: Optional[Union[List[str], Tuple[str]]] = None,
+        scores: Optional[Tuple[float]] = None,
+        font_path: Optional[str] = None,
+    ) -> np.ndarray:
+        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)
+
+    def draw_dt_boxes(self, img_content: InputType, dt_boxes: np.ndarray) -> np.ndarray:
+        img = self.load_img(img_content)
+
+        for idx, box in enumerate(dt_boxes):
+            color = self.get_random_color()
+
+            points = np.array(box)
+            cv2.polylines(img, np.int32([points]), 1, color=color, thickness=1)
+
+            start_point = round(points[0][0]), round(points[0][1])
+            cv2.putText(
+                img, f"{idx}", start_point, cv2.FONT_HERSHEY_SIMPLEX, 1, color, 3
+            )
+        return img
+
+    def draw_ocr_box_txt(
+        self,
+        img_content: InputType,
+        dt_boxes: np.ndarray,
+        txts: Union[List[str], Tuple[str]],
+        scores: Optional[Tuple[float]] = None,
+        font_path: Optional[str] = None,
+    ) -> np.ndarray:
+        font_path = self.get_font_path(font_path)
+
+        image = Image.fromarray(self.load_img(img_content))
+        h, w = image.height, image.width
+        if image.mode == "L":
+            image = image.convert("RGB")
+
+        img_left = image.copy()
+        img_right = Image.new("RGB", (w, h), (255, 255, 255))
+
+        random.seed(0)
+        draw_left = ImageDraw.Draw(img_left)
+        draw_right = ImageDraw.Draw(img_right)
+        for idx, (box, txt) in enumerate(zip(dt_boxes, txts)):
+            if scores is not None and float(scores[idx]) < self.text_score:
+                continue
+
+            color = self.get_random_color()
+
+            box_list = np.array(box).reshape(8).tolist()
+            draw_left.polygon(box_list, fill=color)
+            draw_right.polygon(box_list, outline=color)
+
+            box_height = self.get_box_height(box)
+            box_width = self.get_box_width(box)
+            if box_height > 2 * box_width:
+                font_size = max(int(box_width * 0.9), 10)
+                font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
+                cur_y = box[0][1]
+
+                for c in txt:
+                    draw_right.text(
+                        (box[0][0] + 3, cur_y), c, fill=(0, 0, 0), font=font
+                    )
+                    cur_y += self.get_char_size(font, c)
+            else:
+                font_size = max(int(box_height * 0.8), 10)
+                font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
+                draw_right.text([box[0][0], box[0][1]], txt, fill=(0, 0, 0), font=font)
+
+        img_left = Image.blend(image, img_left, 0.5)
+        img_show = Image.new("RGB", (w * 2, h), (255, 255, 255))
+        img_show.paste(img_left, (0, 0, w, h))
+        img_show.paste(img_right, (w, 0, w * 2, h))
+        return np.array(img_show)
+
+    @staticmethod
+    def get_font_path(font_path: Optional[Union[str, Path]] = None) -> str:
+        if font_path is None or not Path(font_path).exists():
+            raise FileNotFoundError(
+                f"The {font_path} does not exists! \n"
+                f"You could download the file in the https://drive.google.com/file/d/1evWVX38EFNwTq_n5gTFgnlv8tdaNcyIA/view?usp=sharing"
+            )
+        return str(font_path)
+
+    @staticmethod
+    def get_random_color() -> Tuple[int, int, int]:
+        return (
+            random.randint(0, 255),
+            random.randint(0, 255),
+            random.randint(0, 255),
+        )
+
+    @staticmethod
+    def get_box_height(box: List[List[float]]) -> float:
+        return math.sqrt((box[0][0] - box[3][0]) ** 2 + (box[0][1] - box[3][1]) ** 2)
+
+    @staticmethod
+    def get_box_width(box: List[List[float]]) -> float:
+        return math.sqrt((box[0][0] - box[1][0]) ** 2 + (box[0][1] - box[1][1]) ** 2)
+
+    @staticmethod
+    def get_char_size(font, char_str: str) -> float:
+        # compatible with Pillow v9 and v10.
+        if hasattr(font, "getsize"):
+            get_size_func = getattr(font, "getsize")
+            return get_size_func(char_str)[1]
+
+        if hasattr(font, "getlength"):
+            get_size_func = getattr(font, "getlength")
+            return get_size_func(char_str)
+
+        raise ValueError(
+            "The Pillow ImageFont instance has not getsize or getlength func."
+        )

python/requirements_torch.txt

@@ -0,0 +1,10 @@
+pyclipper>=1.2.0
+opencv_python>=4.5.1.48
+numpy>=1.19.5,<3.0.0
+six>=1.15.0
+Shapely>=1.7.1,!=2.0.4  # python3.12 2.0.4 bug
+PyYAML
+Pillow
+tqdm
+torch
+torchvision

python/setup_torch.py

@@ -0,0 +1,77 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+import sys
+from pathlib import Path
+from typing import List, Union
+
+import setuptools
+from get_pypi_latest_version import GetPyPiLatestVersion
+
+
+def read_txt(txt_path: Union[Path, str]) -> List[str]:
+    with open(txt_path, "r", encoding="utf-8") as f:
+        data = [v.rstrip("\n") for v in f]
+    return data
+
+
+def get_readme():
+    root_dir = Path(__file__).resolve().parent.parent
+    readme_path = str(root_dir / "docs" / "doc_whl_rapidocr_ort.md")
+    print(readme_path)
+    with open(readme_path, "r", encoding="utf-8") as f:
+        readme = f.read()
+    return readme
+
+
+MODULE_NAME = "rapidocr_torch"
+
+obtainer = GetPyPiLatestVersion()
+try:
+    latest_version = obtainer(MODULE_NAME)
+except Exception:
+    latest_version = "0.0.0"
+VERSION_NUM = obtainer.version_add_one(latest_version, add_patch=True)
+
+if len(sys.argv) > 2:
+    match_str = " ".join(sys.argv[2:])
+    matched_versions = obtainer.extract_version(match_str)
+    if matched_versions:
+        VERSION_NUM = matched_versions
+sys.argv = sys.argv[:2]
+
+project_urls = {
+    "Documentation": "https://rapidai.github.io/RapidOCRDocs",
+    "Changelog": "https://github.com/RapidAI/RapidOCR/releases",
+}
+
+setuptools.setup(
+    name=MODULE_NAME,
+    version=VERSION_NUM,
+    platforms="Any",
+    description="A cross platform OCR Library based on pytorch.",
+    long_description=get_readme(),
+    long_description_content_type="text/markdown",
+    author="SWHL",
+    author_email="liekkaskono@163.com",
+    url="https://github.com/RapidAI/RapidOCR",
+    project_urls=project_urls,
+    license="Apache-2.0",
+    include_package_data=True,
+    install_requires=read_txt("requirements_torch.txt"),
+    package_dir={"": MODULE_NAME},
+    packages=setuptools.find_namespace_packages(where=MODULE_NAME),
+    package_data={"": ["*.pth", "*.yaml", "*.txt"]},
+    keywords=[
+        "ocr,text_detection,text_recognition,db,onnxruntime,pytorch,paddleocr,openvino,rapidocr"
+    ],
+    classifiers=[
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
+    ],
+    python_requires=">=3.10,<3.13",
+    entry_points={
+        "console_scripts": [f"{MODULE_NAME}={MODULE_NAME}.main:main"],
+    },
+)

python/tests/test_torch.py

@@ -0,0 +1,224 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+import logging
+import sys
+from pathlib import Path
+from typing import List
+
+import cv2
+import numpy as np
+import pytest
+
+root_dir = Path(__file__).resolve().parent.parent
+sys.path.append(str(root_dir))
+
+from rapidocr_torch import LoadImageError, RapidOCR
+
+from .base_module import download_file
+
+engine = RapidOCR()
+tests_dir = root_dir / "tests" / "test_files"
+img_path = tests_dir / "ch_en_num.jpg"
+package_name = "rapidocr_torch"
+
+
+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) >= 53
+
+    img_path.unlink()
+
+
+
+def test_mode_one_img():
+    img_path = tests_dir / "issue_170.png"
+    result, _ = engine(img_path)
+    assert result[0][1] == "TEST"
+
+
+@pytest.mark.parametrize(
+    "img_name,gt",
+    [
+        (
+            "black_font_color_transparent.png",
+            "我是中国人",
+        ),
+        (
+            "white_font_color_transparent.png",
+            "我是中国人",
+        ),
+    ],
+)
+def test_transparent_img(img_name: str, gt: str):
+    img_path = tests_dir / img_name
+    result, _ = engine(str(img_path))
+    assert result[0][1] == gt
+
+
+@pytest.mark.parametrize(
+    "img_name,gt_len,gt_first_len",
+    [
+        ("test_without_det.jpg", 1, "在中国作家协会第三届儿童文学"),
+    ],
+)
+def test_letterbox_like(img_name, gt_len, gt_first_len):
+    img_path = tests_dir / img_name
+    result, _ = engine(str(img_path))
+
+    assert len(result) == gt_len
+    assert result[0][1].lower() == gt_first_len.lower()
+
+
+def test_only_det():
+    result, _ = engine(str(img_path), use_det=True, use_cls=False, use_rec=False)
+
+    assert len(result) == 18
+
+
+def test_only_cls():
+    img_path = tests_dir / "text_cls.jpg"
+    result, _ = engine(img_path, use_det=False, use_cls=True, use_rec=False)
+    assert len(result) == 1
+    assert result[0][0] == "0"
+
+
+def test_only_rec():
+    img_path = tests_dir / "text_rec.jpg"
+    result, _ = engine(img_path, use_det=False, use_cls=False, use_rec=True)
+    assert len(result) == 1
+    assert result[0][0] == "韩国小馆"
+
+
+def test_det_rec():
+    result, _ = engine(img_path, use_det=True, use_cls=False, use_rec=True)
+    assert result[0][1] == "正品促销"
+    assert len(result) == 18
+
+
+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] == "韩国小馆"
+
+
+def test_det_cls_rec():
+    img = cv2.imread(img_path)
+
+    result, _ = engine(img)
+    assert result is not None
+    assert result[0][1] == "正品促销"
+    assert len(result) == 18
+
+
+def test_empty():
+    img = None
+    with pytest.raises(LoadImageError) as exc_info:
+        engine(img)
+        raise LoadImageError
+    assert exc_info.type is LoadImageError
+
+
+def test_zeros():
+    img = np.zeros([640, 640, 3], np.uint8)
+    result, _ = engine(img)
+    assert result is None
+
+
+def test_input_str():
+    result, _ = engine(img_path)
+    assert result[0][1] == "正品促销"
+    assert len(result) == 18
+
+
+def test_input_bytes():
+    with open(img_path, "rb") as f:
+        result, _ = engine(f.read())
+    assert result[0][1] == "正品促销"
+    assert len(result) == 18
+
+
+def test_input_path():
+    result, _ = engine(img_path)
+    assert result[0][1] == "正品促销"
+    assert len(result) == 18
+
+
+def test_input_parameters():
+    img_path = tests_dir / "ch_en_num.jpg"
+    engine = RapidOCR(text_score=1)
+    result, _ = engine(img_path)
+
+    assert result is None
+
+
+def test_input_det_parameters():
+    with pytest.raises(FileNotFoundError) as exc_info:
+        engine = RapidOCR(det_model_path="1.pth")
+        result, _ = engine(img_path)
+        raise FileNotFoundError()
+    assert exc_info.type is FileNotFoundError
+
+
+def test_input_cls_parameters():
+    with pytest.raises(FileNotFoundError) as exc_info:
+        engine = RapidOCR(cls_model_path="1.pth")
+        result, _ = engine(img_path)
+        raise FileNotFoundError()
+    assert exc_info.type is FileNotFoundError
+
+
+def test_input_rec_parameters():
+    with pytest.raises(FileNotFoundError) as exc_info:
+        engine = RapidOCR(rec_model_path="1.pth")
+        result, _ = engine(img_path)
+        raise FileNotFoundError()
+    assert exc_info.type is FileNotFoundError
+
+
+def test_input_three_ndim_two_channel():
+    img_npy = tests_dir / "two_dim_image.npy"
+    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"
+
+
+def test_input_three_ndim_one_channel():
+    img = cv2.imread(img_path)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    img = img[:, :, 0]
+    img = img[..., None]  # (H, W, 1)
+
+    result, _ = engine(img)
+
+    assert len(result) >= 17
+
+
+@pytest.mark.parametrize(
+    "img_name,words",
+    [
+        (
+            "black_font_color_transparent.png",
+            ["我", "是", "中", "国", "人"],
+        ),
+        (
+            "issue_170.png",
+            ["T", "E", "S", "T"],
+        ),
+    ],
+)
+def test_word_ocr(img_name: str, words: List[str]):
+    img_path = tests_dir / img_name
+    result, _ = engine(str(img_path), return_word_box=True)
+    assert result[0][4] == words