diff --git a/python/rapidocr/cal_rec_boxes/main.py b/python/rapidocr/cal_rec_boxes/main.py index 296997f..4fb2736 100644 --- a/python/rapidocr/cal_rec_boxes/main.py +++ b/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)): diff --git a/python/rapidocr/ch_ppocr_rec/__init__.py b/python/rapidocr/ch_ppocr_rec/__init__.py index 46477ba..1faca81 100644 --- a/python/rapidocr/ch_ppocr_rec/__init__.py +++ b/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 diff --git a/python/rapidocr/ch_ppocr_rec/main.py b/python/rapidocr/ch_ppocr_rec/main.py index a37d575..25d1bee 100644 --- a/python/rapidocr/ch_ppocr_rec/main.py +++ b/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 diff --git a/python/rapidocr/ch_ppocr_rec/utils.py b/python/rapidocr/ch_ppocr_rec/utils.py index 1fadcf8..ea36d4b 100644 --- a/python/rapidocr/ch_ppocr_rec/utils.py +++ b/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 diff --git a/python/rapidocr/main.py b/python/rapidocr/main.py index e818b6d..945dacc 100644 --- a/python/rapidocr/main.py +++ b/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