diff --git a/python/rapidocr_onnxruntime/rapid_ocr_api.py b/python/rapidocr_onnxruntime/rapid_ocr_api.py index 31a3d5d..1198d44 100644 --- a/python/rapidocr_onnxruntime/rapid_ocr_api.py +++ b/python/rapidocr_onnxruntime/rapid_ocr_api.py @@ -1,6 +1,7 @@ # -*- encoding: utf-8 -*- # @Author: SWHL # @Contact: liekkaskono@163.com +import argparse import copy import importlib import sys @@ -173,11 +174,20 @@ class RapidOCR(): return filter_boxes, filter_rec_res -if __name__ == '__main__': - rapid_ocr = RapidOCR('config.yaml') +def main(): + parser = argparse.ArgumentParser('RapidOCR') + parser.add_argument('-img', '--img_path', type=str, default=None) + parser.add_argument('-p', '--print_cost', + action='store_true', default=False) + args = parser.parse_args() - import cv2 - img = cv2.imread('resources/test_images/det_images/ch_en_num.jpg') + ocr_engine = RapidOCR() - result = rapid_ocr(img) + result, elapse_list = ocr_engine(args.img_path) print(result) + if args.print_cost: + print(elapse_list) + + +if __name__ == '__main__': + main() diff --git a/python/rapidocr_openvino/rapid_ocr_api.py b/python/rapidocr_openvino/rapid_ocr_api.py index 68ecd67..1198d44 100644 --- a/python/rapidocr_openvino/rapid_ocr_api.py +++ b/python/rapidocr_openvino/rapid_ocr_api.py @@ -1,14 +1,16 @@ # -*- encoding: utf-8 -*- # @Author: SWHL # @Contact: liekkaskono@163.com +import argparse import copy import importlib import sys from pathlib import Path +from typing import Union import cv2 import numpy as np -from typing import Union + from .utils import LoadImage, read_yaml root_dir = Path(__file__).resolve().parent @@ -50,12 +52,10 @@ class RapidOCR(): def __call__(self, img_content: Union[str, np.ndarray, bytes, Path], **kwargs): 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) - # 更新超参数 self.text_detector.postprocess_op.box_thresh = box_thresh self.text_detector.postprocess_op.unclip_ratio = unclip_ratio self.text_score = text_score @@ -100,12 +100,6 @@ class RapidOCR(): for dt, rec in zip(filter_boxes, filter_rec_res)] return fina_result, [det_elapse, cls_elapse, rec_elapse] - @staticmethod - def read_yaml(yaml_path): - with open(yaml_path, 'rb') as f: - data = yaml.load(f, Loader=yaml.Loader) - return data - @staticmethod def init_module(module_name, class_name): module_part = importlib.import_module(module_name) @@ -180,11 +174,20 @@ class RapidOCR(): return filter_boxes, filter_rec_res -if __name__ == '__main__': - rapid_ocr = RapidOCR('config.yaml') +def main(): + parser = argparse.ArgumentParser('RapidOCR') + parser.add_argument('-img', '--img_path', type=str, default=None) + parser.add_argument('-p', '--print_cost', + action='store_true', default=False) + args = parser.parse_args() - import cv2 - img = cv2.imread('resources/test_images/det_images/ch_en_num.jpg') + ocr_engine = RapidOCR() - result = rapid_ocr(img) + result, elapse_list = ocr_engine(args.img_path) print(result) + if args.print_cost: + print(elapse_list) + + +if __name__ == '__main__': + main() diff --git a/python/setup_onnxruntime.py b/python/setup_onnxruntime.py index ce96d48..c95dcf8 100644 --- a/python/setup_onnxruntime.py +++ b/python/setup_onnxruntime.py @@ -84,5 +84,8 @@ setuptools.setup( 'Programming Language :: Python :: 3.8', 'Programming Language :: Python :: 3.9', 'Programming Language :: Python :: 3.10', - ] + ], + entry_points={ + 'console_scripts': [f'{MODULE_NAME}={MODULE_NAME}.rapid_ocr_api:main'], + } ) diff --git a/python/setup_openvino.py b/python/setup_openvino.py index 64e5823..caeb2bc 100644 --- a/python/setup_openvino.py +++ b/python/setup_openvino.py @@ -84,5 +84,8 @@ setuptools.setup( 'Programming Language :: Python :: 3.8', 'Programming Language :: Python :: 3.9', 'Programming Language :: Python :: 3.10', - ] + ], + entry_points={ + 'console_scripts': [f'{MODULE_NAME}={MODULE_NAME}.rapid_ocr_api:main'], + } )