diff --git a/MANIFEST.in b/MANIFEST.in index d674fabc5d714f7e31ed00b32be2d44d6dd10871..1fcf184dacee9dcaf3d5b2e62d12c7b156e068c7 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -5,5 +5,6 @@ recursive-include ppocr/utils *.txt utility.py logging.py network.py recursive-include ppocr/data *.py recursive-include ppocr/postprocess *.py recursive-include tools/infer *.py +recursive-include tools __init__.py recursive-include ppocr/utils/e2e_utils *.py recursive-include ppstructure *.py \ No newline at end of file diff --git a/doc/doc_ch/whl.md b/doc/doc_ch/whl.md index ba5bbae6255382d0c7fa5be319946d6242b1a544..3b88709a2328409a266d0d482baa072dd7aa3824 100644 --- a/doc/doc_ch/whl.md +++ b/doc/doc_ch/whl.md @@ -420,3 +420,5 @@ im_show.save('result.jpg') | cls | 前向时是否启动分类 (命令行模式下使用use_angle_cls控制前向是否启动分类) | FALSE | | show_log | 是否打印det和rec等信息 | FALSE | | type | 执行ocr或者表格结构化, 值可选['ocr','structure'] | ocr | +| ocr_version | OCR模型版本，可选PP-OCRv2, PP-OCR。PP-OCRv2 目前仅支持中文的检测和识别模型，PP-OCR支持中文的检测，识别，多语种识别，方向分类器等模型 | PP-OCRv2 | +| structure_version | 表格结构化模型版本，可选 STRUCTURE。STRUCTURE支持表格结构化模型 | STRUCTURE | diff --git a/doc/doc_en/whl_en.md b/doc/doc_en/whl_en.md index c2577e1e151e4675abab5139da099db9ad20fb4b..62aa452dcd36906c6480031375e6ca94f8a36de3 100644 --- a/doc/doc_en/whl_en.md +++ b/doc/doc_en/whl_en.md @@ -366,4 +366,6 @@ im_show.save('result.jpg') | rec | Enable recognition when `ppocr.ocr` func exec | TRUE | | cls | Enable classification when `ppocr.ocr` func exec((Use use_angle_cls in command line mode to control whether to start classification in the forward direction) | FALSE | | show_log | Whether to print log in det and rec | FALSE | -| type | Perform ocr or table structuring, the value is selected in ['ocr','structure'] | ocr | \ No newline at end of file +| type | Perform ocr or table structuring, the value is selected in ['ocr','structure'] | ocr | +| ocr_version | OCR Model version number, the current model support list is as follows: PP-OCRv2 support Chinese detection and recognition model, PP-OCR support Chinese detection, recognition and direction classifier, multilingual recognition model | PP-OCRv2 | +| structure_version | table structure Model version number, the current model support list is as follows: STRUCTURE support english table structure model | STRUCTURE | diff --git a/paddleocr.py b/paddleocr.py index a98efd34088701d5eb5602743cf75b7d5e80157f..028cfcc1faae3d9cca7d756b55213c030c7496de 100644 --- a/paddleocr.py +++ b/paddleocr.py @@ -16,6 +16,9 @@ import os import sys __dir__ = os.path.dirname(__file__) + +import paddle + sys.path.append(os.path.join(__dir__, '')) import cv2 @@ -29,7 +32,7 @@ from ppocr.utils.logging import get_logger logger = get_logger() from ppocr.utils.utility import check_and_read_gif, get_image_file_list from ppocr.utils.network import maybe_download, download_with_progressbar, is_link, confirm_model_dir_url -from tools.infer.utility import draw_ocr, str2bool +from tools.infer.utility import draw_ocr, str2bool, check_gpu from ppstructure.utility import init_args, draw_structure_result from ppstructure.predict_system import OCRSystem, save_structure_res @@ -39,130 +42,137 @@ __all__ = [ ] SUPPORT_DET_MODEL = ['DB'] -VERSION = '2.2.1' +VERSION = '2.3.0.1' SUPPORT_REC_MODEL = ['CRNN'] BASE_DIR = os.path.expanduser("~/.paddleocr/") -DEFAULT_MODEL_VERSION = '2.0' +DEFAULT_OCR_MODEL_VERSION = 'PP-OCR' +DEFAULT_STRUCTURE_MODEL_VERSION = 'STRUCTURE' MODEL_URLS = { - '2.1': { - 'det': { - 'ch': { - 'url': - 'https://paddleocr.bj.bcebos.com/PP-OCRv2/chinese/ch_PP-OCRv2_det_infer.tar', - }, - }, - 'rec': { - 'ch': { - 'url': - 'https://paddleocr.bj.bcebos.com/PP-OCRv2/chinese/ch_PP-OCRv2_rec_infer.tar', - 'dict_path': './ppocr/utils/ppocr_keys_v1.txt' - } - } - }, - '2.0': { - 'det': { - 'ch': { - 'url': - 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_det_infer.tar', - }, - 'en': { - 'url': - 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/en_ppocr_mobile_v2.0_det_infer.tar', + 'OCR': { + 'PP-OCRv2': { + 'det': { + 'ch': { + 'url': + 'https://paddleocr.bj.bcebos.com/PP-OCRv2/chinese/ch_PP-OCRv2_det_infer.tar', + }, }, - 'structure': { - 'url': - 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_det_infer.tar' + 'rec': { + 'ch': { + 'url': + 'https://paddleocr.bj.bcebos.com/PP-OCRv2/chinese/ch_PP-OCRv2_rec_infer.tar', + 'dict_path': './ppocr/utils/ppocr_keys_v1.txt' + } } }, - 'rec': { - 'ch': { - 'url': - 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_rec_infer.tar', - 'dict_path': './ppocr/utils/ppocr_keys_v1.txt' - }, - 'en': { - 'url': - 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/en_number_mobile_v2.0_rec_infer.tar', - 'dict_path': './ppocr/utils/en_dict.txt' - }, - 'french': { - 'url': - 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/french_mobile_v2.0_rec_infer.tar', - 'dict_path': './ppocr/utils/dict/french_dict.txt' - }, - 'german': { - 'url': - 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/german_mobile_v2.0_rec_infer.tar', - 'dict_path': './ppocr/utils/dict/german_dict.txt' - }, - 'korean': { - 'url': - 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/korean_mobile_v2.0_rec_infer.tar', - 'dict_path': './ppocr/utils/dict/korean_dict.txt' - }, - 'japan': { - 'url': - 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/japan_mobile_v2.0_rec_infer.tar', - 'dict_path': './ppocr/utils/dict/japan_dict.txt' - }, - 'chinese_cht': { - 'url': - 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/chinese_cht_mobile_v2.0_rec_infer.tar', - 'dict_path': './ppocr/utils/dict/chinese_cht_dict.txt' - }, - 'ta': { - 'url': - 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/ta_mobile_v2.0_rec_infer.tar', - 'dict_path': './ppocr/utils/dict/ta_dict.txt' + DEFAULT_OCR_MODEL_VERSION: { + 'det': { + 'ch': { + 'url': + 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_det_infer.tar', + }, + 'en': { + 'url': + 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/en_ppocr_mobile_v2.0_det_infer.tar', + }, + 'structure': { + 'url': + 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_det_infer.tar' + } }, - 'te': { - 'url': - 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/te_mobile_v2.0_rec_infer.tar', - 'dict_path': './ppocr/utils/dict/te_dict.txt' + 'rec': { + 'ch': { + 'url': + 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_rec_infer.tar', + 'dict_path': './ppocr/utils/ppocr_keys_v1.txt' + }, + 'en': { + 'url': + 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/en_number_mobile_v2.0_rec_infer.tar', + 'dict_path': './ppocr/utils/en_dict.txt' + }, + 'french': { + 'url': + 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/french_mobile_v2.0_rec_infer.tar', + 'dict_path': './ppocr/utils/dict/french_dict.txt' + }, + 'german': { + 'url': + 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/german_mobile_v2.0_rec_infer.tar', + 'dict_path': './ppocr/utils/dict/german_dict.txt' + }, + 'korean': { + 'url': + 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/korean_mobile_v2.0_rec_infer.tar', + 'dict_path': './ppocr/utils/dict/korean_dict.txt' + }, + 'japan': { + 'url': + 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/japan_mobile_v2.0_rec_infer.tar', + 'dict_path': './ppocr/utils/dict/japan_dict.txt' + }, + 'chinese_cht': { + 'url': + 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/chinese_cht_mobile_v2.0_rec_infer.tar', + 'dict_path': './ppocr/utils/dict/chinese_cht_dict.txt' + }, + 'ta': { + 'url': + 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/ta_mobile_v2.0_rec_infer.tar', + 'dict_path': './ppocr/utils/dict/ta_dict.txt' + }, + 'te': { + 'url': + 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/te_mobile_v2.0_rec_infer.tar', + 'dict_path': './ppocr/utils/dict/te_dict.txt' + }, + 'ka': { + 'url': + 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/ka_mobile_v2.0_rec_infer.tar', + 'dict_path': './ppocr/utils/dict/ka_dict.txt' + }, + 'latin': { + 'url': + 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/latin_ppocr_mobile_v2.0_rec_infer.tar', + 'dict_path': './ppocr/utils/dict/latin_dict.txt' + }, + 'arabic': { + 'url': + 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/arabic_ppocr_mobile_v2.0_rec_infer.tar', + 'dict_path': './ppocr/utils/dict/arabic_dict.txt' + }, + 'cyrillic': { + 'url': + 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/cyrillic_ppocr_mobile_v2.0_rec_infer.tar', + 'dict_path': './ppocr/utils/dict/cyrillic_dict.txt' + }, + 'devanagari': { + 'url': + 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/devanagari_ppocr_mobile_v2.0_rec_infer.tar', + 'dict_path': './ppocr/utils/dict/devanagari_dict.txt' + }, + 'structure': { + 'url': + 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_rec_infer.tar', + 'dict_path': 'ppocr/utils/dict/table_dict.txt' + } }, - 'ka': { - 'url': - 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/ka_mobile_v2.0_rec_infer.tar', - 'dict_path': './ppocr/utils/dict/ka_dict.txt' + 'cls': { + 'ch': { + 'url': + 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_cls_infer.tar', + } }, - 'latin': { - 'url': - 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/latin_ppocr_mobile_v2.0_rec_infer.tar', - 'dict_path': './ppocr/utils/dict/latin_dict.txt' - }, - 'arabic': { - 'url': - 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/arabic_ppocr_mobile_v2.0_rec_infer.tar', - 'dict_path': './ppocr/utils/dict/arabic_dict.txt' - }, - 'cyrillic': { - 'url': - 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/cyrillic_ppocr_mobile_v2.0_rec_infer.tar', - 'dict_path': './ppocr/utils/dict/cyrillic_dict.txt' - }, - 'devanagari': { - 'url': - 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/devanagari_ppocr_mobile_v2.0_rec_infer.tar', - 'dict_path': './ppocr/utils/dict/devanagari_dict.txt' - }, - 'structure': { - 'url': - 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_rec_infer.tar', - 'dict_path': 'ppocr/utils/dict/table_dict.txt' - } - }, - 'cls': { - 'ch': { - 'url': - 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_cls_infer.tar', - } - }, - 'table': { - 'en': { - 'url': - 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_structure_infer.tar', - 'dict_path': 'ppocr/utils/dict/table_structure_dict.txt' + } + }, + 'STRUCTURE': { + DEFAULT_STRUCTURE_MODEL_VERSION: { + 'table': { + 'en': { + 'url': + 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_structure_infer.tar', + 'dict_path': 'ppocr/utils/dict/table_structure_dict.txt' + } } } } @@ -177,7 +187,20 @@ def parse_args(mMain=True): parser.add_argument("--det", type=str2bool, default=True) parser.add_argument("--rec", type=str2bool, default=True) parser.add_argument("--type", type=str, default='ocr') - parser.add_argument("--version", type=str, default='2.1') + parser.add_argument( + "--ocr_version", + type=str, + default='PP-OCRv2', + help='OCR Model version, the current model support list is as follows: ' + '1. PP-OCRv2 Support Chinese detection and recognition model. ' + '2. PP-OCR support Chinese detection, recognition and direction classifier and multilingual recognition model.' + ) + parser.add_argument( + "--structure_version", + type=str, + default='STRUCTURE', + help='Model version, the current model support list is as follows:' + ' 1. STRUCTURE Support en table structure model.') for action in parser._actions: if action.dest in ['rec_char_dict_path', 'table_char_dict_path']: @@ -215,9 +238,9 @@ def parse_lang(lang): lang = "cyrillic" elif lang in devanagari_lang: lang = "devanagari" - assert lang in MODEL_URLS[DEFAULT_MODEL_VERSION][ + assert lang in MODEL_URLS['OCR'][DEFAULT_OCR_MODEL_VERSION][ 'rec'], 'param lang must in {}, but got {}'.format( - MODEL_URLS[DEFAULT_MODEL_VERSION]['rec'].keys(), lang) + MODEL_URLS['OCR'][DEFAULT_OCR_MODEL_VERSION]['rec'].keys(), lang) if lang == "ch": det_lang = "ch" elif lang == 'structure': @@ -227,33 +250,41 @@ def parse_lang(lang): return lang, det_lang -def get_model_config(version, model_type, lang): - if version not in MODEL_URLS: - logger.warning('version {} not in {}, use version {} instead'.format( - version, MODEL_URLS.keys(), DEFAULT_MODEL_VERSION)) +def get_model_config(type, version, model_type, lang): + if type == 'OCR': + DEFAULT_MODEL_VERSION = DEFAULT_OCR_MODEL_VERSION + elif type == 'STRUCTURE': + DEFAULT_MODEL_VERSION = DEFAULT_STRUCTURE_MODEL_VERSION + else: + raise NotImplementedError + model_urls = MODEL_URLS[type] + if version not in model_urls: + logger.warning('version {} not in {}, auto switch to version {}'.format( + version, model_urls.keys(), DEFAULT_MODEL_VERSION)) version = DEFAULT_MODEL_VERSION - if model_type not in MODEL_URLS[version]: - if model_type in MODEL_URLS[DEFAULT_MODEL_VERSION]: + if model_type not in model_urls[version]: + if model_type in model_urls[DEFAULT_MODEL_VERSION]: logger.warning( - 'version {} not support {} models, use version {} instead'. + 'version {} not support {} models, auto switch to version {}'. format(version, model_type, DEFAULT_MODEL_VERSION)) version = DEFAULT_MODEL_VERSION else: logger.error('{} models is not support, we only support {}'.format( - model_type, MODEL_URLS[DEFAULT_MODEL_VERSION].keys())) + model_type, model_urls[DEFAULT_MODEL_VERSION].keys())) sys.exit(-1) - if lang not in MODEL_URLS[version][model_type]: - if lang in MODEL_URLS[DEFAULT_MODEL_VERSION][model_type]: - logger.warning('lang {} is not support in {}, use {} instead'. - format(lang, version, DEFAULT_MODEL_VERSION)) + if lang not in model_urls[version][model_type]: + if lang in model_urls[DEFAULT_MODEL_VERSION][model_type]: + logger.warning( + 'lang {} is not support in {}, auto switch to version {}'. + format(lang, version, DEFAULT_MODEL_VERSION)) version = DEFAULT_MODEL_VERSION else: logger.error( 'lang {} is not support, we only support {} for {} models'. - format(lang, MODEL_URLS[DEFAULT_MODEL_VERSION][model_type].keys( + format(lang, model_urls[DEFAULT_MODEL_VERSION][model_type].keys( ), model_type)) sys.exit(-1) - return MODEL_URLS[version][model_type][lang] + return model_urls[version][model_type][lang] class PaddleOCR(predict_system.TextSystem): @@ -265,23 +296,28 @@ class PaddleOCR(predict_system.TextSystem): """ params = parse_args(mMain=False) params.__dict__.update(**kwargs) + params.use_gpu = check_gpu(params.use_gpu) + if not params.show_log: logger.setLevel(logging.INFO) self.use_angle_cls = params.use_angle_cls lang, det_lang = parse_lang(params.lang) # init model dir - det_model_config = get_model_config(params.version, 'det', det_lang) + det_model_config = get_model_config('OCR', params.ocr_version, 'det', + det_lang) params.det_model_dir, det_url = confirm_model_dir_url( params.det_model_dir, os.path.join(BASE_DIR, VERSION, 'ocr', 'det', det_lang), det_model_config['url']) - rec_model_config = get_model_config(params.version, 'rec', lang) + rec_model_config = get_model_config('OCR', params.ocr_version, 'rec', + lang) params.rec_model_dir, rec_url = confirm_model_dir_url( params.rec_model_dir, os.path.join(BASE_DIR, VERSION, 'ocr', 'rec', lang), rec_model_config['url']) - cls_model_config = get_model_config(params.version, 'cls', 'ch') + cls_model_config = get_model_config('OCR', params.ocr_version, 'cls', + 'ch') params.cls_model_dir, cls_url = confirm_model_dir_url( params.cls_model_dir, os.path.join(BASE_DIR, VERSION, 'ocr', 'cls'), @@ -362,22 +398,27 @@ class PPStructure(OCRSystem): def __init__(self, **kwargs): params = parse_args(mMain=False) params.__dict__.update(**kwargs) + params.use_gpu = check_gpu(params.use_gpu) + if not params.show_log: logger.setLevel(logging.INFO) lang, det_lang = parse_lang(params.lang) # init model dir - det_model_config = get_model_config(params.version, 'det', det_lang) + det_model_config = get_model_config('OCR', params.ocr_version, 'det', + det_lang) params.det_model_dir, det_url = confirm_model_dir_url( params.det_model_dir, os.path.join(BASE_DIR, VERSION, 'ocr', 'det', det_lang), det_model_config['url']) - rec_model_config = get_model_config(params.version, 'rec', lang) + rec_model_config = get_model_config('OCR', params.ocr_version, 'rec', + lang) params.rec_model_dir, rec_url = confirm_model_dir_url( params.rec_model_dir, os.path.join(BASE_DIR, VERSION, 'ocr', 'rec', lang), rec_model_config['url']) - table_model_config = get_model_config(params.version, 'table', 'en') + table_model_config = get_model_config( + 'STRUCTURE', params.structure_version, 'table', 'en') params.table_model_dir, table_url = confirm_model_dir_url( params.table_model_dir, os.path.join(BASE_DIR, VERSION, 'ocr', 'table'), diff --git a/ppocr/utils/network.py b/ppocr/utils/network.py index 453abb693d4c0ed370c1031b677d5bf51661add9..118d1be364925d9416134cffe21d636fcac753e9 100644 --- a/ppocr/utils/network.py +++ b/ppocr/utils/network.py @@ -24,15 +24,17 @@ from ppocr.utils.logging import get_logger def download_with_progressbar(url, save_path): logger = get_logger() response = requests.get(url, stream=True) - total_size_in_bytes = int(response.headers.get('content-length', 0)) - block_size = 1024 # 1 Kibibyte - progress_bar = tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True) - with open(save_path, 'wb') as file: - for data in response.iter_content(block_size): - progress_bar.update(len(data)) - file.write(data) - progress_bar.close() - if total_size_in_bytes == 0 or progress_bar.n != total_size_in_bytes: + if response.status_code == 200: + total_size_in_bytes = int(response.headers.get('content-length', 1)) + block_size = 1024 # 1 Kibibyte + progress_bar = tqdm( + total=total_size_in_bytes, unit='iB', unit_scale=True) + with open(save_path, 'wb') as file: + for data in response.iter_content(block_size): + progress_bar.update(len(data)) + file.write(data) + progress_bar.close() + else: logger.error("Something went wrong while downloading models") sys.exit(0) @@ -45,7 +47,7 @@ def maybe_download(model_storage_directory, url): if not os.path.exists( os.path.join(model_storage_directory, 'inference.pdiparams') ) or not os.path.exists( - os.path.join(model_storage_directory, 'inference.pdmodel')): + os.path.join(model_storage_directory, 'inference.pdmodel')): assert url.endswith('.tar'), 'Only supports tar compressed package' tmp_path = os.path.join(model_storage_directory, url.split('/')[-1]) print('download {} to {}'.format(url, tmp_path)) diff --git a/tools/__init__.py b/tools/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d56c9dbaa1304b160521da03c05db2352e341bf2 --- /dev/null +++ b/tools/__init__.py @@ -0,0 +1,14 @@ +# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved. +# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team. +# +# 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. diff --git a/tools/infer/utility.py b/tools/infer/utility.py index 98bf0362f6842cb70490e6817ef53fe39109f406..cab918419ab5efbc4a8a11d1669ca6b93e45e789 100755 --- a/tools/infer/utility.py +++ b/tools/infer/utility.py @@ -17,7 +17,7 @@ import os import sys import cv2 import numpy as np -import json +import paddle from PIL import Image, ImageDraw, ImageFont import math from paddle import inference @@ -601,5 +601,12 @@ def get_rotate_crop_image(img, points): return dst_img +def check_gpu(use_gpu): + if use_gpu and not paddle.is_compiled_with_cuda(): + + use_gpu = False + return use_gpu + + if __name__ == '__main__': pass