chore: update files

python/rapidocr/cal_rec_boxes/main.py

@@ -8,7 +8,7 @@ from typing import List, Optional, Tuple
 import cv2
 import numpy as np
 
-from ..ch_ppocr_rec.utils import TextRecognizerOutput
+from ..ch_ppocr_rec.utils import TextRecOutput
 
 
 class CalRecBoxes:
@@ -21,8 +21,8 @@ class CalRecBoxes:
         self,
         imgs: Optional[List[np.ndarray]],
         dt_boxes: Optional[List[np.ndarray]],
-        rec_res: TextRecognizerOutput,
-    ) -> TextRecognizerOutput:
+        rec_res: TextRecOutput,
+    ) -> TextRecOutput:
         # rec_res = list(zip(rec_res.line_results, rec_res.word_results))
         word_results = []
         for idx, (img, box) in enumerate(zip(imgs, dt_boxes)):

python/rapidocr/ch_ppocr_rec/__init__.py

@@ -2,4 +2,4 @@
 # @Author: SWHL
 # @Contact: liekkaskono@163.com
 from .main import TextRecognizer
-from .utils import TextRecognizerOutput
+from .utils import TextRecArguments, TextRecOutput

python/rapidocr/ch_ppocr_rec/main.py

@@ -14,14 +14,14 @@
 import argparse
 import math
 import time
-from typing import Any, Dict, List, Union
+from typing import Any, Dict
 
 import cv2
 import numpy as np
 
 from rapidocr.utils import OrtInferSession, read_yaml
 
-from .utils import CTCLabelDecode, TextRecognizerOutput
+from .utils import CTCLabelDecode, TextRecArguments, TextRecOutput
 
 
 class TextRecognizer:
@@ -40,15 +40,10 @@ class TextRecognizer:
         self.rec_batch_num = config["rec_batch_num"]
         self.rec_image_shape = config["rec_img_shape"]
 
-    def __call__(
-        self,
-        img_list: Union[np.ndarray, List[np.ndarray]],
-        return_word_box: bool = False,
-    ) -> TextRecognizerOutput:
-        if isinstance(img_list, np.ndarray):
-            img_list = [img_list]
+    def __call__(self, args: TextRecArguments) -> TextRecOutput:
+        img_list = [args.img] if isinstance(args.img, np.ndarray) else args.img
+        return_word_box = args.return_word_box
 
-        # 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
@@ -96,7 +91,7 @@ class TextRecognizer:
             elapse += time.time() - starttime
 
         all_line_results, all_word_results = list(zip(*rec_res))
-        return TextRecognizerOutput(all_line_results, all_word_results, elapse)
+        return TextRecOutput(all_line_results, all_word_results, 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

python/rapidocr/ch_ppocr_rec/utils.py

@@ -1,7 +1,7 @@
 # -*- encoding: utf-8 -*-
 # @Author: SWHL
 # @Contact: liekkaskono@163.com
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from pathlib import Path
 from typing import List, Optional, Tuple, Union
 
@@ -9,7 +9,7 @@ import numpy as np
 
 
 @dataclass
-class TextRecognizerConfig:
+class TextRecConfig:
     intra_op_num_threads: int = -1
     inter_op_num_threads: int = -1
     use_cuda: bool = False
@@ -21,18 +21,20 @@ class TextRecognizerConfig:
     rec_keys_path: Union[str, Path, None] = None
 
 
+def cvt(img):
+    if isinstance(img, np.ndarray):
+        return [img]
+    return img
+
+
 @dataclass
-class TextRecognizerInput:
-    img_list: Union[np.ndarray, List[np.ndarray]] = field(init=False)
+class TextRecArguments:
+    img: Union[np.ndarray, List[np.ndarray], None] = None
     return_word_box: bool = False
 
-    def __pose_init__(self):
-        if isinstance(self.img_list, np.ndarray):
-            self.img_list = [self.img_list]
-
 
 @dataclass
-class TextRecognizerOutput:
+class TextRecOutput:
     line_results: Optional[Tuple[List]] = None
     word_results: Optional[List[List]] = None
     elapse: Optional[float] = None

python/rapidocr/main.py

@@ -11,7 +11,7 @@ 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, TextRecognizerOutput
+from .ch_ppocr_rec import TextRecArguments, TextRecognizer, TextRecOutput
 from .utils import (
     LoadImage,
     UpdateParameters,
@@ -108,7 +108,8 @@ class RapidOCR:
             img, cls_res, cls_elapse = self.text_cls(img)
 
         if use_rec:
-            rec_res = self.text_rec(img, return_word_box)
+            rec_input = TextRecArguments(img=img, return_word_box=return_word_box)
+            rec_res = self.text_rec(rec_input)
 
         if return_word_box and dt_boxes is not None and all(rec_res.word_results):
             rec_res = self.cal_rec_boxes(img, dt_boxes, rec_res)
@@ -123,14 +124,7 @@ class RapidOCR:
         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
-            dt_boxes,
-            cls_res,
-            rec_res,
-            det_elapse,
-            cls_elapse,
-        )
+        ocr_res = self.get_final_res(dt_boxes, cls_res, rec_res, det_elapse, cls_elapse)
         return ocr_res
 
     def preprocess(self, img: np.ndarray) -> Tuple[np.ndarray, float, float]:
@@ -283,7 +277,7 @@ class RapidOCR:
         self,
         dt_boxes: Optional[List[np.ndarray]],
         cls_res: Optional[List[List[Union[str, float]]]],
-        rec_res: TextRecognizerOutput,
+        rec_res: TextRecOutput,
         det_elapse: float,
         cls_elapse: float,
     ) -> Tuple[Optional[List[List[Union[Any, str]]]], Optional[List[float]]]:
@@ -316,7 +310,7 @@ class RapidOCR:
         self,
         dt_boxes: Optional[List[np.ndarray]],
         rec_res: Optional[List[Tuple[str, float]]],
-    ) -> Tuple[Optional[List[np.ndarray]], TextRecognizerOutput]:
+    ) -> Tuple[Optional[List[np.ndarray]], TextRecOutput]:
         if dt_boxes is None or rec_res is None:
             return None, None