From 9551d8bb34d14def15794ad643d82e226056d1ec Mon Sep 17 00:00:00 2001 From: SWHL Date: Thu, 21 Sep 2023 16:50:59 +0800 Subject: [PATCH] Add to support visisual result under CLI --- .../gen_whl_to_pypi_rapidocr_ort.yml | 2 +- .../gen_whl_to_pypi_rapidocr_vino.yml | 2 +- python/demo.py | 8 +++--- python/rapidocr_onnxruntime/main.py | 27 ++++++++++++++++++- python/rapidocr_onnxruntime/utils.py | 17 +++++++++++- python/rapidocr_openvino/main.py | 27 ++++++++++++++++++- python/rapidocr_openvino/utils.py | 17 +++++++++++- python/setup_onnxruntime.py | 2 +- python/setup_openvino.py | 2 +- 9 files changed, 92 insertions(+), 12 deletions(-) diff --git a/.github/workflows/gen_whl_to_pypi_rapidocr_ort.yml b/.github/workflows/gen_whl_to_pypi_rapidocr_ort.yml index f8b0c9506..5ca4f686b 100644 --- a/.github/workflows/gen_whl_to_pypi_rapidocr_ort.yml +++ b/.github/workflows/gen_whl_to_pypi_rapidocr_ort.yml @@ -86,7 +86,7 @@ jobs: mv rapidocr_onnxruntime rapidocr_onnxruntime_t mv rapidocr_onnxruntime_t rapidocr_onnxruntime cd rapidocr_onnxruntime - echo "from .rapidocr_onnxruntime.main import RapidOCR" > __init__.py + echo "from .rapidocr_onnxruntime.main import RapidOCR, VisRes" > __init__.py cd .. python -m pip install --upgrade pip diff --git a/.github/workflows/gen_whl_to_pypi_rapidocr_vino.yml b/.github/workflows/gen_whl_to_pypi_rapidocr_vino.yml index 54b39e7df..5dec3d82b 100644 --- a/.github/workflows/gen_whl_to_pypi_rapidocr_vino.yml +++ b/.github/workflows/gen_whl_to_pypi_rapidocr_vino.yml @@ -86,7 +86,7 @@ jobs: mv rapidocr_openvino rapidocr_openvino_t mv rapidocr_openvino_t rapidocr_openvino cd rapidocr_openvino - echo "from .rapidocr_openvino.main import RapidOCR" > __init__.py + echo "from .rapidocr_openvino.main import RapidOCR, VisRes" > __init__.py cd .. python -m pip install --upgrade pip diff --git a/python/demo.py b/python/demo.py index b791ce9fa..07923444d 100644 --- a/python/demo.py +++ b/python/demo.py @@ -8,17 +8,17 @@ # from rapidocr_openvino import RapidOCR, VisRes -rapid_ocr = RapidOCR() +engine = RapidOCR() vis = VisRes(font_path="resources/fonts/FZYTK.TTF") image_path = "tests/test_files/ch_en_num.jpg" with open(image_path, "rb") as f: img = f.read() -result, elapse_list = rapid_ocr(img) +result, elapse_list = engine(img) print(result) print(elapse_list) boxes, txts, scores = list(zip(*result)) -res = vis(img, boxes, txts, scores) -cv2.imwrite("vis.png", res) +vis_img = vis(img, boxes, txts, scores) +cv2.imwrite("vis.png", vis_img) diff --git a/python/rapidocr_onnxruntime/main.py b/python/rapidocr_onnxruntime/main.py index 46e06b7d3..e5b59eb15 100644 --- a/python/rapidocr_onnxruntime/main.py +++ b/python/rapidocr_onnxruntime/main.py @@ -11,7 +11,14 @@ from .ch_ppocr_v2_cls import TextClassifier from .ch_ppocr_v3_det import TextDetector from .ch_ppocr_v3_rec import TextRecognizer -from .utils import LoadImage, UpdateParameters, concat_model_path, init_args, read_yaml +from .utils import ( + LoadImage, + UpdateParameters, + VisRes, + concat_model_path, + init_args, + read_yaml, +) root_dir = Path(__file__).resolve().parent @@ -257,9 +264,27 @@ def main(): args.img_path, use_det=use_det, use_cls=use_cls, use_rec=use_rec ) print(result) + if args.print_cost: print(elapse_list) + if args.vis_res and args.vis_font_path: + vis = VisRes(font_path=args.vis_font_path) + 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, None, None) + cv2.imwrite(str(save_path), vis_img) + print(f"The vis result has saved in {save_path}") + elif use_det and use_rec: + boxes, txts, scores = list(zip(*result)) + vis_img = vis(args.img_path, boxes, txts, scores) + cv2.imwrite(str(save_path), vis_img) + print(f"The vis result has saved in {save_path}") + if __name__ == "__main__": main() diff --git a/python/rapidocr_onnxruntime/utils.py b/python/rapidocr_onnxruntime/utils.py index f4452cea2..b762e45f6 100644 --- a/python/rapidocr_onnxruntime/utils.py +++ b/python/rapidocr_onnxruntime/utils.py @@ -265,6 +265,21 @@ def init_args(): 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 @@ -352,7 +367,7 @@ class VisRes: def __init__( self, font_path: Optional[Union[str, Path]] = None, text_score: float = 0.5 ): - if font_path is None: + 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" diff --git a/python/rapidocr_openvino/main.py b/python/rapidocr_openvino/main.py index 46e06b7d3..e5b59eb15 100644 --- a/python/rapidocr_openvino/main.py +++ b/python/rapidocr_openvino/main.py @@ -11,7 +11,14 @@ from .ch_ppocr_v2_cls import TextClassifier from .ch_ppocr_v3_det import TextDetector from .ch_ppocr_v3_rec import TextRecognizer -from .utils import LoadImage, UpdateParameters, concat_model_path, init_args, read_yaml +from .utils import ( + LoadImage, + UpdateParameters, + VisRes, + concat_model_path, + init_args, + read_yaml, +) root_dir = Path(__file__).resolve().parent @@ -257,9 +264,27 @@ def main(): args.img_path, use_det=use_det, use_cls=use_cls, use_rec=use_rec ) print(result) + if args.print_cost: print(elapse_list) + if args.vis_res and args.vis_font_path: + vis = VisRes(font_path=args.vis_font_path) + 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, None, None) + cv2.imwrite(str(save_path), vis_img) + print(f"The vis result has saved in {save_path}") + elif use_det and use_rec: + boxes, txts, scores = list(zip(*result)) + vis_img = vis(args.img_path, boxes, txts, scores) + cv2.imwrite(str(save_path), vis_img) + print(f"The vis result has saved in {save_path}") + if __name__ == "__main__": main() diff --git a/python/rapidocr_openvino/utils.py b/python/rapidocr_openvino/utils.py index 32aadbb37..01313971a 100644 --- a/python/rapidocr_openvino/utils.py +++ b/python/rapidocr_openvino/utils.py @@ -194,6 +194,21 @@ def init_args(): 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 @@ -281,7 +296,7 @@ class VisRes: def __init__( self, font_path: Optional[Union[str, Path]] = None, text_score: float = 0.5 ): - if font_path is None: + 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" diff --git a/python/setup_onnxruntime.py b/python/setup_onnxruntime.py index c16b0eaa9..f861b740f 100644 --- a/python/setup_onnxruntime.py +++ b/python/setup_onnxruntime.py @@ -38,7 +38,7 @@ def get_readme(): sys.argv = sys.argv[:2] project_urls = { - "Documentation": "https://rapidai.github.io/RapidOCRDocs/docs/install_usage/rapidocr/rapidocr_onnxruntime/", + "Documentation": "https://rapidai.github.io/RapidOCRDocs/docs/install_usage/rapidocr/usage/", "Changelog": "https://rapidai.github.io/RapidOCRDocs/docs/changelog/rapidocr/", } diff --git a/python/setup_openvino.py b/python/setup_openvino.py index 8fba83200..afa4b1731 100644 --- a/python/setup_openvino.py +++ b/python/setup_openvino.py @@ -38,7 +38,7 @@ def get_readme(): sys.argv = sys.argv[:2] project_urls = { - "Documentation": "https://rapidai.github.io/RapidOCRDocs/docs/install_usage/rapidocr/rapidocr_openvino/", + "Documentation": "https://rapidai.github.io/RapidOCRDocs/docs/install_usage/rapidocr/usage/", "Changelog": "https://rapidai.github.io/RapidOCRDocs/docs/changelog/rapidocr/", }