2020-08-22 19:42:14 +08:00
|
|
|
|
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
|
|
|
|
|
#
|
|
|
|
|
# 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.
|
|
|
|
|
|
|
|
|
|
import os
|
|
|
|
|
import sys
|
|
|
|
|
|
|
|
|
|
__dir__ = os.path.dirname(__file__)
|
|
|
|
|
sys.path.append(os.path.join(__dir__, ''))
|
|
|
|
|
|
|
|
|
|
import cv2
|
2021-06-10 14:24:59 +08:00
|
|
|
|
import logging
|
2020-08-22 19:42:14 +08:00
|
|
|
|
import numpy as np
|
|
|
|
|
from pathlib import Path
|
|
|
|
|
|
|
|
|
|
from tools.infer import predict_system
|
2020-12-07 19:10:19 +08:00
|
|
|
|
from ppocr.utils.logging import get_logger
|
2020-08-22 19:42:14 +08:00
|
|
|
|
|
2020-12-07 19:10:19 +08:00
|
|
|
|
logger = get_logger()
|
2020-08-24 11:30:56 +08:00
|
|
|
|
from ppocr.utils.utility import check_and_read_gif, get_image_file_list
|
2021-06-10 14:47:23 +08:00
|
|
|
|
from ppocr.utils.network import maybe_download, download_with_progressbar, is_link, confirm_model_dir_url
|
2021-08-02 15:28:07 +08:00
|
|
|
|
from tools.infer.utility import draw_ocr, str2bool
|
|
|
|
|
from ppstructure.utility import init_args, draw_structure_result
|
|
|
|
|
from ppstructure.predict_system import OCRSystem, save_structure_res
|
2020-08-22 19:42:14 +08:00
|
|
|
|
|
2021-09-06 18:33:21 +08:00
|
|
|
|
__all__ = [
|
|
|
|
|
'PaddleOCR', 'PPStructure', 'draw_ocr', 'draw_structure_result',
|
|
|
|
|
'save_structure_res', 'download_with_progressbar'
|
|
|
|
|
]
|
|
|
|
|
|
|
|
|
|
SUPPORT_DET_MODEL = ['DB']
|
|
|
|
|
VERSION = '2.2.1'
|
|
|
|
|
SUPPORT_REC_MODEL = ['CRNN']
|
|
|
|
|
BASE_DIR = os.path.expanduser("~/.paddleocr/")
|
|
|
|
|
|
|
|
|
|
DEFAULT_MODEL_VERSION = '2.0'
|
|
|
|
|
MODEL_URLS = {
|
|
|
|
|
'2.1': {
|
|
|
|
|
'det': {
|
|
|
|
|
'ch': {
|
|
|
|
|
'url':
|
2021-09-07 14:32:06 +08:00
|
|
|
|
'https://paddleocr.bj.bcebos.com/PP-OCRv2/chinese/ch_PP-OCRv2_det_infer.tar',
|
2021-09-06 18:33:21 +08:00
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
'rec': {
|
|
|
|
|
'ch': {
|
|
|
|
|
'url':
|
2021-09-07 14:32:06 +08:00
|
|
|
|
'https://paddleocr.bj.bcebos.com/PP-OCRv2/chinese/ch_PP-OCRv2_rec_infer.tar',
|
2021-09-06 18:33:21 +08:00
|
|
|
|
'dict_path': './ppocr/utils/ppocr_keys_v1.txt'
|
|
|
|
|
}
|
|
|
|
|
}
|
2021-04-13 17:54:10 +08:00
|
|
|
|
},
|
2021-09-06 18:33:21 +08:00
|
|
|
|
'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',
|
|
|
|
|
},
|
|
|
|
|
'structure': {
|
|
|
|
|
'url':
|
|
|
|
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_det_infer.tar'
|
|
|
|
|
}
|
2020-12-07 19:10:19 +08:00
|
|
|
|
},
|
2021-09-06 18:33:21 +08:00
|
|
|
|
'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':
|
2021-06-10 14:24:59 +08:00
|
|
|
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/en_number_mobile_v2.0_rec_infer.tar',
|
2021-09-06 18:33:21 +08:00
|
|
|
|
'dict_path': './ppocr/utils/en_dict.txt'
|
|
|
|
|
},
|
|
|
|
|
'french': {
|
|
|
|
|
'url':
|
2021-06-10 14:24:59 +08:00
|
|
|
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/french_mobile_v2.0_rec_infer.tar',
|
2021-09-06 18:33:21 +08:00
|
|
|
|
'dict_path': './ppocr/utils/dict/french_dict.txt'
|
|
|
|
|
},
|
|
|
|
|
'german': {
|
|
|
|
|
'url':
|
2021-06-10 14:24:59 +08:00
|
|
|
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/german_mobile_v2.0_rec_infer.tar',
|
2021-09-06 18:33:21 +08:00
|
|
|
|
'dict_path': './ppocr/utils/dict/german_dict.txt'
|
|
|
|
|
},
|
|
|
|
|
'korean': {
|
|
|
|
|
'url':
|
2021-06-10 14:24:59 +08:00
|
|
|
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/korean_mobile_v2.0_rec_infer.tar',
|
2021-09-06 18:33:21 +08:00
|
|
|
|
'dict_path': './ppocr/utils/dict/korean_dict.txt'
|
|
|
|
|
},
|
|
|
|
|
'japan': {
|
|
|
|
|
'url':
|
2021-06-10 14:24:59 +08:00
|
|
|
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/japan_mobile_v2.0_rec_infer.tar',
|
2021-09-06 18:33:21 +08:00
|
|
|
|
'dict_path': './ppocr/utils/dict/japan_dict.txt'
|
|
|
|
|
},
|
|
|
|
|
'chinese_cht': {
|
|
|
|
|
'url':
|
2021-06-10 14:24:59 +08:00
|
|
|
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/chinese_cht_mobile_v2.0_rec_infer.tar',
|
2021-09-06 18:33:21 +08:00
|
|
|
|
'dict_path': './ppocr/utils/dict/chinese_cht_dict.txt'
|
|
|
|
|
},
|
|
|
|
|
'ta': {
|
|
|
|
|
'url':
|
2021-06-10 14:24:59 +08:00
|
|
|
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/ta_mobile_v2.0_rec_infer.tar',
|
2021-09-06 18:33:21 +08:00
|
|
|
|
'dict_path': './ppocr/utils/dict/ta_dict.txt'
|
|
|
|
|
},
|
|
|
|
|
'te': {
|
|
|
|
|
'url':
|
2021-06-10 14:24:59 +08:00
|
|
|
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/te_mobile_v2.0_rec_infer.tar',
|
2021-09-06 18:33:21 +08:00
|
|
|
|
'dict_path': './ppocr/utils/dict/te_dict.txt'
|
|
|
|
|
},
|
|
|
|
|
'ka': {
|
|
|
|
|
'url':
|
2021-06-10 14:24:59 +08:00
|
|
|
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/ka_mobile_v2.0_rec_infer.tar',
|
2021-09-06 18:33:21 +08:00
|
|
|
|
'dict_path': './ppocr/utils/dict/ka_dict.txt'
|
|
|
|
|
},
|
|
|
|
|
'latin': {
|
|
|
|
|
'url':
|
2021-06-10 14:24:59 +08:00
|
|
|
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/latin_ppocr_mobile_v2.0_rec_infer.tar',
|
2021-09-06 18:33:21 +08:00
|
|
|
|
'dict_path': './ppocr/utils/dict/latin_dict.txt'
|
|
|
|
|
},
|
|
|
|
|
'arabic': {
|
|
|
|
|
'url':
|
2021-06-10 14:24:59 +08:00
|
|
|
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/arabic_ppocr_mobile_v2.0_rec_infer.tar',
|
2021-09-06 18:33:21 +08:00
|
|
|
|
'dict_path': './ppocr/utils/dict/arabic_dict.txt'
|
|
|
|
|
},
|
|
|
|
|
'cyrillic': {
|
|
|
|
|
'url':
|
2021-06-10 14:24:59 +08:00
|
|
|
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/cyrillic_ppocr_mobile_v2.0_rec_infer.tar',
|
2021-09-06 18:33:21 +08:00
|
|
|
|
'dict_path': './ppocr/utils/dict/cyrillic_dict.txt'
|
|
|
|
|
},
|
|
|
|
|
'devanagari': {
|
|
|
|
|
'url':
|
2021-06-10 14:24:59 +08:00
|
|
|
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/devanagari_ppocr_mobile_v2.0_rec_infer.tar',
|
2021-09-06 18:33:21 +08:00
|
|
|
|
'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',
|
|
|
|
|
}
|
2021-08-02 15:28:07 +08:00
|
|
|
|
},
|
2021-09-06 18:33:21 +08:00
|
|
|
|
'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'
|
|
|
|
|
}
|
2020-12-07 19:10:19 +08:00
|
|
|
|
}
|
2021-08-02 15:28:07 +08:00
|
|
|
|
}
|
2020-08-22 19:42:14 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
2021-06-02 20:10:59 +08:00
|
|
|
|
def parse_args(mMain=True):
|
2020-08-22 19:42:14 +08:00
|
|
|
|
import argparse
|
2021-06-02 20:10:59 +08:00
|
|
|
|
parser = init_args()
|
|
|
|
|
parser.add_help = mMain
|
|
|
|
|
parser.add_argument("--lang", type=str, default='ch')
|
|
|
|
|
parser.add_argument("--det", type=str2bool, default=True)
|
|
|
|
|
parser.add_argument("--rec", type=str2bool, default=True)
|
2021-08-02 15:28:07 +08:00
|
|
|
|
parser.add_argument("--type", type=str, default='ocr')
|
2021-09-06 18:33:21 +08:00
|
|
|
|
parser.add_argument("--version", type=str, default='2.1')
|
2021-06-02 20:10:59 +08:00
|
|
|
|
|
|
|
|
|
for action in parser._actions:
|
2021-08-02 15:28:07 +08:00
|
|
|
|
if action.dest in ['rec_char_dict_path', 'table_char_dict_path']:
|
2021-06-02 20:10:59 +08:00
|
|
|
|
action.default = None
|
2020-12-07 19:10:19 +08:00
|
|
|
|
if mMain:
|
2021-06-02 20:10:59 +08:00
|
|
|
|
return parser.parse_args()
|
2020-12-07 19:10:19 +08:00
|
|
|
|
else:
|
2021-05-26 17:34:47 +08:00
|
|
|
|
inference_args_dict = {}
|
2021-06-02 20:10:59 +08:00
|
|
|
|
for action in parser._actions:
|
|
|
|
|
inference_args_dict[action.dest] = action.default
|
2021-05-26 17:34:47 +08:00
|
|
|
|
return argparse.Namespace(**inference_args_dict)
|
2020-08-22 19:42:14 +08:00
|
|
|
|
|
|
|
|
|
|
2021-08-02 17:04:53 +08:00
|
|
|
|
def parse_lang(lang):
|
|
|
|
|
latin_lang = [
|
2021-09-06 18:33:21 +08:00
|
|
|
|
'af', 'az', 'bs', 'cs', 'cy', 'da', 'de', 'es', 'et', 'fr', 'ga', 'hr',
|
|
|
|
|
'hu', 'id', 'is', 'it', 'ku', 'la', 'lt', 'lv', 'mi', 'ms', 'mt', 'nl',
|
|
|
|
|
'no', 'oc', 'pi', 'pl', 'pt', 'ro', 'rs_latin', 'sk', 'sl', 'sq', 'sv',
|
|
|
|
|
'sw', 'tl', 'tr', 'uz', 'vi'
|
2021-08-02 17:04:53 +08:00
|
|
|
|
]
|
|
|
|
|
arabic_lang = ['ar', 'fa', 'ug', 'ur']
|
|
|
|
|
cyrillic_lang = [
|
2021-09-06 18:33:21 +08:00
|
|
|
|
'ru', 'rs_cyrillic', 'be', 'bg', 'uk', 'mn', 'abq', 'ady', 'kbd', 'ava',
|
|
|
|
|
'dar', 'inh', 'che', 'lbe', 'lez', 'tab'
|
2021-08-02 17:04:53 +08:00
|
|
|
|
]
|
|
|
|
|
devanagari_lang = [
|
2021-09-06 18:33:21 +08:00
|
|
|
|
'hi', 'mr', 'ne', 'bh', 'mai', 'ang', 'bho', 'mah', 'sck', 'new', 'gom',
|
|
|
|
|
'sa', 'bgc'
|
2021-08-02 17:04:53 +08:00
|
|
|
|
]
|
|
|
|
|
if lang in latin_lang:
|
|
|
|
|
lang = "latin"
|
|
|
|
|
elif lang in arabic_lang:
|
|
|
|
|
lang = "arabic"
|
|
|
|
|
elif lang in cyrillic_lang:
|
|
|
|
|
lang = "cyrillic"
|
|
|
|
|
elif lang in devanagari_lang:
|
|
|
|
|
lang = "devanagari"
|
2021-09-06 18:33:21 +08:00
|
|
|
|
assert lang in MODEL_URLS[DEFAULT_MODEL_VERSION][
|
2021-08-02 17:04:53 +08:00
|
|
|
|
'rec'], 'param lang must in {}, but got {}'.format(
|
2021-09-06 18:33:21 +08:00
|
|
|
|
MODEL_URLS[DEFAULT_MODEL_VERSION]['rec'].keys(), lang)
|
2021-08-02 17:04:53 +08:00
|
|
|
|
if lang == "ch":
|
|
|
|
|
det_lang = "ch"
|
|
|
|
|
elif lang == 'structure':
|
|
|
|
|
det_lang = 'structure'
|
|
|
|
|
else:
|
|
|
|
|
det_lang = "en"
|
|
|
|
|
return lang, det_lang
|
|
|
|
|
|
|
|
|
|
|
2021-09-06 18:33:21 +08:00
|
|
|
|
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))
|
|
|
|
|
version = 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'.
|
|
|
|
|
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()))
|
|
|
|
|
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))
|
|
|
|
|
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(
|
|
|
|
|
), model_type))
|
|
|
|
|
sys.exit(-1)
|
|
|
|
|
return MODEL_URLS[version][model_type][lang]
|
|
|
|
|
|
|
|
|
|
|
2020-08-22 19:42:14 +08:00
|
|
|
|
class PaddleOCR(predict_system.TextSystem):
|
2020-08-24 11:30:56 +08:00
|
|
|
|
def __init__(self, **kwargs):
|
2020-08-22 19:42:14 +08:00
|
|
|
|
"""
|
|
|
|
|
paddleocr package
|
|
|
|
|
args:
|
|
|
|
|
**kwargs: other params show in paddleocr --help
|
|
|
|
|
"""
|
2021-06-10 14:24:59 +08:00
|
|
|
|
params = parse_args(mMain=False)
|
|
|
|
|
params.__dict__.update(**kwargs)
|
2021-06-10 17:12:14 +08:00
|
|
|
|
if not params.show_log:
|
|
|
|
|
logger.setLevel(logging.INFO)
|
2021-06-10 14:24:59 +08:00
|
|
|
|
self.use_angle_cls = params.use_angle_cls
|
2021-08-02 17:04:53 +08:00
|
|
|
|
lang, det_lang = parse_lang(params.lang)
|
2020-08-22 19:42:14 +08:00
|
|
|
|
|
2020-08-24 11:30:56 +08:00
|
|
|
|
# init model dir
|
2021-09-06 18:33:21 +08:00
|
|
|
|
det_model_config = get_model_config(params.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)
|
|
|
|
|
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')
|
|
|
|
|
params.cls_model_dir, cls_url = confirm_model_dir_url(
|
|
|
|
|
params.cls_model_dir,
|
|
|
|
|
os.path.join(BASE_DIR, VERSION, 'ocr', 'cls'),
|
|
|
|
|
cls_model_config['url'])
|
2020-08-22 19:42:14 +08:00
|
|
|
|
# download model
|
2021-06-10 14:47:23 +08:00
|
|
|
|
maybe_download(params.det_model_dir, det_url)
|
|
|
|
|
maybe_download(params.rec_model_dir, rec_url)
|
|
|
|
|
maybe_download(params.cls_model_dir, cls_url)
|
2020-08-22 19:42:14 +08:00
|
|
|
|
|
2021-06-10 14:24:59 +08:00
|
|
|
|
if params.det_algorithm not in SUPPORT_DET_MODEL:
|
2020-08-22 19:42:14 +08:00
|
|
|
|
logger.error('det_algorithm must in {}'.format(SUPPORT_DET_MODEL))
|
|
|
|
|
sys.exit(0)
|
2021-06-10 14:24:59 +08:00
|
|
|
|
if params.rec_algorithm not in SUPPORT_REC_MODEL:
|
2020-08-22 19:42:14 +08:00
|
|
|
|
logger.error('rec_algorithm must in {}'.format(SUPPORT_REC_MODEL))
|
|
|
|
|
sys.exit(0)
|
2021-08-02 17:04:53 +08:00
|
|
|
|
|
|
|
|
|
if params.rec_char_dict_path is None:
|
2021-09-06 18:33:21 +08:00
|
|
|
|
params.rec_char_dict_path = str(
|
|
|
|
|
Path(__file__).parent / rec_model_config['dict_path'])
|
2020-08-22 19:42:14 +08:00
|
|
|
|
|
2021-06-10 14:24:59 +08:00
|
|
|
|
print(params)
|
2020-08-22 19:42:14 +08:00
|
|
|
|
# init det_model and rec_model
|
2021-06-10 14:24:59 +08:00
|
|
|
|
super().__init__(params)
|
2020-08-22 19:42:14 +08:00
|
|
|
|
|
2021-05-26 17:34:47 +08:00
|
|
|
|
def ocr(self, img, det=True, rec=True, cls=True):
|
2020-08-22 19:42:14 +08:00
|
|
|
|
"""
|
|
|
|
|
ocr with paddleocr
|
|
|
|
|
args:
|
|
|
|
|
img: img for ocr, support ndarray, img_path and list or ndarray
|
|
|
|
|
det: use text detection or not, if false, only rec will be exec. default is True
|
|
|
|
|
rec: use text recognition or not, if false, only det will be exec. default is True
|
|
|
|
|
"""
|
|
|
|
|
assert isinstance(img, (np.ndarray, list, str))
|
2020-12-07 19:10:19 +08:00
|
|
|
|
if isinstance(img, list) and det == True:
|
|
|
|
|
logger.error('When input a list of images, det must be false')
|
|
|
|
|
exit(0)
|
2021-05-26 17:34:47 +08:00
|
|
|
|
if cls == True and self.use_angle_cls == False:
|
2021-02-24 15:38:54 +08:00
|
|
|
|
logger.warning(
|
|
|
|
|
'Since the angle classifier is not initialized, the angle classifier will not be uesd during the forward process'
|
|
|
|
|
)
|
2020-12-07 19:10:19 +08:00
|
|
|
|
|
2020-08-22 19:42:14 +08:00
|
|
|
|
if isinstance(img, str):
|
2020-12-07 19:10:19 +08:00
|
|
|
|
# download net image
|
|
|
|
|
if img.startswith('http'):
|
|
|
|
|
download_with_progressbar(img, 'tmp.jpg')
|
|
|
|
|
img = 'tmp.jpg'
|
2020-08-22 19:42:14 +08:00
|
|
|
|
image_file = img
|
|
|
|
|
img, flag = check_and_read_gif(image_file)
|
|
|
|
|
if not flag:
|
2021-01-22 19:15:42 +08:00
|
|
|
|
with open(image_file, 'rb') as f:
|
|
|
|
|
np_arr = np.frombuffer(f.read(), dtype=np.uint8)
|
|
|
|
|
img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
|
2020-08-22 19:42:14 +08:00
|
|
|
|
if img is None:
|
|
|
|
|
logger.error("error in loading image:{}".format(image_file))
|
|
|
|
|
return None
|
2020-12-07 19:10:19 +08:00
|
|
|
|
if isinstance(img, np.ndarray) and len(img.shape) == 2:
|
|
|
|
|
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
|
2020-08-22 19:42:14 +08:00
|
|
|
|
if det and rec:
|
2021-05-26 17:34:47 +08:00
|
|
|
|
dt_boxes, rec_res = self.__call__(img, cls)
|
2020-08-22 19:42:14 +08:00
|
|
|
|
return [[box.tolist(), res] for box, res in zip(dt_boxes, rec_res)]
|
|
|
|
|
elif det and not rec:
|
|
|
|
|
dt_boxes, elapse = self.text_detector(img)
|
|
|
|
|
if dt_boxes is None:
|
|
|
|
|
return None
|
|
|
|
|
return [box.tolist() for box in dt_boxes]
|
|
|
|
|
else:
|
|
|
|
|
if not isinstance(img, list):
|
|
|
|
|
img = [img]
|
2021-05-26 17:34:47 +08:00
|
|
|
|
if self.use_angle_cls and cls:
|
2020-12-07 19:10:19 +08:00
|
|
|
|
img, cls_res, elapse = self.text_classifier(img)
|
|
|
|
|
if not rec:
|
|
|
|
|
return cls_res
|
2020-08-22 19:42:14 +08:00
|
|
|
|
rec_res, elapse = self.text_recognizer(img)
|
|
|
|
|
return rec_res
|
2020-08-24 11:30:56 +08:00
|
|
|
|
|
|
|
|
|
|
2021-08-02 15:28:07 +08:00
|
|
|
|
class PPStructure(OCRSystem):
|
|
|
|
|
def __init__(self, **kwargs):
|
|
|
|
|
params = parse_args(mMain=False)
|
|
|
|
|
params.__dict__.update(**kwargs)
|
|
|
|
|
if not params.show_log:
|
|
|
|
|
logger.setLevel(logging.INFO)
|
2021-08-02 17:04:53 +08:00
|
|
|
|
lang, det_lang = parse_lang(params.lang)
|
|
|
|
|
|
2021-08-02 15:28:07 +08:00
|
|
|
|
# init model dir
|
2021-09-06 18:33:21 +08:00
|
|
|
|
det_model_config = get_model_config(params.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)
|
|
|
|
|
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')
|
|
|
|
|
params.table_model_dir, table_url = confirm_model_dir_url(
|
|
|
|
|
params.table_model_dir,
|
|
|
|
|
os.path.join(BASE_DIR, VERSION, 'ocr', 'table'),
|
|
|
|
|
table_model_config['url'])
|
2021-08-02 15:28:07 +08:00
|
|
|
|
# download model
|
|
|
|
|
maybe_download(params.det_model_dir, det_url)
|
|
|
|
|
maybe_download(params.rec_model_dir, rec_url)
|
|
|
|
|
maybe_download(params.table_model_dir, table_url)
|
|
|
|
|
|
|
|
|
|
if params.rec_char_dict_path is None:
|
2021-09-06 18:33:21 +08:00
|
|
|
|
params.rec_char_dict_path = str(
|
|
|
|
|
Path(__file__).parent / rec_model_config['dict_path'])
|
2021-08-02 15:28:07 +08:00
|
|
|
|
if params.table_char_dict_path is None:
|
2021-09-06 18:33:21 +08:00
|
|
|
|
params.table_char_dict_path = str(
|
|
|
|
|
Path(__file__).parent / table_model_config['dict_path'])
|
2021-08-02 15:28:07 +08:00
|
|
|
|
|
|
|
|
|
print(params)
|
|
|
|
|
super().__init__(params)
|
|
|
|
|
|
|
|
|
|
def __call__(self, img):
|
|
|
|
|
if isinstance(img, str):
|
|
|
|
|
# download net image
|
|
|
|
|
if img.startswith('http'):
|
|
|
|
|
download_with_progressbar(img, 'tmp.jpg')
|
|
|
|
|
img = 'tmp.jpg'
|
|
|
|
|
image_file = img
|
|
|
|
|
img, flag = check_and_read_gif(image_file)
|
|
|
|
|
if not flag:
|
|
|
|
|
with open(image_file, 'rb') as f:
|
|
|
|
|
np_arr = np.frombuffer(f.read(), dtype=np.uint8)
|
|
|
|
|
img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
|
|
|
|
|
if img is None:
|
|
|
|
|
logger.error("error in loading image:{}".format(image_file))
|
|
|
|
|
return None
|
|
|
|
|
if isinstance(img, np.ndarray) and len(img.shape) == 2:
|
|
|
|
|
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
|
|
|
|
|
|
|
|
|
|
res = super().__call__(img)
|
|
|
|
|
return res
|
|
|
|
|
|
|
|
|
|
|
2020-08-24 11:30:56 +08:00
|
|
|
|
def main():
|
2020-12-07 19:10:19 +08:00
|
|
|
|
# for cmd
|
2021-06-02 20:10:59 +08:00
|
|
|
|
args = parse_args(mMain=True)
|
2020-12-07 19:10:19 +08:00
|
|
|
|
image_dir = args.image_dir
|
2021-06-10 14:47:23 +08:00
|
|
|
|
if is_link(image_dir):
|
2020-12-07 19:10:19 +08:00
|
|
|
|
download_with_progressbar(image_dir, 'tmp.jpg')
|
|
|
|
|
image_file_list = ['tmp.jpg']
|
|
|
|
|
else:
|
|
|
|
|
image_file_list = get_image_file_list(args.image_dir)
|
2020-08-24 11:30:56 +08:00
|
|
|
|
if len(image_file_list) == 0:
|
|
|
|
|
logger.error('no images find in {}'.format(args.image_dir))
|
|
|
|
|
return
|
2021-08-02 17:04:53 +08:00
|
|
|
|
if args.type == 'ocr':
|
2021-08-02 15:28:07 +08:00
|
|
|
|
engine = PaddleOCR(**(args.__dict__))
|
2021-08-02 17:04:53 +08:00
|
|
|
|
elif args.type == 'structure':
|
2021-08-02 15:28:07 +08:00
|
|
|
|
engine = PPStructure(**(args.__dict__))
|
|
|
|
|
else:
|
|
|
|
|
raise NotImplementedError
|
2020-12-07 19:10:19 +08:00
|
|
|
|
|
2020-08-24 11:30:56 +08:00
|
|
|
|
for img_path in image_file_list:
|
2021-08-02 15:28:07 +08:00
|
|
|
|
img_name = os.path.basename(img_path).split('.')[0]
|
2020-12-07 19:10:19 +08:00
|
|
|
|
logger.info('{}{}{}'.format('*' * 10, img_path, '*' * 10))
|
2021-08-02 15:28:07 +08:00
|
|
|
|
if args.type == 'ocr':
|
|
|
|
|
result = engine.ocr(img_path,
|
2021-08-02 17:04:53 +08:00
|
|
|
|
det=args.det,
|
|
|
|
|
rec=args.rec,
|
|
|
|
|
cls=args.use_angle_cls)
|
2021-08-02 15:28:07 +08:00
|
|
|
|
if result is not None:
|
|
|
|
|
for line in result:
|
|
|
|
|
logger.info(line)
|
|
|
|
|
elif args.type == 'structure':
|
|
|
|
|
result = engine(img_path)
|
2021-08-02 17:22:31 +08:00
|
|
|
|
save_structure_res(result, args.output, img_name)
|
|
|
|
|
|
2021-08-02 15:28:07 +08:00
|
|
|
|
for item in result:
|
2021-08-02 17:22:31 +08:00
|
|
|
|
item.pop('img')
|
2021-08-03 11:22:49 +08:00
|
|
|
|
logger.info(item)
|