MissPenguin
GitHub
commit
e93735a2ef
|
|
@ -1,7 +1,7 @@
|
||||||
include LICENSE.txt
|
include LICENSE
|
||||||
include README.md
|
include README.md
|
||||||
|
|
||||||
recursive-include ppocr/utils *.txt utility.py logging.py
|
recursive-include ppocr/utils *.txt utility.py logging.py network.py
|
||||||
recursive-include ppocr/data/ *.py
|
recursive-include ppocr/data/ *.py
|
||||||
recursive-include ppocr/postprocess *.py
|
recursive-include ppocr/postprocess *.py
|
||||||
recursive-include tools/infer *.py
|
recursive-include tools/infer *.py
|
||||||
|
|
|
||||||
|
|
@ -355,3 +355,4 @@ im_show.save('result.jpg')
|
||||||
| det | 前向时使用启动检测 | TRUE |
|
| det | 前向时使用启动检测 | TRUE |
|
||||||
| rec | 前向时是否启动识别 | TRUE |
|
| rec | 前向时是否启动识别 | TRUE |
|
||||||
| cls | 前向时是否启动分类 (命令行模式下使用use_angle_cls控制前向是否启动分类) | FALSE |
|
| cls | 前向时是否启动分类 (命令行模式下使用use_angle_cls控制前向是否启动分类) | FALSE |
|
||||||
|
| show_log | 是否打印det和rec等信息 | FALSE |
|
||||||
|
|
|
||||||
|
|
@ -362,3 +362,5 @@ im_show.save('result.jpg')
|
||||||
| det | Enable detction when `ppocr.ocr` func exec | TRUE |
|
| det | Enable detction when `ppocr.ocr` func exec | TRUE |
|
||||||
| rec | Enable recognition when `ppocr.ocr` func exec | TRUE |
|
| 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 |
|
| 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 |
|
||||||
Binary file not shown.
|
After Width: | Height: | Size: 263 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 116 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 26 KiB |
138
paddleocr.py
138
paddleocr.py
|
|
@ -19,17 +19,16 @@ __dir__ = os.path.dirname(__file__)
|
||||||
sys.path.append(os.path.join(__dir__, ''))
|
sys.path.append(os.path.join(__dir__, ''))
|
||||||
|
|
||||||
import cv2
|
import cv2
|
||||||
|
import logging
|
||||||
import numpy as np
|
import numpy as np
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
import tarfile
|
|
||||||
import requests
|
|
||||||
from tqdm import tqdm
|
|
||||||
|
|
||||||
from tools.infer import predict_system
|
from tools.infer import predict_system
|
||||||
from ppocr.utils.logging import get_logger
|
from ppocr.utils.logging import get_logger
|
||||||
|
|
||||||
logger = get_logger()
|
logger = get_logger()
|
||||||
from ppocr.utils.utility import check_and_read_gif, get_image_file_list
|
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, init_args, str2bool
|
from tools.infer.utility import draw_ocr, init_args, str2bool
|
||||||
|
|
||||||
__all__ = ['PaddleOCR']
|
__all__ = ['PaddleOCR']
|
||||||
|
|
@ -37,84 +36,84 @@ __all__ = ['PaddleOCR']
|
||||||
model_urls = {
|
model_urls = {
|
||||||
'det': {
|
'det': {
|
||||||
'ch':
|
'ch':
|
||||||
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_det_infer.tar',
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_det_infer.tar',
|
||||||
'en':
|
'en':
|
||||||
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/en_ppocr_mobile_v2.0_det_infer.tar'
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/en_ppocr_mobile_v2.0_det_infer.tar'
|
||||||
},
|
},
|
||||||
'rec': {
|
'rec': {
|
||||||
'ch': {
|
'ch': {
|
||||||
'url':
|
'url':
|
||||||
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_rec_infer.tar',
|
'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'
|
'dict_path': './ppocr/utils/ppocr_keys_v1.txt'
|
||||||
},
|
},
|
||||||
'en': {
|
'en': {
|
||||||
'url':
|
'url':
|
||||||
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/en_number_mobile_v2.0_rec_infer.tar',
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/en_number_mobile_v2.0_rec_infer.tar',
|
||||||
'dict_path': './ppocr/utils/en_dict.txt'
|
'dict_path': './ppocr/utils/en_dict.txt'
|
||||||
},
|
},
|
||||||
'french': {
|
'french': {
|
||||||
'url':
|
'url':
|
||||||
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/french_mobile_v2.0_rec_infer.tar',
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/french_mobile_v2.0_rec_infer.tar',
|
||||||
'dict_path': './ppocr/utils/dict/french_dict.txt'
|
'dict_path': './ppocr/utils/dict/french_dict.txt'
|
||||||
},
|
},
|
||||||
'german': {
|
'german': {
|
||||||
'url':
|
'url':
|
||||||
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/german_mobile_v2.0_rec_infer.tar',
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/german_mobile_v2.0_rec_infer.tar',
|
||||||
'dict_path': './ppocr/utils/dict/german_dict.txt'
|
'dict_path': './ppocr/utils/dict/german_dict.txt'
|
||||||
},
|
},
|
||||||
'korean': {
|
'korean': {
|
||||||
'url':
|
'url':
|
||||||
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/korean_mobile_v2.0_rec_infer.tar',
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/korean_mobile_v2.0_rec_infer.tar',
|
||||||
'dict_path': './ppocr/utils/dict/korean_dict.txt'
|
'dict_path': './ppocr/utils/dict/korean_dict.txt'
|
||||||
},
|
},
|
||||||
'japan': {
|
'japan': {
|
||||||
'url':
|
'url':
|
||||||
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/japan_mobile_v2.0_rec_infer.tar',
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/japan_mobile_v2.0_rec_infer.tar',
|
||||||
'dict_path': './ppocr/utils/dict/japan_dict.txt'
|
'dict_path': './ppocr/utils/dict/japan_dict.txt'
|
||||||
},
|
},
|
||||||
'chinese_cht': {
|
'chinese_cht': {
|
||||||
'url':
|
'url':
|
||||||
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/chinese_cht_mobile_v2.0_rec_infer.tar',
|
'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'
|
'dict_path': './ppocr/utils/dict/chinese_cht_dict.txt'
|
||||||
},
|
},
|
||||||
'ta': {
|
'ta': {
|
||||||
'url':
|
'url':
|
||||||
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/ta_mobile_v2.0_rec_infer.tar',
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/ta_mobile_v2.0_rec_infer.tar',
|
||||||
'dict_path': './ppocr/utils/dict/ta_dict.txt'
|
'dict_path': './ppocr/utils/dict/ta_dict.txt'
|
||||||
},
|
},
|
||||||
'te': {
|
'te': {
|
||||||
'url':
|
'url':
|
||||||
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/te_mobile_v2.0_rec_infer.tar',
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/te_mobile_v2.0_rec_infer.tar',
|
||||||
'dict_path': './ppocr/utils/dict/te_dict.txt'
|
'dict_path': './ppocr/utils/dict/te_dict.txt'
|
||||||
},
|
},
|
||||||
'ka': {
|
'ka': {
|
||||||
'url':
|
'url':
|
||||||
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/ka_mobile_v2.0_rec_infer.tar',
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/ka_mobile_v2.0_rec_infer.tar',
|
||||||
'dict_path': './ppocr/utils/dict/ka_dict.txt'
|
'dict_path': './ppocr/utils/dict/ka_dict.txt'
|
||||||
},
|
},
|
||||||
'latin': {
|
'latin': {
|
||||||
'url':
|
'url':
|
||||||
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/latin_ppocr_mobile_v2.0_rec_infer.tar',
|
'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'
|
'dict_path': './ppocr/utils/dict/latin_dict.txt'
|
||||||
},
|
},
|
||||||
'arabic': {
|
'arabic': {
|
||||||
'url':
|
'url':
|
||||||
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/arabic_ppocr_mobile_v2.0_rec_infer.tar',
|
'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'
|
'dict_path': './ppocr/utils/dict/arabic_dict.txt'
|
||||||
},
|
},
|
||||||
'cyrillic': {
|
'cyrillic': {
|
||||||
'url':
|
'url':
|
||||||
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/cyrillic_ppocr_mobile_v2.0_rec_infer.tar',
|
'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'
|
'dict_path': './ppocr/utils/dict/cyrillic_dict.txt'
|
||||||
},
|
},
|
||||||
'devanagari': {
|
'devanagari': {
|
||||||
'url':
|
'url':
|
||||||
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/devanagari_ppocr_mobile_v2.0_rec_infer.tar',
|
'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'
|
'dict_path': './ppocr/utils/dict/devanagari_dict.txt'
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
'cls':
|
'cls':
|
||||||
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_cls_infer.tar'
|
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_cls_infer.tar'
|
||||||
}
|
}
|
||||||
|
|
||||||
SUPPORT_DET_MODEL = ['DB']
|
SUPPORT_DET_MODEL = ['DB']
|
||||||
|
|
@ -123,50 +122,6 @@ SUPPORT_REC_MODEL = ['CRNN']
|
||||||
BASE_DIR = os.path.expanduser("~/.paddleocr/")
|
BASE_DIR = os.path.expanduser("~/.paddleocr/")
|
||||||
|
|
||||||
|
|
||||||
def download_with_progressbar(url, save_path):
|
|
||||||
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:
|
|
||||||
logger.error("Something went wrong while downloading models")
|
|
||||||
sys.exit(0)
|
|
||||||
|
|
||||||
|
|
||||||
def maybe_download(model_storage_directory, url):
|
|
||||||
# using custom model
|
|
||||||
tar_file_name_list = [
|
|
||||||
'inference.pdiparams', 'inference.pdiparams.info', 'inference.pdmodel'
|
|
||||||
]
|
|
||||||
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')):
|
|
||||||
tmp_path = os.path.join(model_storage_directory, url.split('/')[-1])
|
|
||||||
print('download {} to {}'.format(url, tmp_path))
|
|
||||||
os.makedirs(model_storage_directory, exist_ok=True)
|
|
||||||
download_with_progressbar(url, tmp_path)
|
|
||||||
with tarfile.open(tmp_path, 'r') as tarObj:
|
|
||||||
for member in tarObj.getmembers():
|
|
||||||
filename = None
|
|
||||||
for tar_file_name in tar_file_name_list:
|
|
||||||
if tar_file_name in member.name:
|
|
||||||
filename = tar_file_name
|
|
||||||
if filename is None:
|
|
||||||
continue
|
|
||||||
file = tarObj.extractfile(member)
|
|
||||||
with open(
|
|
||||||
os.path.join(model_storage_directory, filename),
|
|
||||||
'wb') as f:
|
|
||||||
f.write(file.read())
|
|
||||||
os.remove(tmp_path)
|
|
||||||
|
|
||||||
|
|
||||||
def parse_args(mMain=True):
|
def parse_args(mMain=True):
|
||||||
import argparse
|
import argparse
|
||||||
parser = init_args()
|
parser = init_args()
|
||||||
|
|
@ -194,10 +149,12 @@ class PaddleOCR(predict_system.TextSystem):
|
||||||
args:
|
args:
|
||||||
**kwargs: other params show in paddleocr --help
|
**kwargs: other params show in paddleocr --help
|
||||||
"""
|
"""
|
||||||
postprocess_params = parse_args(mMain=False)
|
params = parse_args(mMain=False)
|
||||||
postprocess_params.__dict__.update(**kwargs)
|
params.__dict__.update(**kwargs)
|
||||||
self.use_angle_cls = postprocess_params.use_angle_cls
|
if not params.show_log:
|
||||||
lang = postprocess_params.lang
|
logger.setLevel(logging.INFO)
|
||||||
|
self.use_angle_cls = params.use_angle_cls
|
||||||
|
lang = params.lang
|
||||||
latin_lang = [
|
latin_lang = [
|
||||||
'af', 'az', 'bs', 'cs', 'cy', 'da', 'de', 'es', 'et', 'fr', 'ga',
|
'af', 'az', 'bs', 'cs', 'cy', 'da', 'de', 'es', 'et', 'fr', 'ga',
|
||||||
'hr', 'hu', 'id', 'is', 'it', 'ku', 'la', 'lt', 'lv', 'mi', 'ms',
|
'hr', 'hu', 'id', 'is', 'it', 'ku', 'la', 'lt', 'lv', 'mi', 'ms',
|
||||||
|
|
@ -223,46 +180,45 @@ class PaddleOCR(predict_system.TextSystem):
|
||||||
lang = "devanagari"
|
lang = "devanagari"
|
||||||
assert lang in model_urls[
|
assert lang in model_urls[
|
||||||
'rec'], 'param lang must in {}, but got {}'.format(
|
'rec'], 'param lang must in {}, but got {}'.format(
|
||||||
model_urls['rec'].keys(), lang)
|
model_urls['rec'].keys(), lang)
|
||||||
if lang == "ch":
|
if lang == "ch":
|
||||||
det_lang = "ch"
|
det_lang = "ch"
|
||||||
else:
|
else:
|
||||||
det_lang = "en"
|
det_lang = "en"
|
||||||
use_inner_dict = False
|
use_inner_dict = False
|
||||||
if postprocess_params.rec_char_dict_path is None:
|
if params.rec_char_dict_path is None:
|
||||||
use_inner_dict = True
|
use_inner_dict = True
|
||||||
postprocess_params.rec_char_dict_path = model_urls['rec'][lang][
|
params.rec_char_dict_path = model_urls['rec'][lang][
|
||||||
'dict_path']
|
'dict_path']
|
||||||
|
|
||||||
# init model dir
|
# init model dir
|
||||||
if postprocess_params.det_model_dir is None:
|
params.det_model_dir, det_url = confirm_model_dir_url(params.det_model_dir,
|
||||||
postprocess_params.det_model_dir = os.path.join(BASE_DIR, VERSION,
|
os.path.join(BASE_DIR, VERSION, 'det', det_lang),
|
||||||
'det', det_lang)
|
model_urls['det'][det_lang])
|
||||||
if postprocess_params.rec_model_dir is None:
|
params.rec_model_dir, rec_url = confirm_model_dir_url(params.rec_model_dir,
|
||||||
postprocess_params.rec_model_dir = os.path.join(BASE_DIR, VERSION,
|
os.path.join(BASE_DIR, VERSION, 'rec', lang),
|
||||||
'rec', lang)
|
model_urls['rec'][lang]['url'])
|
||||||
if postprocess_params.cls_model_dir is None:
|
params.cls_model_dir, cls_url = confirm_model_dir_url(params.cls_model_dir,
|
||||||
postprocess_params.cls_model_dir = os.path.join(BASE_DIR, 'cls')
|
os.path.join(BASE_DIR, VERSION, 'cls'),
|
||||||
print(postprocess_params)
|
model_urls['cls'])
|
||||||
# download model
|
# download model
|
||||||
maybe_download(postprocess_params.det_model_dir,
|
maybe_download(params.det_model_dir, det_url)
|
||||||
model_urls['det'][det_lang])
|
maybe_download(params.rec_model_dir, rec_url)
|
||||||
maybe_download(postprocess_params.rec_model_dir,
|
maybe_download(params.cls_model_dir, cls_url)
|
||||||
model_urls['rec'][lang]['url'])
|
|
||||||
maybe_download(postprocess_params.cls_model_dir, model_urls['cls'])
|
|
||||||
|
|
||||||
if postprocess_params.det_algorithm not in SUPPORT_DET_MODEL:
|
if params.det_algorithm not in SUPPORT_DET_MODEL:
|
||||||
logger.error('det_algorithm must in {}'.format(SUPPORT_DET_MODEL))
|
logger.error('det_algorithm must in {}'.format(SUPPORT_DET_MODEL))
|
||||||
sys.exit(0)
|
sys.exit(0)
|
||||||
if postprocess_params.rec_algorithm not in SUPPORT_REC_MODEL:
|
if params.rec_algorithm not in SUPPORT_REC_MODEL:
|
||||||
logger.error('rec_algorithm must in {}'.format(SUPPORT_REC_MODEL))
|
logger.error('rec_algorithm must in {}'.format(SUPPORT_REC_MODEL))
|
||||||
sys.exit(0)
|
sys.exit(0)
|
||||||
if use_inner_dict:
|
if use_inner_dict:
|
||||||
postprocess_params.rec_char_dict_path = str(
|
params.rec_char_dict_path = str(
|
||||||
Path(__file__).parent / postprocess_params.rec_char_dict_path)
|
Path(__file__).parent / params.rec_char_dict_path)
|
||||||
|
|
||||||
|
print(params)
|
||||||
# init det_model and rec_model
|
# init det_model and rec_model
|
||||||
super().__init__(postprocess_params)
|
super().__init__(params)
|
||||||
|
|
||||||
def ocr(self, img, det=True, rec=True, cls=True):
|
def ocr(self, img, det=True, rec=True, cls=True):
|
||||||
"""
|
"""
|
||||||
|
|
@ -320,7 +276,7 @@ def main():
|
||||||
# for cmd
|
# for cmd
|
||||||
args = parse_args(mMain=True)
|
args = parse_args(mMain=True)
|
||||||
image_dir = args.image_dir
|
image_dir = args.image_dir
|
||||||
if image_dir.startswith('http'):
|
if is_link(image_dir):
|
||||||
download_with_progressbar(image_dir, 'tmp.jpg')
|
download_with_progressbar(image_dir, 'tmp.jpg')
|
||||||
image_file_list = ['tmp.jpg']
|
image_file_list = ['tmp.jpg']
|
||||||
else:
|
else:
|
||||||
|
|
|
||||||
|
|
@ -29,6 +29,7 @@ from .label_ops import *
|
||||||
from .east_process import *
|
from .east_process import *
|
||||||
from .sast_process import *
|
from .sast_process import *
|
||||||
from .pg_process import *
|
from .pg_process import *
|
||||||
|
from .gen_table_mask import *
|
||||||
|
|
||||||
|
|
||||||
def transform(data, ops=None):
|
def transform(data, ops=None):
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,244 @@
|
||||||
|
"""
|
||||||
|
# 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.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import absolute_import
|
||||||
|
from __future__ import division
|
||||||
|
from __future__ import print_function
|
||||||
|
from __future__ import unicode_literals
|
||||||
|
|
||||||
|
import sys
|
||||||
|
import six
|
||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
|
||||||
|
class GenTableMask(object):
|
||||||
|
""" gen table mask """
|
||||||
|
|
||||||
|
def __init__(self, shrink_h_max, shrink_w_max, mask_type=0, **kwargs):
|
||||||
|
self.shrink_h_max = 5
|
||||||
|
self.shrink_w_max = 5
|
||||||
|
self.mask_type = mask_type
|
||||||
|
|
||||||
|
def projection(self, erosion, h, w, spilt_threshold=0):
|
||||||
|
# 水平投影
|
||||||
|
projection_map = np.ones_like(erosion)
|
||||||
|
project_val_array = [0 for _ in range(0, h)]
|
||||||
|
|
||||||
|
for j in range(0, h):
|
||||||
|
for i in range(0, w):
|
||||||
|
if erosion[j, i] == 255:
|
||||||
|
project_val_array[j] += 1
|
||||||
|
# 根据数组，获取切割点
|
||||||
|
start_idx = 0 # 记录进入字符区的索引
|
||||||
|
end_idx = 0 # 记录进入空白区域的索引
|
||||||
|
in_text = False # 是否遍历到了字符区内
|
||||||
|
box_list = []
|
||||||
|
for i in range(len(project_val_array)):
|
||||||
|
if in_text == False and project_val_array[i] > spilt_threshold: # 进入字符区了
|
||||||
|
in_text = True
|
||||||
|
start_idx = i
|
||||||
|
elif project_val_array[i] <= spilt_threshold and in_text == True: # 进入空白区了
|
||||||
|
end_idx = i
|
||||||
|
in_text = False
|
||||||
|
if end_idx - start_idx <= 2:
|
||||||
|
continue
|
||||||
|
box_list.append((start_idx, end_idx + 1))
|
||||||
|
|
||||||
|
if in_text:
|
||||||
|
box_list.append((start_idx, h - 1))
|
||||||
|
# 绘制投影直方图
|
||||||
|
for j in range(0, h):
|
||||||
|
for i in range(0, project_val_array[j]):
|
||||||
|
projection_map[j, i] = 0
|
||||||
|
return box_list, projection_map
|
||||||
|
|
||||||
|
def projection_cx(self, box_img):
|
||||||
|
box_gray_img = cv2.cvtColor(box_img, cv2.COLOR_BGR2GRAY)
|
||||||
|
h, w = box_gray_img.shape
|
||||||
|
# 灰度图片进行二值化处理
|
||||||
|
ret, thresh1 = cv2.threshold(box_gray_img, 200, 255, cv2.THRESH_BINARY_INV)
|
||||||
|
# 纵向腐蚀
|
||||||
|
if h < w:
|
||||||
|
kernel = np.ones((2, 1), np.uint8)
|
||||||
|
erode = cv2.erode(thresh1, kernel, iterations=1)
|
||||||
|
else:
|
||||||
|
erode = thresh1
|
||||||
|
# 水平膨胀
|
||||||
|
kernel = np.ones((1, 5), np.uint8)
|
||||||
|
erosion = cv2.dilate(erode, kernel, iterations=1)
|
||||||
|
# 水平投影
|
||||||
|
projection_map = np.ones_like(erosion)
|
||||||
|
project_val_array = [0 for _ in range(0, h)]
|
||||||
|
|
||||||
|
for j in range(0, h):
|
||||||
|
for i in range(0, w):
|
||||||
|
if erosion[j, i] == 255:
|
||||||
|
project_val_array[j] += 1
|
||||||
|
# 根据数组，获取切割点
|
||||||
|
start_idx = 0 # 记录进入字符区的索引
|
||||||
|
end_idx = 0 # 记录进入空白区域的索引
|
||||||
|
in_text = False # 是否遍历到了字符区内
|
||||||
|
box_list = []
|
||||||
|
spilt_threshold = 0
|
||||||
|
for i in range(len(project_val_array)):
|
||||||
|
if in_text == False and project_val_array[i] > spilt_threshold: # 进入字符区了
|
||||||
|
in_text = True
|
||||||
|
start_idx = i
|
||||||
|
elif project_val_array[i] <= spilt_threshold and in_text == True: # 进入空白区了
|
||||||
|
end_idx = i
|
||||||
|
in_text = False
|
||||||
|
if end_idx - start_idx <= 2:
|
||||||
|
continue
|
||||||
|
box_list.append((start_idx, end_idx + 1))
|
||||||
|
|
||||||
|
if in_text:
|
||||||
|
box_list.append((start_idx, h - 1))
|
||||||
|
# 绘制投影直方图
|
||||||
|
for j in range(0, h):
|
||||||
|
for i in range(0, project_val_array[j]):
|
||||||
|
projection_map[j, i] = 0
|
||||||
|
split_bbox_list = []
|
||||||
|
if len(box_list) > 1:
|
||||||
|
for i, (h_start, h_end) in enumerate(box_list):
|
||||||
|
if i == 0:
|
||||||
|
h_start = 0
|
||||||
|
if i == len(box_list):
|
||||||
|
h_end = h
|
||||||
|
word_img = erosion[h_start:h_end + 1, :]
|
||||||
|
word_h, word_w = word_img.shape
|
||||||
|
w_split_list, w_projection_map = self.projection(word_img.T, word_w, word_h)
|
||||||
|
w_start, w_end = w_split_list[0][0], w_split_list[-1][1]
|
||||||
|
if h_start > 0:
|
||||||
|
h_start -= 1
|
||||||
|
h_end += 1
|
||||||
|
word_img = box_img[h_start:h_end + 1:, w_start:w_end + 1, :]
|
||||||
|
split_bbox_list.append([w_start, h_start, w_end, h_end])
|
||||||
|
else:
|
||||||
|
split_bbox_list.append([0, 0, w, h])
|
||||||
|
return split_bbox_list
|
||||||
|
|
||||||
|
def shrink_bbox(self, bbox):
|
||||||
|
left, top, right, bottom = bbox
|
||||||
|
sh_h = min(max(int((bottom - top) * 0.1), 1), self.shrink_h_max)
|
||||||
|
sh_w = min(max(int((right - left) * 0.1), 1), self.shrink_w_max)
|
||||||
|
left_new = left + sh_w
|
||||||
|
right_new = right - sh_w
|
||||||
|
top_new = top + sh_h
|
||||||
|
bottom_new = bottom - sh_h
|
||||||
|
if left_new >= right_new:
|
||||||
|
left_new = left
|
||||||
|
right_new = right
|
||||||
|
if top_new >= bottom_new:
|
||||||
|
top_new = top
|
||||||
|
bottom_new = bottom
|
||||||
|
return [left_new, top_new, right_new, bottom_new]
|
||||||
|
|
||||||
|
def __call__(self, data):
|
||||||
|
img = data['image']
|
||||||
|
cells = data['cells']
|
||||||
|
height, width = img.shape[0:2]
|
||||||
|
if self.mask_type == 1:
|
||||||
|
mask_img = np.zeros((height, width), dtype=np.float32)
|
||||||
|
else:
|
||||||
|
mask_img = np.zeros((height, width, 3), dtype=np.float32)
|
||||||
|
cell_num = len(cells)
|
||||||
|
for cno in range(cell_num):
|
||||||
|
if "bbox" in cells[cno]:
|
||||||
|
bbox = cells[cno]['bbox']
|
||||||
|
left, top, right, bottom = bbox
|
||||||
|
box_img = img[top:bottom, left:right, :].copy()
|
||||||
|
split_bbox_list = self.projection_cx(box_img)
|
||||||
|
for sno in range(len(split_bbox_list)):
|
||||||
|
split_bbox_list[sno][0] += left
|
||||||
|
split_bbox_list[sno][1] += top
|
||||||
|
split_bbox_list[sno][2] += left
|
||||||
|
split_bbox_list[sno][3] += top
|
||||||
|
|
||||||
|
for sno in range(len(split_bbox_list)):
|
||||||
|
left, top, right, bottom = split_bbox_list[sno]
|
||||||
|
left, top, right, bottom = self.shrink_bbox([left, top, right, bottom])
|
||||||
|
if self.mask_type == 1:
|
||||||
|
mask_img[top:bottom, left:right] = 1.0
|
||||||
|
data['mask_img'] = mask_img
|
||||||
|
else:
|
||||||
|
mask_img[top:bottom, left:right, :] = (255, 255, 255)
|
||||||
|
data['image'] = mask_img
|
||||||
|
return data
|
||||||
|
|
||||||
|
class ResizeTableImage(object):
|
||||||
|
def __init__(self, max_len, **kwargs):
|
||||||
|
super(ResizeTableImage, self).__init__()
|
||||||
|
self.max_len = max_len
|
||||||
|
|
||||||
|
def get_img_bbox(self, cells):
|
||||||
|
bbox_list = []
|
||||||
|
if len(cells) == 0:
|
||||||
|
return bbox_list
|
||||||
|
cell_num = len(cells)
|
||||||
|
for cno in range(cell_num):
|
||||||
|
if "bbox" in cells[cno]:
|
||||||
|
bbox = cells[cno]['bbox']
|
||||||
|
bbox_list.append(bbox)
|
||||||
|
return bbox_list
|
||||||
|
|
||||||
|
def resize_img_table(self, img, bbox_list, max_len):
|
||||||
|
height, width = img.shape[0:2]
|
||||||
|
ratio = max_len / (max(height, width) * 1.0)
|
||||||
|
resize_h = int(height * ratio)
|
||||||
|
resize_w = int(width * ratio)
|
||||||
|
img_new = cv2.resize(img, (resize_w, resize_h))
|
||||||
|
bbox_list_new = []
|
||||||
|
for bno in range(len(bbox_list)):
|
||||||
|
left, top, right, bottom = bbox_list[bno].copy()
|
||||||
|
left = int(left * ratio)
|
||||||
|
top = int(top * ratio)
|
||||||
|
right = int(right * ratio)
|
||||||
|
bottom = int(bottom * ratio)
|
||||||
|
bbox_list_new.append([left, top, right, bottom])
|
||||||
|
return img_new, bbox_list_new
|
||||||
|
|
||||||
|
def __call__(self, data):
|
||||||
|
img = data['image']
|
||||||
|
if 'cells' not in data:
|
||||||
|
cells = []
|
||||||
|
else:
|
||||||
|
cells = data['cells']
|
||||||
|
bbox_list = self.get_img_bbox(cells)
|
||||||
|
img_new, bbox_list_new = self.resize_img_table(img, bbox_list, self.max_len)
|
||||||
|
data['image'] = img_new
|
||||||
|
cell_num = len(cells)
|
||||||
|
bno = 0
|
||||||
|
for cno in range(cell_num):
|
||||||
|
if "bbox" in data['cells'][cno]:
|
||||||
|
data['cells'][cno]['bbox'] = bbox_list_new[bno]
|
||||||
|
bno += 1
|
||||||
|
data['max_len'] = self.max_len
|
||||||
|
return data
|
||||||
|
|
||||||
|
class PaddingTableImage(object):
|
||||||
|
def __init__(self, **kwargs):
|
||||||
|
super(PaddingTableImage, self).__init__()
|
||||||
|
|
||||||
|
def __call__(self, data):
|
||||||
|
img = data['image']
|
||||||
|
max_len = data['max_len']
|
||||||
|
padding_img = np.zeros((max_len, max_len, 3), dtype=np.float32)
|
||||||
|
height, width = img.shape[0:2]
|
||||||
|
padding_img[0:height, 0:width, :] = img.copy()
|
||||||
|
data['image'] = padding_img
|
||||||
|
return data
|
||||||
|
|
||||||
|
|
@ -81,7 +81,7 @@ class NormalizeImage(object):
|
||||||
assert isinstance(img,
|
assert isinstance(img,
|
||||||
np.ndarray), "invalid input 'img' in NormalizeImage"
|
np.ndarray), "invalid input 'img' in NormalizeImage"
|
||||||
data['image'] = (
|
data['image'] = (
|
||||||
img.astype('float32') * self.scale - self.mean) / self.std
|
img.astype('float32') * self.scale - self.mean) / self.std
|
||||||
return data
|
return data
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -163,7 +163,7 @@ class DetResizeForTest(object):
|
||||||
img, (ratio_h, ratio_w)
|
img, (ratio_h, ratio_w)
|
||||||
"""
|
"""
|
||||||
limit_side_len = self.limit_side_len
|
limit_side_len = self.limit_side_len
|
||||||
h, w, _ = img.shape
|
h, w, c = img.shape
|
||||||
|
|
||||||
# limit the max side
|
# limit the max side
|
||||||
if self.limit_type == 'max':
|
if self.limit_type == 'max':
|
||||||
|
|
@ -174,7 +174,7 @@ class DetResizeForTest(object):
|
||||||
ratio = float(limit_side_len) / w
|
ratio = float(limit_side_len) / w
|
||||||
else:
|
else:
|
||||||
ratio = 1.
|
ratio = 1.
|
||||||
else:
|
elif self.limit_type == 'min':
|
||||||
if min(h, w) < limit_side_len:
|
if min(h, w) < limit_side_len:
|
||||||
if h < w:
|
if h < w:
|
||||||
ratio = float(limit_side_len) / h
|
ratio = float(limit_side_len) / h
|
||||||
|
|
@ -182,6 +182,10 @@ class DetResizeForTest(object):
|
||||||
ratio = float(limit_side_len) / w
|
ratio = float(limit_side_len) / w
|
||||||
else:
|
else:
|
||||||
ratio = 1.
|
ratio = 1.
|
||||||
|
elif self.limit_type == 'resize_long':
|
||||||
|
ratio = float(limit_side_len) / max(h,w)
|
||||||
|
else:
|
||||||
|
raise Exception('not support limit type, image ')
|
||||||
resize_h = int(h * ratio)
|
resize_h = int(h * ratio)
|
||||||
resize_w = int(w * ratio)
|
resize_w = int(w * ratio)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -24,7 +24,8 @@ __all__ = ['build_post_process']
|
||||||
from .db_postprocess import DBPostProcess
|
from .db_postprocess import DBPostProcess
|
||||||
from .east_postprocess import EASTPostProcess
|
from .east_postprocess import EASTPostProcess
|
||||||
from .sast_postprocess import SASTPostProcess
|
from .sast_postprocess import SASTPostProcess
|
||||||
from .rec_postprocess import CTCLabelDecode, AttnLabelDecode, SRNLabelDecode, DistillationCTCLabelDecode
|
from .rec_postprocess import CTCLabelDecode, AttnLabelDecode, SRNLabelDecode, DistillationCTCLabelDecode, \
|
||||||
|
TableLabelDecode
|
||||||
from .cls_postprocess import ClsPostProcess
|
from .cls_postprocess import ClsPostProcess
|
||||||
from .pg_postprocess import PGPostProcess
|
from .pg_postprocess import PGPostProcess
|
||||||
|
|
||||||
|
|
@ -33,7 +34,7 @@ def build_post_process(config, global_config=None):
|
||||||
support_dict = [
|
support_dict = [
|
||||||
'DBPostProcess', 'EASTPostProcess', 'SASTPostProcess', 'CTCLabelDecode',
|
'DBPostProcess', 'EASTPostProcess', 'SASTPostProcess', 'CTCLabelDecode',
|
||||||
'AttnLabelDecode', 'ClsPostProcess', 'SRNLabelDecode', 'PGPostProcess',
|
'AttnLabelDecode', 'ClsPostProcess', 'SRNLabelDecode', 'PGPostProcess',
|
||||||
'DistillationCTCLabelDecode'
|
'DistillationCTCLabelDecode', 'TableLabelDecode'
|
||||||
]
|
]
|
||||||
|
|
||||||
config = copy.deepcopy(config)
|
config = copy.deepcopy(config)
|
||||||
|
|
|
||||||
|
|
@ -44,16 +44,16 @@ class BaseRecLabelDecode(object):
|
||||||
self.character_str = string.printable[:-6]
|
self.character_str = string.printable[:-6]
|
||||||
dict_character = list(self.character_str)
|
dict_character = list(self.character_str)
|
||||||
elif character_type in support_character_type:
|
elif character_type in support_character_type:
|
||||||
self.character_str = ""
|
self.character_str = []
|
||||||
assert character_dict_path is not None, "character_dict_path should not be None when character_type is {}".format(
|
assert character_dict_path is not None, "character_dict_path should not be None when character_type is {}".format(
|
||||||
character_type)
|
character_type)
|
||||||
with open(character_dict_path, "rb") as fin:
|
with open(character_dict_path, "rb") as fin:
|
||||||
lines = fin.readlines()
|
lines = fin.readlines()
|
||||||
for line in lines:
|
for line in lines:
|
||||||
line = line.decode('utf-8').strip("\n").strip("\r\n")
|
line = line.decode('utf-8').strip("\n").strip("\r\n")
|
||||||
self.character_str += line
|
self.character_str.append(line)
|
||||||
if use_space_char:
|
if use_space_char:
|
||||||
self.character_str += " "
|
self.character_str.append(" ")
|
||||||
dict_character = list(self.character_str)
|
dict_character = list(self.character_str)
|
||||||
|
|
||||||
else:
|
else:
|
||||||
|
|
@ -319,3 +319,138 @@ class SRNLabelDecode(BaseRecLabelDecode):
|
||||||
assert False, "unsupport type %s in get_beg_end_flag_idx" \
|
assert False, "unsupport type %s in get_beg_end_flag_idx" \
|
||||||
% beg_or_end
|
% beg_or_end
|
||||||
return idx
|
return idx
|
||||||
|
|
||||||
|
|
||||||
|
class TableLabelDecode(object):
|
||||||
|
""" """
|
||||||
|
|
||||||
|
def __init__(self,
|
||||||
|
character_dict_path,
|
||||||
|
**kwargs):
|
||||||
|
list_character, list_elem = self.load_char_elem_dict(character_dict_path)
|
||||||
|
list_character = self.add_special_char(list_character)
|
||||||
|
list_elem = self.add_special_char(list_elem)
|
||||||
|
self.dict_character = {}
|
||||||
|
self.dict_idx_character = {}
|
||||||
|
for i, char in enumerate(list_character):
|
||||||
|
self.dict_idx_character[i] = char
|
||||||
|
self.dict_character[char] = i
|
||||||
|
self.dict_elem = {}
|
||||||
|
self.dict_idx_elem = {}
|
||||||
|
for i, elem in enumerate(list_elem):
|
||||||
|
self.dict_idx_elem[i] = elem
|
||||||
|
self.dict_elem[elem] = i
|
||||||
|
|
||||||
|
def load_char_elem_dict(self, character_dict_path):
|
||||||
|
list_character = []
|
||||||
|
list_elem = []
|
||||||
|
with open(character_dict_path, "rb") as fin:
|
||||||
|
lines = fin.readlines()
|
||||||
|
substr = lines[0].decode('utf-8').strip("\n").split("\t")
|
||||||
|
character_num = int(substr[0])
|
||||||
|
elem_num = int(substr[1])
|
||||||
|
for cno in range(1, 1 + character_num):
|
||||||
|
character = lines[cno].decode('utf-8').strip("\n")
|
||||||
|
list_character.append(character)
|
||||||
|
for eno in range(1 + character_num, 1 + character_num + elem_num):
|
||||||
|
elem = lines[eno].decode('utf-8').strip("\n")
|
||||||
|
list_elem.append(elem)
|
||||||
|
return list_character, list_elem
|
||||||
|
|
||||||
|
def add_special_char(self, list_character):
|
||||||
|
self.beg_str = "sos"
|
||||||
|
self.end_str = "eos"
|
||||||
|
list_character = [self.beg_str] + list_character + [self.end_str]
|
||||||
|
return list_character
|
||||||
|
|
||||||
|
def __call__(self, preds):
|
||||||
|
structure_probs = preds['structure_probs']
|
||||||
|
loc_preds = preds['loc_preds']
|
||||||
|
if isinstance(structure_probs,paddle.Tensor):
|
||||||
|
structure_probs = structure_probs.numpy()
|
||||||
|
if isinstance(loc_preds,paddle.Tensor):
|
||||||
|
loc_preds = loc_preds.numpy()
|
||||||
|
structure_idx = structure_probs.argmax(axis=2)
|
||||||
|
structure_probs = structure_probs.max(axis=2)
|
||||||
|
structure_str, structure_pos, result_score_list, result_elem_idx_list = self.decode(structure_idx,
|
||||||
|
structure_probs, 'elem')
|
||||||
|
res_html_code_list = []
|
||||||
|
res_loc_list = []
|
||||||
|
batch_num = len(structure_str)
|
||||||
|
for bno in range(batch_num):
|
||||||
|
res_loc = []
|
||||||
|
for sno in range(len(structure_str[bno])):
|
||||||
|
text = structure_str[bno][sno]
|
||||||
|
if text in ['<td>', '<td']:
|
||||||
|
pos = structure_pos[bno][sno]
|
||||||
|
res_loc.append(loc_preds[bno, pos])
|
||||||
|
res_html_code = ''.join(structure_str[bno])
|
||||||
|
res_loc = np.array(res_loc)
|
||||||
|
res_html_code_list.append(res_html_code)
|
||||||
|
res_loc_list.append(res_loc)
|
||||||
|
return {'res_html_code': res_html_code_list, 'res_loc': res_loc_list, 'res_score_list': result_score_list,
|
||||||
|
'res_elem_idx_list': result_elem_idx_list,'structure_str_list':structure_str}
|
||||||
|
|
||||||
|
def decode(self, text_index, structure_probs, char_or_elem):
|
||||||
|
"""convert text-label into text-index.
|
||||||
|
"""
|
||||||
|
if char_or_elem == "char":
|
||||||
|
current_dict = self.dict_idx_character
|
||||||
|
else:
|
||||||
|
current_dict = self.dict_idx_elem
|
||||||
|
ignored_tokens = self.get_ignored_tokens('elem')
|
||||||
|
beg_idx, end_idx = ignored_tokens
|
||||||
|
|
||||||
|
result_list = []
|
||||||
|
result_pos_list = []
|
||||||
|
result_score_list = []
|
||||||
|
result_elem_idx_list = []
|
||||||
|
batch_size = len(text_index)
|
||||||
|
for batch_idx in range(batch_size):
|
||||||
|
char_list = []
|
||||||
|
elem_pos_list = []
|
||||||
|
elem_idx_list = []
|
||||||
|
score_list = []
|
||||||
|
for idx in range(len(text_index[batch_idx])):
|
||||||
|
tmp_elem_idx = int(text_index[batch_idx][idx])
|
||||||
|
if idx > 0 and tmp_elem_idx == end_idx:
|
||||||
|
break
|
||||||
|
if tmp_elem_idx in ignored_tokens:
|
||||||
|
continue
|
||||||
|
|
||||||
|
char_list.append(current_dict[tmp_elem_idx])
|
||||||
|
elem_pos_list.append(idx)
|
||||||
|
score_list.append(structure_probs[batch_idx, idx])
|
||||||
|
elem_idx_list.append(tmp_elem_idx)
|
||||||
|
result_list.append(char_list)
|
||||||
|
result_pos_list.append(elem_pos_list)
|
||||||
|
result_score_list.append(score_list)
|
||||||
|
result_elem_idx_list.append(elem_idx_list)
|
||||||
|
return result_list, result_pos_list, result_score_list, result_elem_idx_list
|
||||||
|
|
||||||
|
def get_ignored_tokens(self, char_or_elem):
|
||||||
|
beg_idx = self.get_beg_end_flag_idx("beg", char_or_elem)
|
||||||
|
end_idx = self.get_beg_end_flag_idx("end", char_or_elem)
|
||||||
|
return [beg_idx, end_idx]
|
||||||
|
|
||||||
|
def get_beg_end_flag_idx(self, beg_or_end, char_or_elem):
|
||||||
|
if char_or_elem == "char":
|
||||||
|
if beg_or_end == "beg":
|
||||||
|
idx = self.dict_character[self.beg_str]
|
||||||
|
elif beg_or_end == "end":
|
||||||
|
idx = self.dict_character[self.end_str]
|
||||||
|
else:
|
||||||
|
assert False, "Unsupport type %s in get_beg_end_flag_idx of char" \
|
||||||
|
% beg_or_end
|
||||||
|
elif char_or_elem == "elem":
|
||||||
|
if beg_or_end == "beg":
|
||||||
|
idx = self.dict_elem[self.beg_str]
|
||||||
|
elif beg_or_end == "end":
|
||||||
|
idx = self.dict_elem[self.end_str]
|
||||||
|
else:
|
||||||
|
assert False, "Unsupport type %s in get_beg_end_flag_idx of elem" \
|
||||||
|
% beg_or_end
|
||||||
|
else:
|
||||||
|
assert False, "Unsupport type %s in char_or_elem" \
|
||||||
|
% char_or_elem
|
||||||
|
return idx
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,277 @@
|
||||||
|
←
|
||||||
|
</overline>
|
||||||
|
☆
|
||||||
|
─
|
||||||
|
α
|
||||||
|
|
||||||
|
|
||||||
|
⋅
|
||||||
|
$
|
||||||
|
ω
|
||||||
|
ψ
|
||||||
|
χ
|
||||||
|
(
|
||||||
|
υ
|
||||||
|
≥
|
||||||
|
σ
|
||||||
|
,
|
||||||
|
ρ
|
||||||
|
ε
|
||||||
|
0
|
||||||
|
■
|
||||||
|
4
|
||||||
|
8
|
||||||
|
✗
|
||||||
|
b
|
||||||
|
<
|
||||||
|
✓
|
||||||
|
Ψ
|
||||||
|
Ω
|
||||||
|
€
|
||||||
|
D
|
||||||
|
3
|
||||||
|
Π
|
||||||
|
H
|
||||||
|
║
|
||||||
|
</strike>
|
||||||
|
L
|
||||||
|
Φ
|
||||||
|
Χ
|
||||||
|
θ
|
||||||
|
P
|
||||||
|
κ
|
||||||
|
λ
|
||||||
|
μ
|
||||||
|
T
|
||||||
|
ξ
|
||||||
|
X
|
||||||
|
β
|
||||||
|
γ
|
||||||
|
δ
|
||||||
|
\
|
||||||
|
ζ
|
||||||
|
η
|
||||||
|
`
|
||||||
|
d
|
||||||
|
<strike>
|
||||||
|
h
|
||||||
|
f
|
||||||
|
l
|
||||||
|
Θ
|
||||||
|
p
|
||||||
|
√
|
||||||
|
t
|
||||||
|
</sub>
|
||||||
|
x
|
||||||
|
Β
|
||||||
|
Γ
|
||||||
|
Δ
|
||||||
|
|
|
||||||
|
ǂ
|
||||||
|
ɛ
|
||||||
|
j
|
||||||
|
̧
|
||||||
|
➢
|
||||||
|
|
||||||
|
̌
|
||||||
|
′
|
||||||
|
«
|
||||||
|
△
|
||||||
|
▲
|
||||||
|
#
|
||||||
|
</b>
|
||||||
|
'
|
||||||
|
Ι
|
||||||
|
+
|
||||||
|
¶
|
||||||
|
/
|
||||||
|
▼
|
||||||
|
⇑
|
||||||
|
□
|
||||||
|
·
|
||||||
|
7
|
||||||
|
▪
|
||||||
|
;
|
||||||
|
?
|
||||||
|
➔
|
||||||
|
∩
|
||||||
|
C
|
||||||
|
÷
|
||||||
|
G
|
||||||
|
⇒
|
||||||
|
K
|
||||||
|
<sup>
|
||||||
|
O
|
||||||
|
S
|
||||||
|
С
|
||||||
|
W
|
||||||
|
Α
|
||||||
|
[
|
||||||
|
○
|
||||||
|
_
|
||||||
|
●
|
||||||
|
‡
|
||||||
|
c
|
||||||
|
z
|
||||||
|
g
|
||||||
|
<i>
|
||||||
|
o
|
||||||
|
<sub>
|
||||||
|
〈
|
||||||
|
〉
|
||||||
|
s
|
||||||
|
⩽
|
||||||
|
w
|
||||||
|
φ
|
||||||
|
ʹ
|
||||||
|
{
|
||||||
|
»
|
||||||
|
∣
|
||||||
|
̆
|
||||||
|
e
|
||||||
|
ˆ
|
||||||
|
∈
|
||||||
|
τ
|
||||||
|
◆
|
||||||
|
ι
|
||||||
|
∅
|
||||||
|
∆
|
||||||
|
∙
|
||||||
|
∘
|
||||||
|
Ø
|
||||||
|
ß
|
||||||
|
✔
|
||||||
|
∞
|
||||||
|
∑
|
||||||
|
−
|
||||||
|
×
|
||||||
|
◊
|
||||||
|
∗
|
||||||
|
∖
|
||||||
|
˃
|
||||||
|
˂
|
||||||
|
∫
|
||||||
|
"
|
||||||
|
i
|
||||||
|
&
|
||||||
|
π
|
||||||
|
↔
|
||||||
|
*
|
||||||
|
∥
|
||||||
|
æ
|
||||||
|
∧
|
||||||
|
.
|
||||||
|
⁄
|
||||||
|
ø
|
||||||
|
Q
|
||||||
|
∼
|
||||||
|
6
|
||||||
|
⁎
|
||||||
|
:
|
||||||
|
★
|
||||||
|
>
|
||||||
|
a
|
||||||
|
B
|
||||||
|
≈
|
||||||
|
F
|
||||||
|
J
|
||||||
|
̄
|
||||||
|
N
|
||||||
|
♯
|
||||||
|
R
|
||||||
|
V
|
||||||
|
<overline>
|
||||||
|
―
|
||||||
|
Z
|
||||||
|
♣
|
||||||
|
^
|
||||||
|
¤
|
||||||
|
¥
|
||||||
|
§
|
||||||
|
<underline>
|
||||||
|
¢
|
||||||
|
£
|
||||||
|
≦
|
||||||
|
|
||||||
|
≤
|
||||||
|
‖
|
||||||
|
Λ
|
||||||
|
©
|
||||||
|
n
|
||||||
|
↓
|
||||||
|
→
|
||||||
|
↑
|
||||||
|
r
|
||||||
|
°
|
||||||
|
±
|
||||||
|
v
|
||||||
|
<b>
|
||||||
|
♂
|
||||||
|
k
|
||||||
|
♀
|
||||||
|
~
|
||||||
|
ᅟ
|
||||||
|
̇
|
||||||
|
@
|
||||||
|
”
|
||||||
|
♦
|
||||||
|
ł
|
||||||
|
®
|
||||||
|
⊕
|
||||||
|
„
|
||||||
|
!
|
||||||
|
</sup>
|
||||||
|
%
|
||||||
|
⇓
|
||||||
|
)
|
||||||
|
-
|
||||||
|
1
|
||||||
|
5
|
||||||
|
9
|
||||||
|
=
|
||||||
|
А
|
||||||
|
A
|
||||||
|
‰
|
||||||
|
⋆
|
||||||
|
Σ
|
||||||
|
E
|
||||||
|
◦
|
||||||
|
I
|
||||||
|
※
|
||||||
|
M
|
||||||
|
m
|
||||||
|
̨
|
||||||
|
⩾
|
||||||
|
†
|
||||||
|
</i>
|
||||||
|
•
|
||||||
|
U
|
||||||
|
Y
|
||||||
|
|
||||||
|
]
|
||||||
|
̸
|
||||||
|
2
|
||||||
|
‐
|
||||||
|
–
|
||||||
|
‒
|
||||||
|
̂
|
||||||
|
—
|
||||||
|
̀
|
||||||
|
́
|
||||||
|
’
|
||||||
|
‘
|
||||||
|
⋮
|
||||||
|
⋯
|
||||||
|
̊
|
||||||
|
“
|
||||||
|
̈
|
||||||
|
≧
|
||||||
|
q
|
||||||
|
u
|
||||||
|
ı
|
||||||
|
y
|
||||||
|
</underline>
|
||||||
|
|
||||||
|
̃
|
||||||
|
}
|
||||||
|
ν
|
||||||
File diff suppressed because it is too large
Load Diff
|
|
@ -22,7 +22,7 @@ logger_initialized = {}
|
||||||
|
|
||||||
|
|
||||||
@functools.lru_cache()
|
@functools.lru_cache()
|
||||||
def get_logger(name='root', log_file=None, log_level=logging.INFO):
|
def get_logger(name='root', log_file=None, log_level=logging.DEBUG):
|
||||||
"""Initialize and get a logger by name.
|
"""Initialize and get a logger by name.
|
||||||
If the logger has not been initialized, this method will initialize the
|
If the logger has not been initialized, this method will initialize the
|
||||||
logger by adding one or two handlers, otherwise the initialized logger will
|
logger by adding one or two handlers, otherwise the initialized logger will
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,82 @@
|
||||||
|
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||||
|
#
|
||||||
|
# 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
|
||||||
|
import tarfile
|
||||||
|
import requests
|
||||||
|
from tqdm import tqdm
|
||||||
|
|
||||||
|
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:
|
||||||
|
logger.error("Something went wrong while downloading models")
|
||||||
|
sys.exit(0)
|
||||||
|
|
||||||
|
|
||||||
|
def maybe_download(model_storage_directory, url):
|
||||||
|
# using custom model
|
||||||
|
tar_file_name_list = [
|
||||||
|
'inference.pdiparams', 'inference.pdiparams.info', 'inference.pdmodel'
|
||||||
|
]
|
||||||
|
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')):
|
||||||
|
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))
|
||||||
|
os.makedirs(model_storage_directory, exist_ok=True)
|
||||||
|
download_with_progressbar(url, tmp_path)
|
||||||
|
with tarfile.open(tmp_path, 'r') as tarObj:
|
||||||
|
for member in tarObj.getmembers():
|
||||||
|
filename = None
|
||||||
|
for tar_file_name in tar_file_name_list:
|
||||||
|
if tar_file_name in member.name:
|
||||||
|
filename = tar_file_name
|
||||||
|
if filename is None:
|
||||||
|
continue
|
||||||
|
file = tarObj.extractfile(member)
|
||||||
|
with open(
|
||||||
|
os.path.join(model_storage_directory, filename),
|
||||||
|
'wb') as f:
|
||||||
|
f.write(file.read())
|
||||||
|
os.remove(tmp_path)
|
||||||
|
|
||||||
|
|
||||||
|
def is_link(s):
|
||||||
|
return s is not None and s.startswith('http')
|
||||||
|
|
||||||
|
|
||||||
|
def confirm_model_dir_url(model_dir, default_model_dir, default_url):
|
||||||
|
url = default_url
|
||||||
|
if model_dir is None or is_link(model_dir):
|
||||||
|
if is_link(model_dir):
|
||||||
|
url = model_dir
|
||||||
|
file_name = url.split('/')[-1][:-4]
|
||||||
|
model_dir = default_model_dir
|
||||||
|
model_dir = os.path.join(model_dir, file_name)
|
||||||
|
return model_dir, url
|
||||||
|
|
@ -0,0 +1,9 @@
|
||||||
|
include LICENSE
|
||||||
|
include README.md
|
||||||
|
|
||||||
|
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 ppstructure *.py
|
||||||
|
|
||||||
|
|
@ -0,0 +1,17 @@
|
||||||
|
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||||
|
#
|
||||||
|
# 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.
|
||||||
|
|
||||||
|
from .paddlestructure import PaddleStructure, draw_result, to_excel
|
||||||
|
|
||||||
|
__all__ = ['PaddleStructure', 'draw_result', 'to_excel']
|
||||||
|
|
@ -0,0 +1,70 @@
|
||||||
|
# PaddleStructure
|
||||||
|
|
||||||
|
## 1. Introduction to pipeline
|
||||||
|
|
||||||
|
PaddleStructure is a toolkit for complex layout text OCR, the process is as follows
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
In PaddleStructure, the image will be analyzed by layoutparser first. In the layout analysis, the area in the image will be classified, and the OCR process will be carried out according to the category.
|
||||||
|
|
||||||
|
Currently layoutparser will output five categories:
|
||||||
|
1. Text
|
||||||
|
2. Title
|
||||||
|
3. Figure
|
||||||
|
4. List
|
||||||
|
5. Table
|
||||||
|
|
||||||
|
Types 1-4 follow the traditional OCR process, and 5 follow the Table OCR process.
|
||||||
|
|
||||||
|
## 2. LayoutParser
|
||||||
|
|
||||||
|
|
||||||
|
## 3. Table OCR
|
||||||
|
|
||||||
|
[doc](table/README.md)
|
||||||
|
|
||||||
|
## 4. PaddleStructure whl package introduction
|
||||||
|
|
||||||
|
### 4.1 Use
|
||||||
|
|
||||||
|
4.1.1 Use by code
|
||||||
|
```python
|
||||||
|
import cv2
|
||||||
|
from paddlestructure import PaddleStructure,draw_result
|
||||||
|
|
||||||
|
table_engine = PaddleStructure(
|
||||||
|
output='./output/table',
|
||||||
|
show_log=True)
|
||||||
|
|
||||||
|
img_path = '../doc/table/1.png'
|
||||||
|
img = cv2.imread(img_path)
|
||||||
|
result = table_engine(img)
|
||||||
|
for line in result:
|
||||||
|
print(line)
|
||||||
|
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
|
font_path = 'path/tp/PaddleOCR/doc/fonts/simfang.ttf'
|
||||||
|
image = Image.open(img_path).convert('RGB')
|
||||||
|
im_show = draw_result(image, result,font_path=font_path)
|
||||||
|
im_show = Image.fromarray(im_show)
|
||||||
|
im_show.save('result.jpg')
|
||||||
|
```
|
||||||
|
|
||||||
|
4.1.2 Use by command line
|
||||||
|
```bash
|
||||||
|
paddlestructure --image_dir=../doc/table/1.png
|
||||||
|
```
|
||||||
|
|
||||||
|
### 参数说明
|
||||||
|
大部分参数和paddleocr whl包保持一致，见 [whl包文档](../doc/doc_ch/whl.md)
|
||||||
|
|
||||||
|
| 字段 | 说明 | 默认值 |
|
||||||
|
|------------------------|------------------------------------------------------|------------------|
|
||||||
|
| output | excel和识别结果保存的地址 | ./output/table |
|
||||||
|
| structure_max_len | structure模型预测时，图像的长边resize尺度 | 488 |
|
||||||
|
| structure_model_dir | structure inference 模型地址 | None |
|
||||||
|
| structure_char_type | structure 模型所用字典地址 | ../ppocr/utils/dict/table_structure_dict.tx |
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -0,0 +1,69 @@
|
||||||
|
# PaddleStructure
|
||||||
|
|
||||||
|
## 1. pipeline介绍
|
||||||
|
|
||||||
|
PaddleStructure 是一个用于复杂板式文字OCR的工具包，流程如下
|
||||||
|
|
||||||
|
|
||||||
|
在PaddleStructure中，图片会先经由layoutparser进行版面分析，在版面分析中，会对图片里的区域进行分类，根据根据类别进行对于的ocr流程。
|
||||||
|
|
||||||
|
目前layoutparser会输出五个类别:
|
||||||
|
1. Text
|
||||||
|
2. Title
|
||||||
|
3. Figure
|
||||||
|
4. List
|
||||||
|
5. Table
|
||||||
|
|
||||||
|
1-4类走传统的OCR流程，5走表格的OCR流程。
|
||||||
|
|
||||||
|
## 2. LayoutParser
|
||||||
|
|
||||||
|
|
||||||
|
## 3. Table OCR
|
||||||
|
|
||||||
|
[文档](table/README_ch.md)
|
||||||
|
|
||||||
|
## 4. PaddleStructure whl包介绍
|
||||||
|
|
||||||
|
### 4.1 使用
|
||||||
|
|
||||||
|
4.1.1 代码使用
|
||||||
|
```python
|
||||||
|
import cv2
|
||||||
|
from paddlestructure import PaddleStructure,draw_result
|
||||||
|
|
||||||
|
table_engine = PaddleStructure(
|
||||||
|
output='./output/table',
|
||||||
|
show_log=True)
|
||||||
|
|
||||||
|
img_path = '../doc/table/1.png'
|
||||||
|
img = cv2.imread(img_path)
|
||||||
|
result = table_engine(img)
|
||||||
|
for line in result:
|
||||||
|
print(line)
|
||||||
|
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
|
font_path = 'path/tp/PaddleOCR/doc/fonts/simfang.ttf'
|
||||||
|
image = Image.open(img_path).convert('RGB')
|
||||||
|
im_show = draw_result(image, result,font_path=font_path)
|
||||||
|
im_show = Image.fromarray(im_show)
|
||||||
|
im_show.save('result.jpg')
|
||||||
|
```
|
||||||
|
|
||||||
|
4.1.2 命令行使用
|
||||||
|
```bash
|
||||||
|
paddlestructure --image_dir=../doc/table/1.png
|
||||||
|
```
|
||||||
|
|
||||||
|
### 参数说明
|
||||||
|
大部分参数和paddleocr whl包保持一致，见 [whl包文档](../doc/doc_ch/whl.md)
|
||||||
|
|
||||||
|
| 字段 | 说明 | 默认值 |
|
||||||
|
|------------------------|------------------------------------------------------|------------------|
|
||||||
|
| output | excel和识别结果保存的地址 | ./output/table |
|
||||||
|
| structure_max_len | structure模型预测时，图像的长边resize尺度 | 488 |
|
||||||
|
| structure_model_dir | structure inference 模型地址 | None |
|
||||||
|
| structure_char_type | structure 模型所用字典地址 | ../ppocr/utils/dict/table_structure_dict.tx |
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -0,0 +1,148 @@
|
||||||
|
# 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 logging
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
|
||||||
|
__dir__ = os.path.dirname(__file__)
|
||||||
|
sys.path.append(__dir__)
|
||||||
|
sys.path.append(os.path.join(__dir__, '..'))
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from ppocr.utils.logging import get_logger
|
||||||
|
from test.predict_system import OCRSystem, save_res
|
||||||
|
from test.table.predict_table import to_excel
|
||||||
|
from test.utility import init_args, draw_result
|
||||||
|
|
||||||
|
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, confirm_model_dir_url, is_link
|
||||||
|
|
||||||
|
__all__ = ['PaddleStructure', 'draw_result', 'to_excel']
|
||||||
|
|
||||||
|
VERSION = '2.1'
|
||||||
|
BASE_DIR = os.path.expanduser("~/.paddlestructure/")
|
||||||
|
|
||||||
|
model_urls = {
|
||||||
|
'det': 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_det_infer.tar',
|
||||||
|
'rec': 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_rec_infer.tar',
|
||||||
|
'structure': 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_structure_infer.tar'
|
||||||
|
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def parse_args(mMain=True):
|
||||||
|
import argparse
|
||||||
|
parser = init_args()
|
||||||
|
parser.add_help = mMain
|
||||||
|
|
||||||
|
for action in parser._actions:
|
||||||
|
if action.dest in ['rec_char_dict_path', 'structure_char_dict_path']:
|
||||||
|
action.default = None
|
||||||
|
if mMain:
|
||||||
|
return parser.parse_args()
|
||||||
|
else:
|
||||||
|
inference_args_dict = {}
|
||||||
|
for action in parser._actions:
|
||||||
|
inference_args_dict[action.dest] = action.default
|
||||||
|
return argparse.Namespace(**inference_args_dict)
|
||||||
|
|
||||||
|
|
||||||
|
class PaddleStructure(OCRSystem):
|
||||||
|
def __init__(self, **kwargs):
|
||||||
|
params = parse_args(mMain=False)
|
||||||
|
params.__dict__.update(**kwargs)
|
||||||
|
if not params.show_log:
|
||||||
|
logger.setLevel(logging.INFO)
|
||||||
|
params.use_angle_cls = False
|
||||||
|
# init model dir
|
||||||
|
params.det_model_dir, det_url = confirm_model_dir_url(params.det_model_dir,
|
||||||
|
os.path.join(BASE_DIR, VERSION, 'det'),
|
||||||
|
model_urls['det'])
|
||||||
|
params.rec_model_dir, rec_url = confirm_model_dir_url(params.rec_model_dir,
|
||||||
|
os.path.join(BASE_DIR, VERSION, 'rec'),
|
||||||
|
model_urls['rec'])
|
||||||
|
params.structure_model_dir, structure_url = confirm_model_dir_url(params.structure_model_dir,
|
||||||
|
os.path.join(BASE_DIR, VERSION, 'structure'),
|
||||||
|
model_urls['structure'])
|
||||||
|
# download model
|
||||||
|
maybe_download(params.det_model_dir, det_url)
|
||||||
|
maybe_download(params.rec_model_dir, rec_url)
|
||||||
|
maybe_download(params.structure_model_dir, structure_url)
|
||||||
|
|
||||||
|
if params.rec_char_dict_path is None:
|
||||||
|
params.rec_char_type = 'EN'
|
||||||
|
if os.path.exists(str(Path(__file__).parent / 'ppocr/utils/dict/table_dict.txt')):
|
||||||
|
params.rec_char_dict_path = str(Path(__file__).parent / 'ppocr/utils/dict/table_dict.txt')
|
||||||
|
else:
|
||||||
|
params.rec_char_dict_path = str(Path(__file__).parent.parent / 'ppocr/utils/dict/table_dict.txt')
|
||||||
|
if params.structure_char_dict_path is None:
|
||||||
|
if os.path.exists(str(Path(__file__).parent / 'ppocr/utils/dict/table_structure_dict.txt')):
|
||||||
|
params.structure_char_dict_path = str(
|
||||||
|
Path(__file__).parent / 'ppocr/utils/dict/table_structure_dict.txt')
|
||||||
|
else:
|
||||||
|
params.structure_char_dict_path = str(
|
||||||
|
Path(__file__).parent.parent / 'ppocr/utils/dict/table_structure_dict.txt')
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
# for cmd
|
||||||
|
args = parse_args(mMain=True)
|
||||||
|
image_dir = args.image_dir
|
||||||
|
save_folder = args.output
|
||||||
|
if image_dir.startswith('http'):
|
||||||
|
download_with_progressbar(image_dir, 'tmp.jpg')
|
||||||
|
image_file_list = ['tmp.jpg']
|
||||||
|
else:
|
||||||
|
image_file_list = get_image_file_list(args.image_dir)
|
||||||
|
if len(image_file_list) == 0:
|
||||||
|
logger.error('no images find in {}'.format(args.image_dir))
|
||||||
|
return
|
||||||
|
|
||||||
|
structure_engine = PaddleStructure(**(args.__dict__))
|
||||||
|
for img_path in image_file_list:
|
||||||
|
img_name = os.path.basename(img_path).split('.')[0]
|
||||||
|
logger.info('{}{}{}'.format('*' * 10, img_path, '*' * 10))
|
||||||
|
result = structure_engine(img_path)
|
||||||
|
for item in result:
|
||||||
|
logger.info(item['res'])
|
||||||
|
save_res(result, save_folder, img_name)
|
||||||
|
logger.info('result save to {}'.format(os.path.join(save_folder, img_name)))
|
||||||
|
|
@ -0,0 +1,134 @@
|
||||||
|
# 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
|
||||||
|
import subprocess
|
||||||
|
|
||||||
|
__dir__ = os.path.dirname(os.path.abspath(__file__))
|
||||||
|
sys.path.append(__dir__)
|
||||||
|
sys.path.append(os.path.abspath(os.path.join(__dir__, '..')))
|
||||||
|
|
||||||
|
os.environ["FLAGS_allocator_strategy"] = 'auto_growth'
|
||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
import time
|
||||||
|
import logging
|
||||||
|
|
||||||
|
import layoutparser as lp
|
||||||
|
|
||||||
|
from ppocr.utils.utility import get_image_file_list, check_and_read_gif
|
||||||
|
from ppocr.utils.logging import get_logger
|
||||||
|
from tools.infer.predict_system import TextSystem
|
||||||
|
from test.table.predict_table import TableSystem, to_excel
|
||||||
|
from test.utility import parse_args, draw_result
|
||||||
|
|
||||||
|
logger = get_logger()
|
||||||
|
|
||||||
|
|
||||||
|
class OCRSystem(object):
|
||||||
|
def __init__(self, args):
|
||||||
|
args.det_limit_type = 'resize_long'
|
||||||
|
args.drop_score = 0
|
||||||
|
if not args.show_log:
|
||||||
|
logger.setLevel(logging.INFO)
|
||||||
|
self.text_system = TextSystem(args)
|
||||||
|
self.table_system = TableSystem(args, self.text_system.text_detector, self.text_system.text_recognizer)
|
||||||
|
self.table_layout = lp.PaddleDetectionLayoutModel("lp://PubLayNet/ppyolov2_r50vd_dcn_365e_publaynet/config",
|
||||||
|
threshold=0.5, enable_mkldnn=args.enable_mkldnn,
|
||||||
|
enforce_cpu=not args.use_gpu, thread_num=args.cpu_threads)
|
||||||
|
self.use_angle_cls = args.use_angle_cls
|
||||||
|
self.drop_score = args.drop_score
|
||||||
|
|
||||||
|
def __call__(self, img):
|
||||||
|
ori_im = img.copy()
|
||||||
|
layout_res = self.table_layout.detect(img[..., ::-1])
|
||||||
|
res_list = []
|
||||||
|
for region in layout_res:
|
||||||
|
x1, y1, x2, y2 = region.coordinates
|
||||||
|
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
|
||||||
|
roi_img = ori_im[y1:y2, x1:x2, :]
|
||||||
|
if region.type == 'Table':
|
||||||
|
res = self.table_system(roi_img)
|
||||||
|
else:
|
||||||
|
filter_boxes, filter_rec_res = self.text_system(roi_img)
|
||||||
|
filter_boxes = [x + [x1, y1] for x in filter_boxes]
|
||||||
|
filter_boxes = [x.reshape(-1).tolist() for x in filter_boxes]
|
||||||
|
|
||||||
|
res = (filter_boxes, filter_rec_res)
|
||||||
|
res_list.append({'type': region.type, 'bbox': [x1, y1, x2, y2], 'res': res})
|
||||||
|
return res_list
|
||||||
|
|
||||||
|
|
||||||
|
def save_res(res, save_folder, img_name):
|
||||||
|
excel_save_folder = os.path.join(save_folder, img_name)
|
||||||
|
os.makedirs(excel_save_folder, exist_ok=True)
|
||||||
|
# save res
|
||||||
|
for region in res:
|
||||||
|
if region['type'] == 'Table':
|
||||||
|
excel_path = os.path.join(excel_save_folder, '{}.xlsx'.format(region['bbox']))
|
||||||
|
to_excel(region['res'], excel_path)
|
||||||
|
elif region['type'] == 'Figure':
|
||||||
|
pass
|
||||||
|
else:
|
||||||
|
with open(os.path.join(excel_save_folder, 'res.txt'), 'a', encoding='utf8') as f:
|
||||||
|
for box, rec_res in zip(region['res'][0], region['res'][1]):
|
||||||
|
f.write('{}\t{}\n'.format(np.array(box).reshape(-1).tolist(), rec_res))
|
||||||
|
|
||||||
|
|
||||||
|
def main(args):
|
||||||
|
image_file_list = get_image_file_list(args.image_dir)
|
||||||
|
image_file_list = image_file_list
|
||||||
|
image_file_list = image_file_list[args.process_id::args.total_process_num]
|
||||||
|
save_folder = args.output
|
||||||
|
os.makedirs(save_folder, exist_ok=True)
|
||||||
|
|
||||||
|
structure_sys = OCRSystem(args)
|
||||||
|
img_num = len(image_file_list)
|
||||||
|
for i, image_file in enumerate(image_file_list):
|
||||||
|
logger.info("[{}/{}] {}".format(i, img_num, image_file))
|
||||||
|
img, flag = check_and_read_gif(image_file)
|
||||||
|
img_name = os.path.basename(image_file).split('.')[0]
|
||||||
|
|
||||||
|
if not flag:
|
||||||
|
img = cv2.imread(image_file)
|
||||||
|
if img is None:
|
||||||
|
logger.error("error in loading image:{}".format(image_file))
|
||||||
|
continue
|
||||||
|
starttime = time.time()
|
||||||
|
res = structure_sys(img)
|
||||||
|
save_res(res, save_folder, img_name)
|
||||||
|
draw_img = draw_result(img, res, args.vis_font_path)
|
||||||
|
cv2.imwrite(os.path.join(save_folder, img_name, 'show.jpg'), draw_img)
|
||||||
|
logger.info('result save to {}'.format(os.path.join(save_folder, img_name)))
|
||||||
|
elapse = time.time() - starttime
|
||||||
|
logger.info("Predict time : {:.3f}s".format(elapse))
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
args = parse_args()
|
||||||
|
if args.use_mp:
|
||||||
|
p_list = []
|
||||||
|
total_process_num = args.total_process_num
|
||||||
|
for process_id in range(total_process_num):
|
||||||
|
cmd = [sys.executable, "-u"] + sys.argv + [
|
||||||
|
"--process_id={}".format(process_id),
|
||||||
|
"--use_mp={}".format(False)
|
||||||
|
]
|
||||||
|
p = subprocess.Popen(cmd, stdout=sys.stdout, stderr=sys.stdout)
|
||||||
|
p_list.append(p)
|
||||||
|
for p in p_list:
|
||||||
|
p.wait()
|
||||||
|
else:
|
||||||
|
main(args)
|
||||||
|
|
@ -0,0 +1,72 @@
|
||||||
|
# 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
|
||||||
|
|
||||||
|
from setuptools import setup
|
||||||
|
from io import open
|
||||||
|
import shutil
|
||||||
|
|
||||||
|
with open('../requirements.txt', encoding="utf-8-sig") as f:
|
||||||
|
requirements = f.readlines()
|
||||||
|
requirements.append('tqdm')
|
||||||
|
requirements.append('layoutparser')
|
||||||
|
requirements.append('iopath')
|
||||||
|
|
||||||
|
|
||||||
|
def readme():
|
||||||
|
with open('api_ch.md', encoding="utf-8-sig") as f:
|
||||||
|
README = f.read()
|
||||||
|
return README
|
||||||
|
|
||||||
|
|
||||||
|
shutil.copytree('/table', './test/table')
|
||||||
|
shutil.copyfile('/predict_system.py', './test/predict_system.py')
|
||||||
|
shutil.copyfile('/utility.py', './test/utility.py')
|
||||||
|
shutil.copytree('../ppocr', './ppocr')
|
||||||
|
shutil.copytree('../tools', './tools')
|
||||||
|
shutil.copyfile('../LICENSE', './LICENSE')
|
||||||
|
|
||||||
|
setup(
|
||||||
|
name='paddlestructure',
|
||||||
|
packages=['paddlestructure'],
|
||||||
|
package_dir={'paddlestructure': ''},
|
||||||
|
include_package_data=True,
|
||||||
|
entry_points={"console_scripts": ["paddlestructure= paddlestructure.paddlestructure:main"]},
|
||||||
|
version='1.0',
|
||||||
|
install_requires=requirements,
|
||||||
|
license='Apache License 2.0',
|
||||||
|
description='Awesome OCR toolkits based on PaddlePaddle （8.6M ultra-lightweight pre-trained model, support training and deployment among server, mobile, embeded and IoT devices',
|
||||||
|
long_description=readme(),
|
||||||
|
long_description_content_type='text/markdown',
|
||||||
|
url='https://github.com/PaddlePaddle/PaddleOCR',
|
||||||
|
download_url='https://github.com/PaddlePaddle/PaddleOCR.git',
|
||||||
|
keywords=[
|
||||||
|
'ocr textdetection textrecognition paddleocr crnn east star-net rosetta ocrlite db chineseocr chinesetextdetection chinesetextrecognition'
|
||||||
|
],
|
||||||
|
classifiers=[
|
||||||
|
'Intended Audience :: Developers', 'Operating System :: OS Independent',
|
||||||
|
'Natural Language :: Chinese (Simplified)',
|
||||||
|
'Programming Language :: Python :: 3',
|
||||||
|
'Programming Language :: Python :: 3.2',
|
||||||
|
'Programming Language :: Python :: 3.3',
|
||||||
|
'Programming Language :: Python :: 3.4',
|
||||||
|
'Programming Language :: Python :: 3.5',
|
||||||
|
'Programming Language :: Python :: 3.6',
|
||||||
|
'Programming Language :: Python :: 3.7', 'Topic :: Utilities'
|
||||||
|
], )
|
||||||
|
|
||||||
|
shutil.rmtree('ppocr')
|
||||||
|
shutil.rmtree('tools')
|
||||||
|
shutil.rmtree('test')
|
||||||
|
os.remove('LICENSE')
|
||||||
|
|
@ -0,0 +1,49 @@
|
||||||
|
# Table structure and content prediction
|
||||||
|
|
||||||
|
## 1. pipeline
|
||||||
|
The ocr of the table mainly contains three models
|
||||||
|
1. Single line text detection-DB
|
||||||
|
2. Single line text recognition-CRNN
|
||||||
|
3. Table structure and cell coordinate prediction-RARE
|
||||||
|
|
||||||
|
The table ocr flow chart is as follows
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
1. The coordinates of single-line text is detected by DB model, and then sends it to the recognition model to get the recognition result.
|
||||||
|
2. The table structure and cell coordinates is predicted by RARE model.
|
||||||
|
3. The recognition result of the cell is combined by the coordinates, recognition result of the single line and the coordinates of the cell.
|
||||||
|
4. The cell recognition result and the table structure together construct the html string of the table.
|
||||||
|
|
||||||
|
## 2. How to use
|
||||||
|
|
||||||
|
|
||||||
|
### 2.1 Train
|
||||||
|
TBD
|
||||||
|
|
||||||
|
### 2.2 Eval
|
||||||
|
First cd to the PaddleOCR/ppstructure directory
|
||||||
|
|
||||||
|
The table uses TEDS (Tree-Edit-Distance-based Similarity) as the evaluation metric of the model. Before the model evaluation, the three models in the pipeline need to be exported as inference models (we have provided them), and the gt for evaluation needs to be prepared. Examples of gt are as follows:
|
||||||
|
```json
|
||||||
|
{"PMC4289340_004_00.png": [["<html>", "<body>", "<table>", "<thead>", "<tr>", "<td>", "</td>", "<td>", "</td>", "<td>", "</td>", "</tr>", "</thead>", "<tbody>", "<tr>", "<td>", "</td>", "<td>", "</td>", "<td>", "</td>", "</tr>", "</tbody>", "</table>", "</body>", "</html>"], [[1, 4, 29, 13], [137, 4, 161, 13], [215, 4, 236, 13], [1, 17, 30, 27], [137, 17, 147, 27], [215, 17, 225, 27]], [["<b>", "F", "e", "a", "t", "u", "r", "e", "</b>"], ["<b>", "G", "b", "3", " ", "+", "</b>"], ["<b>", "G", "b", "3", " ", "-", "</b>"], ["<b>", "P", "a", "t", "i", "e", "n", "t", "s", "</b>"], ["6", "2"], ["4", "5"]]]}
|
||||||
|
```
|
||||||
|
In gt json, the key is the image name, the value is the corresponding gt, and gt is a list composed of four items, and each item is
|
||||||
|
1. HTML string list of table structure
|
||||||
|
2. The coordinates of each cell (not including the empty text in the cell)
|
||||||
|
3. The text information in each cell (not including the empty text in the cell)
|
||||||
|
4. The text information in each cell (including the empty text in the cell)
|
||||||
|
|
||||||
|
Use the following command to evaluate. After the evaluation is completed, the teds indicator will be output.
|
||||||
|
```python
|
||||||
|
python3 table/eval_table.py --det_model_dir=path/to/det_model_dir --rec_model_dir=path/to/rec_model_dir --structure_model_dir=path/to/structure_model_dir --image_dir=../doc/table/1.png --rec_char_dict_path=../ppocr/utils/dict/table_dict.txt --structure_char_dict_path=../ppocr/utils/dict/table_structure_dict.txt --rec_char_type=EN --det_limit_side_len=736 --det_limit_type=min --gt_path=path/to/gt.json
|
||||||
|
```
|
||||||
|
|
||||||
|
|
||||||
|
### 2.3 Inference
|
||||||
|
First cd to the PaddleOCR/ppstructure directory
|
||||||
|
|
||||||
|
```python
|
||||||
|
python3 table/predict_table.py --det_model_dir=path/to/det_model_dir --rec_model_dir=path/to/rec_model_dir --structure_model_dir=path/to/structure_model_dir --image_dir=../doc/table/1.png --rec_char_dict_path=../ppocr/utils/dict/table_dict.txt --structure_char_dict_path=../ppocr/utils/dict/table_structure_dict.txt --rec_char_type=EN --det_limit_side_len=736 --det_limit_type=min --output ../output/table
|
||||||
|
```
|
||||||
|
After running, the excel sheet of each picture will be saved in the directory specified by the table_output field
|
||||||
|
|
@ -0,0 +1,49 @@
|
||||||
|
# 表格结构和内容预测
|
||||||
|
|
||||||
|
## 1. pipeline
|
||||||
|
表格的ocr主要包含三个模型
|
||||||
|
1. 单行文本检测-DB
|
||||||
|
2. 单行文本识别-CRNN
|
||||||
|
3. 表格结构和cell坐标预测-RARE
|
||||||
|
|
||||||
|
具体流程图如下
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
1. 图片由单行文字检测检测模型到单行文字的坐标，然后送入识别模型拿到识别结果。
|
||||||
|
2. 图片由表格结构和cell坐标预测模型拿到表格的结构信息和单元格的坐标信息。
|
||||||
|
3. 由单行文字的坐标、识别结果和单元格的坐标一起组合出单元格的识别结果。
|
||||||
|
4. 单元格的识别结果和表格结构一起构造表格的html字符串。
|
||||||
|
|
||||||
|
## 2. 使用
|
||||||
|
|
||||||
|
|
||||||
|
### 2.1 训练
|
||||||
|
TBD
|
||||||
|
|
||||||
|
### 2.2 评估
|
||||||
|
先cd到PaddleOCR/ppstructure目录下
|
||||||
|
|
||||||
|
表格使用 TEDS(Tree-Edit-Distance-based Similarity) 作为模型的评估指标。在进行模型评估之前，需要将pipeline中的三个模型分别导出为inference模型(我们已经提供好)，还需要准备评估的gt， gt示例如下:
|
||||||
|
```json
|
||||||
|
{"PMC4289340_004_00.png": [["<html>", "<body>", "<table>", "<thead>", "<tr>", "<td>", "</td>", "<td>", "</td>", "<td>", "</td>", "</tr>", "</thead>", "<tbody>", "<tr>", "<td>", "</td>", "<td>", "</td>", "<td>", "</td>", "</tr>", "</tbody>", "</table>", "</body>", "</html>"], [[1, 4, 29, 13], [137, 4, 161, 13], [215, 4, 236, 13], [1, 17, 30, 27], [137, 17, 147, 27], [215, 17, 225, 27]], [["<b>", "F", "e", "a", "t", "u", "r", "e", "</b>"], ["<b>", "G", "b", "3", " ", "+", "</b>"], ["<b>", "G", "b", "3", " ", "-", "</b>"], ["<b>", "P", "a", "t", "i", "e", "n", "t", "s", "</b>"], ["6", "2"], ["4", "5"]]]}
|
||||||
|
```
|
||||||
|
json 中，key为图片名，value为对于的gt，gt是一个由四个item组成的list，每个item分别为
|
||||||
|
1. 表格结构的html字符串list
|
||||||
|
2. 每个cell的坐标 (不包括cell里文字为空的)
|
||||||
|
3. 每个cell里的文字信息 (不包括cell里文字为空的)
|
||||||
|
4. 每个cell里的文字信息 (包括cell里文字为空的)
|
||||||
|
|
||||||
|
准备完成后使用如下命令进行评估，评估完成后会输出teds指标。
|
||||||
|
```python
|
||||||
|
python3 table/eval_table.py --det_model_dir=path/to/det_model_dir --rec_model_dir=path/to/rec_model_dir --structure_model_dir=path/to/structure_model_dir --image_dir=../doc/table/1.png --rec_char_dict_path=../ppocr/utils/dict/table_dict.txt --structure_char_dict_path=../ppocr/utils/dict/table_structure_dict.txt --rec_char_type=EN --det_limit_side_len=736 --det_limit_type=min --gt_path=path/to/gt.json
|
||||||
|
```
|
||||||
|
|
||||||
|
|
||||||
|
### 2.3 预测
|
||||||
|
先cd到PaddleOCR/ppstructure目录下
|
||||||
|
|
||||||
|
```python
|
||||||
|
python3 table/predict_table.py --det_model_dir=path/to/det_model_dir --rec_model_dir=path/to/rec_model_dir --structure_model_dir=path/to/structure_model_dir --image_dir=../doc/table/1.png --rec_char_dict_path=../ppocr/utils/dict/table_dict.txt --structure_char_dict_path=../ppocr/utils/dict/table_structure_dict.txt --rec_char_type=EN --det_limit_side_len=736 --det_limit_type=min --output ../output/table
|
||||||
|
```
|
||||||
|
运行完成后，每张图片的excel表格会保存到table_output字段指定的目录下
|
||||||
|
|
@ -0,0 +1,13 @@
|
||||||
|
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||||
|
#
|
||||||
|
# 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.
|
||||||
|
|
@ -0,0 +1,72 @@
|
||||||
|
# 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(os.path.abspath(__file__))
|
||||||
|
sys.path.append(__dir__)
|
||||||
|
sys.path.append(os.path.abspath(os.path.join(__dir__, '../..')))
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
import json
|
||||||
|
from tqdm import tqdm
|
||||||
|
from test.table.table_metric import TEDS
|
||||||
|
from test.table.predict_table import TableSystem
|
||||||
|
from test.utility import init_args
|
||||||
|
from ppocr.utils.logging import get_logger
|
||||||
|
|
||||||
|
logger = get_logger()
|
||||||
|
|
||||||
|
|
||||||
|
def parse_args():
|
||||||
|
parser = init_args()
|
||||||
|
parser.add_argument("--gt_path", type=str)
|
||||||
|
return parser.parse_args()
|
||||||
|
|
||||||
|
def main(gt_path, img_root, args):
|
||||||
|
teds = TEDS(n_jobs=16)
|
||||||
|
|
||||||
|
text_sys = TableSystem(args)
|
||||||
|
jsons_gt = json.load(open(gt_path)) # gt
|
||||||
|
pred_htmls = []
|
||||||
|
gt_htmls = []
|
||||||
|
for img_name in tqdm(jsons_gt):
|
||||||
|
# read image
|
||||||
|
img = cv2.imread(os.path.join(img_root,img_name))
|
||||||
|
pred_html = text_sys(img)
|
||||||
|
pred_htmls.append(pred_html)
|
||||||
|
|
||||||
|
gt_structures, gt_bboxes, gt_contents, contents_with_block = jsons_gt[img_name]
|
||||||
|
gt_html, gt = get_gt_html(gt_structures, contents_with_block)
|
||||||
|
gt_htmls.append(gt_html)
|
||||||
|
scores = teds.batch_evaluate_html(gt_htmls, pred_htmls)
|
||||||
|
logger.info('teds:', sum(scores) / len(scores))
|
||||||
|
|
||||||
|
|
||||||
|
def get_gt_html(gt_structures, contents_with_block):
|
||||||
|
end_html = []
|
||||||
|
td_index = 0
|
||||||
|
for tag in gt_structures:
|
||||||
|
if '</td>' in tag:
|
||||||
|
if contents_with_block[td_index] != []:
|
||||||
|
end_html.extend(contents_with_block[td_index])
|
||||||
|
end_html.append(tag)
|
||||||
|
td_index += 1
|
||||||
|
else:
|
||||||
|
end_html.append(tag)
|
||||||
|
return ''.join(end_html), end_html
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
args = parse_args()
|
||||||
|
main(args.gt_path,args.image_dir, args)
|
||||||
|
|
@ -0,0 +1,192 @@
|
||||||
|
import json
|
||||||
|
def distance(box_1, box_2):
|
||||||
|
x1, y1, x2, y2 = box_1
|
||||||
|
x3, y3, x4, y4 = box_2
|
||||||
|
dis = abs(x3 - x1) + abs(y3 - y1) + abs(x4- x2) + abs(y4 - y2)
|
||||||
|
dis_2 = abs(x3 - x1) + abs(y3 - y1)
|
||||||
|
dis_3 = abs(x4- x2) + abs(y4 - y2)
|
||||||
|
return dis + min(dis_2, dis_3)
|
||||||
|
|
||||||
|
def compute_iou(rec1, rec2):
|
||||||
|
"""
|
||||||
|
computing IoU
|
||||||
|
:param rec1: (y0, x0, y1, x1), which reflects
|
||||||
|
(top, left, bottom, right)
|
||||||
|
:param rec2: (y0, x0, y1, x1)
|
||||||
|
:return: scala value of IoU
|
||||||
|
"""
|
||||||
|
# computing area of each rectangles
|
||||||
|
S_rec1 = (rec1[2] - rec1[0]) * (rec1[3] - rec1[1])
|
||||||
|
S_rec2 = (rec2[2] - rec2[0]) * (rec2[3] - rec2[1])
|
||||||
|
|
||||||
|
# computing the sum_area
|
||||||
|
sum_area = S_rec1 + S_rec2
|
||||||
|
|
||||||
|
# find the each edge of intersect rectangle
|
||||||
|
left_line = max(rec1[1], rec2[1])
|
||||||
|
right_line = min(rec1[3], rec2[3])
|
||||||
|
top_line = max(rec1[0], rec2[0])
|
||||||
|
bottom_line = min(rec1[2], rec2[2])
|
||||||
|
|
||||||
|
# judge if there is an intersect
|
||||||
|
if left_line >= right_line or top_line >= bottom_line:
|
||||||
|
return 0.0
|
||||||
|
else:
|
||||||
|
intersect = (right_line - left_line) * (bottom_line - top_line)
|
||||||
|
return (intersect / (sum_area - intersect))*1.0
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
def matcher_merge(ocr_bboxes, pred_bboxes):
|
||||||
|
all_dis = []
|
||||||
|
ious = []
|
||||||
|
matched = {}
|
||||||
|
for i, gt_box in enumerate(ocr_bboxes):
|
||||||
|
distances = []
|
||||||
|
for j, pred_box in enumerate(pred_bboxes):
|
||||||
|
# compute l1 distence and IOU between two boxes
|
||||||
|
distances.append((distance(gt_box, pred_box), 1. - compute_iou(gt_box, pred_box)))
|
||||||
|
sorted_distances = distances.copy()
|
||||||
|
# select nearest cell
|
||||||
|
sorted_distances = sorted(sorted_distances, key = lambda item: (item[1], item[0]))
|
||||||
|
if distances.index(sorted_distances[0]) not in matched.keys():
|
||||||
|
matched[distances.index(sorted_distances[0])] = [i]
|
||||||
|
else:
|
||||||
|
matched[distances.index(sorted_distances[0])].append(i)
|
||||||
|
return matched#, sum(ious) / len(ious)
|
||||||
|
|
||||||
|
def complex_num(pred_bboxes):
|
||||||
|
complex_nums = []
|
||||||
|
for bbox in pred_bboxes:
|
||||||
|
distances = []
|
||||||
|
temp_ious = []
|
||||||
|
for pred_bbox in pred_bboxes:
|
||||||
|
if bbox != pred_bbox:
|
||||||
|
distances.append(distance(bbox, pred_bbox))
|
||||||
|
temp_ious.append(compute_iou(bbox, pred_bbox))
|
||||||
|
complex_nums.append(temp_ious[distances.index(min(distances))])
|
||||||
|
return sum(complex_nums) / len(complex_nums)
|
||||||
|
|
||||||
|
def get_rows(pred_bboxes):
|
||||||
|
pre_bbox = pred_bboxes[0]
|
||||||
|
res = []
|
||||||
|
step = 0
|
||||||
|
for i in range(len(pred_bboxes)):
|
||||||
|
bbox = pred_bboxes[i]
|
||||||
|
if bbox[1] - pre_bbox[1] > 2 or bbox[0] - pre_bbox[0] < 0:
|
||||||
|
break
|
||||||
|
else:
|
||||||
|
res.append(bbox)
|
||||||
|
step += 1
|
||||||
|
for i in range(step):
|
||||||
|
pred_bboxes.pop(0)
|
||||||
|
return res, pred_bboxes
|
||||||
|
def refine_rows(pred_bboxes): # 微调整行的框，使在一条水平线上
|
||||||
|
ys_1 = []
|
||||||
|
ys_2 = []
|
||||||
|
for box in pred_bboxes:
|
||||||
|
ys_1.append(box[1])
|
||||||
|
ys_2.append(box[3])
|
||||||
|
min_y_1 = sum(ys_1) / len(ys_1)
|
||||||
|
min_y_2 = sum(ys_2) / len(ys_2)
|
||||||
|
re_boxes = []
|
||||||
|
for box in pred_bboxes:
|
||||||
|
box[1] = min_y_1
|
||||||
|
box[3] = min_y_2
|
||||||
|
re_boxes.append(box)
|
||||||
|
return re_boxes
|
||||||
|
|
||||||
|
def matcher_refine_row(gt_bboxes, pred_bboxes):
|
||||||
|
before_refine_pred_bboxes = pred_bboxes.copy()
|
||||||
|
pred_bboxes = []
|
||||||
|
while(len(before_refine_pred_bboxes) != 0):
|
||||||
|
row_bboxes, before_refine_pred_bboxes = get_rows(before_refine_pred_bboxes)
|
||||||
|
print(row_bboxes)
|
||||||
|
pred_bboxes.extend(refine_rows(row_bboxes))
|
||||||
|
all_dis = []
|
||||||
|
ious = []
|
||||||
|
matched = {}
|
||||||
|
for i, gt_box in enumerate(gt_bboxes):
|
||||||
|
distances = []
|
||||||
|
#temp_ious = []
|
||||||
|
for j, pred_box in enumerate(pred_bboxes):
|
||||||
|
distances.append(distance(gt_box, pred_box))
|
||||||
|
#temp_ious.append(compute_iou(gt_box, pred_box))
|
||||||
|
#all_dis.append(min(distances))
|
||||||
|
#ious.append(temp_ious[distances.index(min(distances))])
|
||||||
|
if distances.index(min(distances)) not in matched.keys():
|
||||||
|
matched[distances.index(min(distances))] = [i]
|
||||||
|
else:
|
||||||
|
matched[distances.index(min(distances))].append(i)
|
||||||
|
return matched#, sum(ious) / len(ious)
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
#先挑选出一行，再进行匹配
|
||||||
|
def matcher_structure_1(gt_bboxes, pred_bboxes_rows, pred_bboxes):
|
||||||
|
gt_box_index = 0
|
||||||
|
delete_gt_bboxes = gt_bboxes.copy()
|
||||||
|
match_bboxes_ready = []
|
||||||
|
matched = {}
|
||||||
|
while(len(delete_gt_bboxes) != 0):
|
||||||
|
row_bboxes, delete_gt_bboxes = get_rows(delete_gt_bboxes)
|
||||||
|
row_bboxes = sorted(row_bboxes, key = lambda key: key[0])
|
||||||
|
if len(pred_bboxes_rows) > 0:
|
||||||
|
match_bboxes_ready.extend(pred_bboxes_rows.pop(0))
|
||||||
|
print(row_bboxes)
|
||||||
|
for i, gt_box in enumerate(row_bboxes):
|
||||||
|
#print(gt_box)
|
||||||
|
pred_distances = []
|
||||||
|
distances = []
|
||||||
|
for pred_bbox in pred_bboxes:
|
||||||
|
pred_distances.append(distance(gt_box, pred_bbox))
|
||||||
|
for j, pred_box in enumerate(match_bboxes_ready):
|
||||||
|
distances.append(distance(gt_box, pred_box))
|
||||||
|
index = pred_distances.index(min(distances))
|
||||||
|
#print('index', index)
|
||||||
|
if index not in matched.keys():
|
||||||
|
matched[index] = [gt_box_index]
|
||||||
|
else:
|
||||||
|
matched[index].append(gt_box_index)
|
||||||
|
gt_box_index += 1
|
||||||
|
return matched
|
||||||
|
|
||||||
|
def matcher_structure(gt_bboxes, pred_bboxes_rows, pred_bboxes):
|
||||||
|
'''
|
||||||
|
gt_bboxes: 排序后
|
||||||
|
pred_bboxes:
|
||||||
|
'''
|
||||||
|
pre_bbox = gt_bboxes[0]
|
||||||
|
matched = {}
|
||||||
|
match_bboxes_ready = []
|
||||||
|
match_bboxes_ready.extend(pred_bboxes_rows.pop(0))
|
||||||
|
for i, gt_box in enumerate(gt_bboxes):
|
||||||
|
|
||||||
|
pred_distances = []
|
||||||
|
for pred_bbox in pred_bboxes:
|
||||||
|
pred_distances.append(distance(gt_box, pred_bbox))
|
||||||
|
distances = []
|
||||||
|
gap_pre = gt_box[1] - pre_bbox[1]
|
||||||
|
gap_pre_1 = gt_box[0] - pre_bbox[2]
|
||||||
|
#print(gap_pre, len(pred_bboxes_rows))
|
||||||
|
if (gap_pre_1 < 0 and len(pred_bboxes_rows) > 0):
|
||||||
|
match_bboxes_ready.extend(pred_bboxes_rows.pop(0))
|
||||||
|
if len(pred_bboxes_rows) == 1:
|
||||||
|
match_bboxes_ready.extend(pred_bboxes_rows.pop(0))
|
||||||
|
if len(match_bboxes_ready) == 0 and len(pred_bboxes_rows) > 0:
|
||||||
|
match_bboxes_ready.extend(pred_bboxes_rows.pop(0))
|
||||||
|
if len(match_bboxes_ready) == 0 and len(pred_bboxes_rows) == 0:
|
||||||
|
break
|
||||||
|
#print(match_bboxes_ready)
|
||||||
|
for j, pred_box in enumerate(match_bboxes_ready):
|
||||||
|
distances.append(distance(gt_box, pred_box))
|
||||||
|
index = pred_distances.index(min(distances))
|
||||||
|
#print(gt_box, index)
|
||||||
|
#match_bboxes_ready.pop(distances.index(min(distances)))
|
||||||
|
print(gt_box, match_bboxes_ready[distances.index(min(distances))])
|
||||||
|
if index not in matched.keys():
|
||||||
|
matched[index] = [i]
|
||||||
|
else:
|
||||||
|
matched[index].append(i)
|
||||||
|
pre_bbox = gt_box
|
||||||
|
return matched
|
||||||
|
|
@ -0,0 +1,139 @@
|
||||||
|
# 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(os.path.abspath(__file__))
|
||||||
|
sys.path.append(__dir__)
|
||||||
|
sys.path.append(os.path.abspath(os.path.join(__dir__, '../..')))
|
||||||
|
|
||||||
|
os.environ["FLAGS_allocator_strategy"] = 'auto_growth'
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
import math
|
||||||
|
import time
|
||||||
|
import traceback
|
||||||
|
import paddle
|
||||||
|
|
||||||
|
import tools.infer.utility as utility
|
||||||
|
from ppocr.data import create_operators, transform
|
||||||
|
from ppocr.postprocess import build_post_process
|
||||||
|
from ppocr.utils.logging import get_logger
|
||||||
|
from ppocr.utils.utility import get_image_file_list, check_and_read_gif
|
||||||
|
from test.utility import parse_args
|
||||||
|
|
||||||
|
logger = get_logger()
|
||||||
|
|
||||||
|
|
||||||
|
class TableStructurer(object):
|
||||||
|
def __init__(self, args):
|
||||||
|
pre_process_list = [{
|
||||||
|
'ResizeTableImage': {
|
||||||
|
'max_len': args.structure_max_len
|
||||||
|
}
|
||||||
|
}, {
|
||||||
|
'NormalizeImage': {
|
||||||
|
'std': [0.229, 0.224, 0.225],
|
||||||
|
'mean': [0.485, 0.456, 0.406],
|
||||||
|
'scale': '1./255.',
|
||||||
|
'order': 'hwc'
|
||||||
|
}
|
||||||
|
}, {
|
||||||
|
'PaddingTableImage': None
|
||||||
|
}, {
|
||||||
|
'ToCHWImage': None
|
||||||
|
}, {
|
||||||
|
'KeepKeys': {
|
||||||
|
'keep_keys': ['image']
|
||||||
|
}
|
||||||
|
}]
|
||||||
|
postprocess_params = {
|
||||||
|
'name': 'TableLabelDecode',
|
||||||
|
"character_type": args.structure_char_type,
|
||||||
|
"character_dict_path": args.structure_char_dict_path,
|
||||||
|
}
|
||||||
|
|
||||||
|
self.preprocess_op = create_operators(pre_process_list)
|
||||||
|
self.postprocess_op = build_post_process(postprocess_params)
|
||||||
|
self.predictor, self.input_tensor, self.output_tensors, self.config = \
|
||||||
|
utility.create_predictor(args, 'structure', logger)
|
||||||
|
|
||||||
|
def __call__(self, img):
|
||||||
|
ori_im = img.copy()
|
||||||
|
data = {'image': img}
|
||||||
|
data = transform(data, self.preprocess_op)
|
||||||
|
img = data[0]
|
||||||
|
if img is None:
|
||||||
|
return None, 0
|
||||||
|
img = np.expand_dims(img, axis=0)
|
||||||
|
img = img.copy()
|
||||||
|
starttime = time.time()
|
||||||
|
|
||||||
|
self.input_tensor.copy_from_cpu(img)
|
||||||
|
self.predictor.run()
|
||||||
|
outputs = []
|
||||||
|
for output_tensor in self.output_tensors:
|
||||||
|
output = output_tensor.copy_to_cpu()
|
||||||
|
outputs.append(output)
|
||||||
|
|
||||||
|
preds = {}
|
||||||
|
preds['structure_probs'] = outputs[1]
|
||||||
|
preds['loc_preds'] = outputs[0]
|
||||||
|
|
||||||
|
post_result = self.postprocess_op(preds)
|
||||||
|
|
||||||
|
structure_str_list = post_result['structure_str_list']
|
||||||
|
res_loc = post_result['res_loc']
|
||||||
|
imgh, imgw = ori_im.shape[0:2]
|
||||||
|
res_loc_final = []
|
||||||
|
for rno in range(len(res_loc[0])):
|
||||||
|
x0, y0, x1, y1 = res_loc[0][rno]
|
||||||
|
left = max(int(imgw * x0), 0)
|
||||||
|
top = max(int(imgh * y0), 0)
|
||||||
|
right = min(int(imgw * x1), imgw - 1)
|
||||||
|
bottom = min(int(imgh * y1), imgh - 1)
|
||||||
|
res_loc_final.append([left, top, right, bottom])
|
||||||
|
|
||||||
|
structure_str_list = structure_str_list[0][:-1]
|
||||||
|
structure_str_list = ['<html>', '<body>', '<table>'] + structure_str_list + ['</table>', '</body>', '</html>']
|
||||||
|
|
||||||
|
elapse = time.time() - starttime
|
||||||
|
return (structure_str_list, res_loc_final), elapse
|
||||||
|
|
||||||
|
|
||||||
|
def main(args):
|
||||||
|
image_file_list = get_image_file_list(args.image_dir)
|
||||||
|
table_structurer = TableStructurer(args)
|
||||||
|
count = 0
|
||||||
|
total_time = 0
|
||||||
|
for image_file in image_file_list:
|
||||||
|
img, flag = check_and_read_gif(image_file)
|
||||||
|
if not flag:
|
||||||
|
img = cv2.imread(image_file)
|
||||||
|
if img is None:
|
||||||
|
logger.info("error in loading image:{}".format(image_file))
|
||||||
|
continue
|
||||||
|
structure_res, elapse = table_structurer(img)
|
||||||
|
|
||||||
|
logger.info("result: {}".format(structure_res))
|
||||||
|
|
||||||
|
if count > 0:
|
||||||
|
total_time += elapse
|
||||||
|
count += 1
|
||||||
|
logger.info("Predict time of {}: {}".format(image_file, elapse))
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main(parse_args())
|
||||||
|
|
@ -0,0 +1,221 @@
|
||||||
|
# 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
|
||||||
|
import subprocess
|
||||||
|
|
||||||
|
__dir__ = os.path.dirname(os.path.abspath(__file__))
|
||||||
|
sys.path.append(__dir__)
|
||||||
|
sys.path.append(os.path.abspath(os.path.join(__dir__, '..')))
|
||||||
|
sys.path.append(os.path.abspath(os.path.join(__dir__, '../..')))
|
||||||
|
|
||||||
|
os.environ["FLAGS_allocator_strategy"] = 'auto_growth'
|
||||||
|
import cv2
|
||||||
|
import copy
|
||||||
|
import numpy as np
|
||||||
|
import time
|
||||||
|
import tools.infer.predict_rec as predict_rec
|
||||||
|
import tools.infer.predict_det as predict_det
|
||||||
|
from ppocr.utils.utility import get_image_file_list, check_and_read_gif
|
||||||
|
from ppocr.utils.logging import get_logger
|
||||||
|
from test.table.matcher import distance, compute_iou
|
||||||
|
from test.utility import parse_args
|
||||||
|
import test.table.predict_structure as predict_strture
|
||||||
|
|
||||||
|
logger = get_logger()
|
||||||
|
|
||||||
|
|
||||||
|
def expand(pix, det_box, shape):
|
||||||
|
x0, y0, x1, y1 = det_box
|
||||||
|
# print(shape)
|
||||||
|
h, w, c = shape
|
||||||
|
tmp_x0 = x0 - pix
|
||||||
|
tmp_x1 = x1 + pix
|
||||||
|
tmp_y0 = y0 - pix
|
||||||
|
tmp_y1 = y1 + pix
|
||||||
|
x0_ = tmp_x0 if tmp_x0 >= 0 else 0
|
||||||
|
x1_ = tmp_x1 if tmp_x1 <= w else w
|
||||||
|
y0_ = tmp_y0 if tmp_y0 >= 0 else 0
|
||||||
|
y1_ = tmp_y1 if tmp_y1 <= h else h
|
||||||
|
return x0_, y0_, x1_, y1_
|
||||||
|
|
||||||
|
|
||||||
|
class TableSystem(object):
|
||||||
|
def __init__(self, args, text_detector=None, text_recognizer=None):
|
||||||
|
self.text_detector = predict_det.TextDetector(args) if text_detector is None else text_detector
|
||||||
|
self.text_recognizer = predict_rec.TextRecognizer(args) if text_recognizer is None else text_recognizer
|
||||||
|
self.table_structurer = predict_strture.TableStructurer(args)
|
||||||
|
|
||||||
|
def __call__(self, img):
|
||||||
|
ori_im = img.copy()
|
||||||
|
structure_res, elapse = self.table_structurer(copy.deepcopy(img))
|
||||||
|
dt_boxes, elapse = self.text_detector(copy.deepcopy(img))
|
||||||
|
dt_boxes = sorted_boxes(dt_boxes)
|
||||||
|
|
||||||
|
r_boxes = []
|
||||||
|
for box in dt_boxes:
|
||||||
|
x_min = box[:, 0].min() - 1
|
||||||
|
x_max = box[:, 0].max() + 1
|
||||||
|
y_min = box[:, 1].min() - 1
|
||||||
|
y_max = box[:, 1].max() + 1
|
||||||
|
box = [x_min, y_min, x_max, y_max]
|
||||||
|
r_boxes.append(box)
|
||||||
|
dt_boxes = np.array(r_boxes)
|
||||||
|
|
||||||
|
logger.debug("dt_boxes num : {}, elapse : {}".format(
|
||||||
|
len(dt_boxes), elapse))
|
||||||
|
if dt_boxes is None:
|
||||||
|
return None, None
|
||||||
|
img_crop_list = []
|
||||||
|
|
||||||
|
for i in range(len(dt_boxes)):
|
||||||
|
det_box = dt_boxes[i]
|
||||||
|
x0, y0, x1, y1 = expand(2, det_box, ori_im.shape)
|
||||||
|
text_rect = ori_im[int(y0):int(y1), int(x0):int(x1), :]
|
||||||
|
img_crop_list.append(text_rect)
|
||||||
|
rec_res, elapse = self.text_recognizer(img_crop_list)
|
||||||
|
logger.debug("rec_res num : {}, elapse : {}".format(
|
||||||
|
len(rec_res), elapse))
|
||||||
|
|
||||||
|
pred_html, pred = self.rebuild_table(structure_res, dt_boxes, rec_res)
|
||||||
|
return pred_html
|
||||||
|
|
||||||
|
def rebuild_table(self, structure_res, dt_boxes, rec_res):
|
||||||
|
pred_structures, pred_bboxes = structure_res
|
||||||
|
matched_index = self.match_result(dt_boxes, pred_bboxes)
|
||||||
|
pred_html, pred = self.get_pred_html(pred_structures, matched_index, rec_res)
|
||||||
|
return pred_html, pred
|
||||||
|
|
||||||
|
def match_result(self, dt_boxes, pred_bboxes):
|
||||||
|
matched = {}
|
||||||
|
for i, gt_box in enumerate(dt_boxes):
|
||||||
|
# gt_box = [np.min(gt_box[:, 0]), np.min(gt_box[:, 1]), np.max(gt_box[:, 0]), np.max(gt_box[:, 1])]
|
||||||
|
distances = []
|
||||||
|
for j, pred_box in enumerate(pred_bboxes):
|
||||||
|
distances.append(
|
||||||
|
(distance(gt_box, pred_box), 1. - compute_iou(gt_box, pred_box))) # 获取两两cell之间的L1距离和 1- IOU
|
||||||
|
sorted_distances = distances.copy()
|
||||||
|
# 根据距离和IOU挑选最"近"的cell
|
||||||
|
sorted_distances = sorted(sorted_distances, key=lambda item: (item[1], item[0]))
|
||||||
|
if distances.index(sorted_distances[0]) not in matched.keys():
|
||||||
|
matched[distances.index(sorted_distances[0])] = [i]
|
||||||
|
else:
|
||||||
|
matched[distances.index(sorted_distances[0])].append(i)
|
||||||
|
return matched
|
||||||
|
|
||||||
|
def get_pred_html(self, pred_structures, matched_index, ocr_contents):
|
||||||
|
end_html = []
|
||||||
|
td_index = 0
|
||||||
|
for tag in pred_structures:
|
||||||
|
if '</td>' in tag:
|
||||||
|
if td_index in matched_index.keys():
|
||||||
|
b_with = False
|
||||||
|
if '<b>' in ocr_contents[matched_index[td_index][0]] and len(matched_index[td_index]) > 1:
|
||||||
|
b_with = True
|
||||||
|
end_html.extend('<b>')
|
||||||
|
for i, td_index_index in enumerate(matched_index[td_index]):
|
||||||
|
content = ocr_contents[td_index_index][0]
|
||||||
|
if len(matched_index[td_index]) > 1:
|
||||||
|
if len(content) == 0:
|
||||||
|
continue
|
||||||
|
if content[0] == ' ':
|
||||||
|
content = content[1:]
|
||||||
|
if '<b>' in content:
|
||||||
|
content = content[3:]
|
||||||
|
if '</b>' in content:
|
||||||
|
content = content[:-4]
|
||||||
|
if len(content) == 0:
|
||||||
|
continue
|
||||||
|
if i != len(matched_index[td_index]) - 1 and ' ' != content[-1]:
|
||||||
|
content += ' '
|
||||||
|
end_html.extend(content)
|
||||||
|
if b_with:
|
||||||
|
end_html.extend('</b>')
|
||||||
|
|
||||||
|
end_html.append(tag)
|
||||||
|
td_index += 1
|
||||||
|
else:
|
||||||
|
end_html.append(tag)
|
||||||
|
return ''.join(end_html), end_html
|
||||||
|
|
||||||
|
|
||||||
|
def sorted_boxes(dt_boxes):
|
||||||
|
"""
|
||||||
|
Sort text boxes in order from top to bottom, left to right
|
||||||
|
args:
|
||||||
|
dt_boxes(array):detected text boxes with shape [4, 2]
|
||||||
|
return:
|
||||||
|
sorted boxes(array) with shape [4, 2]
|
||||||
|
"""
|
||||||
|
num_boxes = dt_boxes.shape[0]
|
||||||
|
sorted_boxes = sorted(dt_boxes, key=lambda x: (x[0][1], x[0][0]))
|
||||||
|
_boxes = list(sorted_boxes)
|
||||||
|
|
||||||
|
for i in range(num_boxes - 1):
|
||||||
|
if abs(_boxes[i + 1][0][1] - _boxes[i][0][1]) < 10 and \
|
||||||
|
(_boxes[i + 1][0][0] < _boxes[i][0][0]):
|
||||||
|
tmp = _boxes[i]
|
||||||
|
_boxes[i] = _boxes[i + 1]
|
||||||
|
_boxes[i + 1] = tmp
|
||||||
|
return _boxes
|
||||||
|
|
||||||
|
|
||||||
|
def to_excel(html_table, excel_path):
|
||||||
|
from tablepyxl import tablepyxl
|
||||||
|
tablepyxl.document_to_xl(html_table, excel_path)
|
||||||
|
|
||||||
|
|
||||||
|
def main(args):
|
||||||
|
image_file_list = get_image_file_list(args.image_dir)
|
||||||
|
image_file_list = image_file_list[args.process_id::args.total_process_num]
|
||||||
|
os.makedirs(args.output, exist_ok=True)
|
||||||
|
|
||||||
|
text_sys = TableSystem(args)
|
||||||
|
img_num = len(image_file_list)
|
||||||
|
for i, image_file in enumerate(image_file_list):
|
||||||
|
logger.info("[{}/{}] {}".format(i, img_num, image_file))
|
||||||
|
img, flag = check_and_read_gif(image_file)
|
||||||
|
excel_path = os.path.join(args.output, os.path.basename(image_file).split('.')[0] + '.xlsx')
|
||||||
|
if not flag:
|
||||||
|
img = cv2.imread(image_file)
|
||||||
|
if img is None:
|
||||||
|
logger.error("error in loading image:{}".format(image_file))
|
||||||
|
continue
|
||||||
|
starttime = time.time()
|
||||||
|
pred_html = text_sys(img)
|
||||||
|
|
||||||
|
to_excel(pred_html, excel_path)
|
||||||
|
logger.info('excel saved to {}'.format(excel_path))
|
||||||
|
logger.info(pred_html)
|
||||||
|
elapse = time.time() - starttime
|
||||||
|
logger.info("Predict time : {:.3f}s".format(elapse))
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
args = parse_args()
|
||||||
|
if args.use_mp:
|
||||||
|
p_list = []
|
||||||
|
total_process_num = args.total_process_num
|
||||||
|
for process_id in range(total_process_num):
|
||||||
|
cmd = [sys.executable, "-u"] + sys.argv + [
|
||||||
|
"--process_id={}".format(process_id),
|
||||||
|
"--use_mp={}".format(False)
|
||||||
|
]
|
||||||
|
p = subprocess.Popen(cmd, stdout=sys.stdout, stderr=sys.stdout)
|
||||||
|
p_list.append(p)
|
||||||
|
for p in p_list:
|
||||||
|
p.wait()
|
||||||
|
else:
|
||||||
|
main(args)
|
||||||
|
|
@ -0,0 +1,16 @@
|
||||||
|
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||||
|
#
|
||||||
|
# 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.
|
||||||
|
|
||||||
|
__all__ = ['TEDS']
|
||||||
|
from .table_metric import TEDS
|
||||||
|
|
@ -0,0 +1,51 @@
|
||||||
|
from tqdm import tqdm
|
||||||
|
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||||||
|
|
||||||
|
|
||||||
|
def parallel_process(array, function, n_jobs=16, use_kwargs=False, front_num=0):
|
||||||
|
"""
|
||||||
|
A parallel version of the map function with a progress bar.
|
||||||
|
Args:
|
||||||
|
array (array-like): An array to iterate over.
|
||||||
|
function (function): A python function to apply to the elements of array
|
||||||
|
n_jobs (int, default=16): The number of cores to use
|
||||||
|
use_kwargs (boolean, default=False): Whether to consider the elements of array as dictionaries of
|
||||||
|
keyword arguments to function
|
||||||
|
front_num (int, default=3): The number of iterations to run serially before kicking off the parallel job.
|
||||||
|
Useful for catching bugs
|
||||||
|
Returns:
|
||||||
|
[function(array[0]), function(array[1]), ...]
|
||||||
|
"""
|
||||||
|
# We run the first few iterations serially to catch bugs
|
||||||
|
if front_num > 0:
|
||||||
|
front = [function(**a) if use_kwargs else function(a)
|
||||||
|
for a in array[:front_num]]
|
||||||
|
else:
|
||||||
|
front = []
|
||||||
|
# If we set n_jobs to 1, just run a list comprehension. This is useful for benchmarking and debugging.
|
||||||
|
if n_jobs == 1:
|
||||||
|
return front + [function(**a) if use_kwargs else function(a) for a in tqdm(array[front_num:])]
|
||||||
|
# Assemble the workers
|
||||||
|
with ProcessPoolExecutor(max_workers=n_jobs) as pool:
|
||||||
|
# Pass the elements of array into function
|
||||||
|
if use_kwargs:
|
||||||
|
futures = [pool.submit(function, **a) for a in array[front_num:]]
|
||||||
|
else:
|
||||||
|
futures = [pool.submit(function, a) for a in array[front_num:]]
|
||||||
|
kwargs = {
|
||||||
|
'total': len(futures),
|
||||||
|
'unit': 'it',
|
||||||
|
'unit_scale': True,
|
||||||
|
'leave': True
|
||||||
|
}
|
||||||
|
# Print out the progress as tasks complete
|
||||||
|
for f in tqdm(as_completed(futures), **kwargs):
|
||||||
|
pass
|
||||||
|
out = []
|
||||||
|
# Get the results from the futures.
|
||||||
|
for i, future in tqdm(enumerate(futures)):
|
||||||
|
try:
|
||||||
|
out.append(future.result())
|
||||||
|
except Exception as e:
|
||||||
|
out.append(e)
|
||||||
|
return front + out
|
||||||
|
|
@ -0,0 +1,247 @@
|
||||||
|
# Copyright 2020 IBM
|
||||||
|
# Author: peter.zhong@au1.ibm.com
|
||||||
|
#
|
||||||
|
# This is free software; you can redistribute it and/or modify
|
||||||
|
# it under the terms of the Apache 2.0 License.
|
||||||
|
#
|
||||||
|
# This software is distributed in the hope that it will be useful,
|
||||||
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||||
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||||
|
# Apache 2.0 License for more details.
|
||||||
|
|
||||||
|
import distance
|
||||||
|
from apted import APTED, Config
|
||||||
|
from apted.helpers import Tree
|
||||||
|
from lxml import etree, html
|
||||||
|
from collections import deque
|
||||||
|
from .parallel import parallel_process
|
||||||
|
from tqdm import tqdm
|
||||||
|
|
||||||
|
|
||||||
|
class TableTree(Tree):
|
||||||
|
def __init__(self, tag, colspan=None, rowspan=None, content=None, *children):
|
||||||
|
self.tag = tag
|
||||||
|
self.colspan = colspan
|
||||||
|
self.rowspan = rowspan
|
||||||
|
self.content = content
|
||||||
|
self.children = list(children)
|
||||||
|
|
||||||
|
def bracket(self):
|
||||||
|
"""Show tree using brackets notation"""
|
||||||
|
if self.tag == 'td':
|
||||||
|
result = '"tag": %s, "colspan": %d, "rowspan": %d, "text": %s' % \
|
||||||
|
(self.tag, self.colspan, self.rowspan, self.content)
|
||||||
|
else:
|
||||||
|
result = '"tag": %s' % self.tag
|
||||||
|
for child in self.children:
|
||||||
|
result += child.bracket()
|
||||||
|
return "{{{}}}".format(result)
|
||||||
|
|
||||||
|
|
||||||
|
class CustomConfig(Config):
|
||||||
|
@staticmethod
|
||||||
|
def maximum(*sequences):
|
||||||
|
"""Get maximum possible value
|
||||||
|
"""
|
||||||
|
return max(map(len, sequences))
|
||||||
|
|
||||||
|
def normalized_distance(self, *sequences):
|
||||||
|
"""Get distance from 0 to 1
|
||||||
|
"""
|
||||||
|
return float(distance.levenshtein(*sequences)) / self.maximum(*sequences)
|
||||||
|
|
||||||
|
def rename(self, node1, node2):
|
||||||
|
"""Compares attributes of trees"""
|
||||||
|
#print(node1.tag)
|
||||||
|
if (node1.tag != node2.tag) or (node1.colspan != node2.colspan) or (node1.rowspan != node2.rowspan):
|
||||||
|
return 1.
|
||||||
|
if node1.tag == 'td':
|
||||||
|
if node1.content or node2.content:
|
||||||
|
#print(node1.content, )
|
||||||
|
return self.normalized_distance(node1.content, node2.content)
|
||||||
|
return 0.
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
class CustomConfig_del_short(Config):
|
||||||
|
@staticmethod
|
||||||
|
def maximum(*sequences):
|
||||||
|
"""Get maximum possible value
|
||||||
|
"""
|
||||||
|
return max(map(len, sequences))
|
||||||
|
|
||||||
|
def normalized_distance(self, *sequences):
|
||||||
|
"""Get distance from 0 to 1
|
||||||
|
"""
|
||||||
|
return float(distance.levenshtein(*sequences)) / self.maximum(*sequences)
|
||||||
|
|
||||||
|
def rename(self, node1, node2):
|
||||||
|
"""Compares attributes of trees"""
|
||||||
|
if (node1.tag != node2.tag) or (node1.colspan != node2.colspan) or (node1.rowspan != node2.rowspan):
|
||||||
|
return 1.
|
||||||
|
if node1.tag == 'td':
|
||||||
|
if node1.content or node2.content:
|
||||||
|
#print('before')
|
||||||
|
#print(node1.content, node2.content)
|
||||||
|
#print('after')
|
||||||
|
node1_content = node1.content
|
||||||
|
node2_content = node2.content
|
||||||
|
if len(node1_content) < 3:
|
||||||
|
node1_content = ['####']
|
||||||
|
if len(node2_content) < 3:
|
||||||
|
node2_content = ['####']
|
||||||
|
return self.normalized_distance(node1_content, node2_content)
|
||||||
|
return 0.
|
||||||
|
|
||||||
|
class CustomConfig_del_block(Config):
|
||||||
|
@staticmethod
|
||||||
|
def maximum(*sequences):
|
||||||
|
"""Get maximum possible value
|
||||||
|
"""
|
||||||
|
return max(map(len, sequences))
|
||||||
|
|
||||||
|
def normalized_distance(self, *sequences):
|
||||||
|
"""Get distance from 0 to 1
|
||||||
|
"""
|
||||||
|
return float(distance.levenshtein(*sequences)) / self.maximum(*sequences)
|
||||||
|
|
||||||
|
def rename(self, node1, node2):
|
||||||
|
"""Compares attributes of trees"""
|
||||||
|
if (node1.tag != node2.tag) or (node1.colspan != node2.colspan) or (node1.rowspan != node2.rowspan):
|
||||||
|
return 1.
|
||||||
|
if node1.tag == 'td':
|
||||||
|
if node1.content or node2.content:
|
||||||
|
|
||||||
|
node1_content = node1.content
|
||||||
|
node2_content = node2.content
|
||||||
|
while ' ' in node1_content:
|
||||||
|
print(node1_content.index(' '))
|
||||||
|
node1_content.pop(node1_content.index(' '))
|
||||||
|
while ' ' in node2_content:
|
||||||
|
print(node2_content.index(' '))
|
||||||
|
node2_content.pop(node2_content.index(' '))
|
||||||
|
return self.normalized_distance(node1_content, node2_content)
|
||||||
|
return 0.
|
||||||
|
|
||||||
|
class TEDS(object):
|
||||||
|
''' Tree Edit Distance basead Similarity
|
||||||
|
'''
|
||||||
|
|
||||||
|
def __init__(self, structure_only=False, n_jobs=1, ignore_nodes=None):
|
||||||
|
assert isinstance(n_jobs, int) and (
|
||||||
|
n_jobs >= 1), 'n_jobs must be an integer greather than 1'
|
||||||
|
self.structure_only = structure_only
|
||||||
|
self.n_jobs = n_jobs
|
||||||
|
self.ignore_nodes = ignore_nodes
|
||||||
|
self.__tokens__ = []
|
||||||
|
|
||||||
|
def tokenize(self, node):
|
||||||
|
''' Tokenizes table cells
|
||||||
|
'''
|
||||||
|
self.__tokens__.append('<%s>' % node.tag)
|
||||||
|
if node.text is not None:
|
||||||
|
self.__tokens__ += list(node.text)
|
||||||
|
for n in node.getchildren():
|
||||||
|
self.tokenize(n)
|
||||||
|
if node.tag != 'unk':
|
||||||
|
self.__tokens__.append('</%s>' % node.tag)
|
||||||
|
if node.tag != 'td' and node.tail is not None:
|
||||||
|
self.__tokens__ += list(node.tail)
|
||||||
|
|
||||||
|
def load_html_tree(self, node, parent=None):
|
||||||
|
''' Converts HTML tree to the format required by apted
|
||||||
|
'''
|
||||||
|
global __tokens__
|
||||||
|
if node.tag == 'td':
|
||||||
|
if self.structure_only:
|
||||||
|
cell = []
|
||||||
|
else:
|
||||||
|
self.__tokens__ = []
|
||||||
|
self.tokenize(node)
|
||||||
|
cell = self.__tokens__[1:-1].copy()
|
||||||
|
new_node = TableTree(node.tag,
|
||||||
|
int(node.attrib.get('colspan', '1')),
|
||||||
|
int(node.attrib.get('rowspan', '1')),
|
||||||
|
cell, *deque())
|
||||||
|
else:
|
||||||
|
new_node = TableTree(node.tag, None, None, None, *deque())
|
||||||
|
if parent is not None:
|
||||||
|
parent.children.append(new_node)
|
||||||
|
if node.tag != 'td':
|
||||||
|
for n in node.getchildren():
|
||||||
|
self.load_html_tree(n, new_node)
|
||||||
|
if parent is None:
|
||||||
|
return new_node
|
||||||
|
|
||||||
|
def evaluate(self, pred, true):
|
||||||
|
''' Computes TEDS score between the prediction and the ground truth of a
|
||||||
|
given sample
|
||||||
|
'''
|
||||||
|
if (not pred) or (not true):
|
||||||
|
return 0.0
|
||||||
|
parser = html.HTMLParser(remove_comments=True, encoding='utf-8')
|
||||||
|
pred = html.fromstring(pred, parser=parser)
|
||||||
|
true = html.fromstring(true, parser=parser)
|
||||||
|
if pred.xpath('body/table') and true.xpath('body/table'):
|
||||||
|
pred = pred.xpath('body/table')[0]
|
||||||
|
true = true.xpath('body/table')[0]
|
||||||
|
if self.ignore_nodes:
|
||||||
|
etree.strip_tags(pred, *self.ignore_nodes)
|
||||||
|
etree.strip_tags(true, *self.ignore_nodes)
|
||||||
|
n_nodes_pred = len(pred.xpath(".//*"))
|
||||||
|
n_nodes_true = len(true.xpath(".//*"))
|
||||||
|
n_nodes = max(n_nodes_pred, n_nodes_true)
|
||||||
|
tree_pred = self.load_html_tree(pred)
|
||||||
|
tree_true = self.load_html_tree(true)
|
||||||
|
distance = APTED(tree_pred, tree_true,
|
||||||
|
CustomConfig()).compute_edit_distance()
|
||||||
|
return 1.0 - (float(distance) / n_nodes)
|
||||||
|
else:
|
||||||
|
return 0.0
|
||||||
|
|
||||||
|
def batch_evaluate(self, pred_json, true_json):
|
||||||
|
''' Computes TEDS score between the prediction and the ground truth of
|
||||||
|
a batch of samples
|
||||||
|
@params pred_json: {'FILENAME': 'HTML CODE', ...}
|
||||||
|
@params true_json: {'FILENAME': {'html': 'HTML CODE'}, ...}
|
||||||
|
@output: {'FILENAME': 'TEDS SCORE', ...}
|
||||||
|
'''
|
||||||
|
samples = true_json.keys()
|
||||||
|
if self.n_jobs == 1:
|
||||||
|
scores = [self.evaluate(pred_json.get(
|
||||||
|
filename, ''), true_json[filename]['html']) for filename in tqdm(samples)]
|
||||||
|
else:
|
||||||
|
inputs = [{'pred': pred_json.get(
|
||||||
|
filename, ''), 'true': true_json[filename]['html']} for filename in samples]
|
||||||
|
scores = parallel_process(
|
||||||
|
inputs, self.evaluate, use_kwargs=True, n_jobs=self.n_jobs, front_num=1)
|
||||||
|
scores = dict(zip(samples, scores))
|
||||||
|
return scores
|
||||||
|
|
||||||
|
def batch_evaluate_html(self, pred_htmls, true_htmls):
|
||||||
|
''' Computes TEDS score between the prediction and the ground truth of
|
||||||
|
a batch of samples
|
||||||
|
'''
|
||||||
|
if self.n_jobs == 1:
|
||||||
|
scores = [self.evaluate(pred_html, true_html) for (
|
||||||
|
pred_html, true_html) in zip(pred_htmls, true_htmls)]
|
||||||
|
else:
|
||||||
|
inputs = [{"pred": pred_html, "true": true_html} for(
|
||||||
|
pred_html, true_html) in zip(pred_htmls, true_htmls)]
|
||||||
|
|
||||||
|
scores = parallel_process(
|
||||||
|
inputs, self.evaluate, use_kwargs=True, n_jobs=self.n_jobs, front_num=1)
|
||||||
|
return scores
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
import json
|
||||||
|
import pprint
|
||||||
|
with open('sample_pred.json') as fp:
|
||||||
|
pred_json = json.load(fp)
|
||||||
|
with open('sample_gt.json') as fp:
|
||||||
|
true_json = json.load(fp)
|
||||||
|
teds = TEDS(n_jobs=4)
|
||||||
|
scores = teds.batch_evaluate(pred_json, true_json)
|
||||||
|
pp = pprint.PrettyPrinter()
|
||||||
|
pp.pprint(scores)
|
||||||
|
|
@ -0,0 +1,13 @@
|
||||||
|
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||||
|
#
|
||||||
|
# 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.
|
||||||
|
|
@ -0,0 +1,283 @@
|
||||||
|
# This is where we handle translating css styles into openpyxl styles
|
||||||
|
# and cascading those from parent to child in the dom.
|
||||||
|
|
||||||
|
from openpyxl.cell import cell
|
||||||
|
from openpyxl.styles import Font, Alignment, PatternFill, NamedStyle, Border, Side, Color
|
||||||
|
from openpyxl.styles.fills import FILL_SOLID
|
||||||
|
from openpyxl.styles.numbers import FORMAT_CURRENCY_USD_SIMPLE, FORMAT_PERCENTAGE
|
||||||
|
from openpyxl.styles.colors import BLACK
|
||||||
|
|
||||||
|
FORMAT_DATE_MMDDYYYY = 'mm/dd/yyyy'
|
||||||
|
|
||||||
|
|
||||||
|
def colormap(color):
|
||||||
|
"""
|
||||||
|
Convenience for looking up known colors
|
||||||
|
"""
|
||||||
|
cmap = {'black': BLACK}
|
||||||
|
return cmap.get(color, color)
|
||||||
|
|
||||||
|
|
||||||
|
def style_string_to_dict(style):
|
||||||
|
"""
|
||||||
|
Convert css style string to a python dictionary
|
||||||
|
"""
|
||||||
|
def clean_split(string, delim):
|
||||||
|
return (s.strip() for s in string.split(delim))
|
||||||
|
styles = [clean_split(s, ":") for s in style.split(";") if ":" in s]
|
||||||
|
return dict(styles)
|
||||||
|
|
||||||
|
|
||||||
|
def get_side(style, name):
|
||||||
|
return {'border_style': style.get('border-{}-style'.format(name)),
|
||||||
|
'color': colormap(style.get('border-{}-color'.format(name)))}
|
||||||
|
|
||||||
|
known_styles = {}
|
||||||
|
|
||||||
|
|
||||||
|
def style_dict_to_named_style(style_dict, number_format=None):
|
||||||
|
"""
|
||||||
|
Change css style (stored in a python dictionary) to openpyxl NamedStyle
|
||||||
|
"""
|
||||||
|
|
||||||
|
style_and_format_string = str({
|
||||||
|
'style_dict': style_dict,
|
||||||
|
'parent': style_dict.parent,
|
||||||
|
'number_format': number_format,
|
||||||
|
})
|
||||||
|
|
||||||
|
if style_and_format_string not in known_styles:
|
||||||
|
# Font
|
||||||
|
font = Font(bold=style_dict.get('font-weight') == 'bold',
|
||||||
|
color=style_dict.get_color('color', None),
|
||||||
|
size=style_dict.get('font-size'))
|
||||||
|
|
||||||
|
# Alignment
|
||||||
|
alignment = Alignment(horizontal=style_dict.get('text-align', 'general'),
|
||||||
|
vertical=style_dict.get('vertical-align'),
|
||||||
|
wrap_text=style_dict.get('white-space', 'nowrap') == 'normal')
|
||||||
|
|
||||||
|
# Fill
|
||||||
|
bg_color = style_dict.get_color('background-color')
|
||||||
|
fg_color = style_dict.get_color('foreground-color', Color())
|
||||||
|
fill_type = style_dict.get('fill-type')
|
||||||
|
if bg_color and bg_color != 'transparent':
|
||||||
|
fill = PatternFill(fill_type=fill_type or FILL_SOLID,
|
||||||
|
start_color=bg_color,
|
||||||
|
end_color=fg_color)
|
||||||
|
else:
|
||||||
|
fill = PatternFill()
|
||||||
|
|
||||||
|
# Border
|
||||||
|
border = Border(left=Side(**get_side(style_dict, 'left')),
|
||||||
|
right=Side(**get_side(style_dict, 'right')),
|
||||||
|
top=Side(**get_side(style_dict, 'top')),
|
||||||
|
bottom=Side(**get_side(style_dict, 'bottom')),
|
||||||
|
diagonal=Side(**get_side(style_dict, 'diagonal')),
|
||||||
|
diagonal_direction=None,
|
||||||
|
outline=Side(**get_side(style_dict, 'outline')),
|
||||||
|
vertical=None,
|
||||||
|
horizontal=None)
|
||||||
|
|
||||||
|
name = 'Style {}'.format(len(known_styles) + 1)
|
||||||
|
|
||||||
|
pyxl_style = NamedStyle(name=name, font=font, fill=fill, alignment=alignment, border=border,
|
||||||
|
number_format=number_format)
|
||||||
|
|
||||||
|
known_styles[style_and_format_string] = pyxl_style
|
||||||
|
|
||||||
|
return known_styles[style_and_format_string]
|
||||||
|
|
||||||
|
|
||||||
|
class StyleDict(dict):
|
||||||
|
"""
|
||||||
|
It's like a dictionary, but it looks for items in the parent dictionary
|
||||||
|
"""
|
||||||
|
def __init__(self, *args, **kwargs):
|
||||||
|
self.parent = kwargs.pop('parent', None)
|
||||||
|
super(StyleDict, self).__init__(*args, **kwargs)
|
||||||
|
|
||||||
|
def __getitem__(self, item):
|
||||||
|
if item in self:
|
||||||
|
return super(StyleDict, self).__getitem__(item)
|
||||||
|
elif self.parent:
|
||||||
|
return self.parent[item]
|
||||||
|
else:
|
||||||
|
raise KeyError('{} not found'.format(item))
|
||||||
|
|
||||||
|
def __hash__(self):
|
||||||
|
return hash(tuple([(k, self.get(k)) for k in self._keys()]))
|
||||||
|
|
||||||
|
# Yielding the keys avoids creating unnecessary data structures
|
||||||
|
# and happily works with both python2 and python3 where the
|
||||||
|
# .keys() method is a dictionary_view in python3 and a list in python2.
|
||||||
|
def _keys(self):
|
||||||
|
yielded = set()
|
||||||
|
for k in self.keys():
|
||||||
|
yielded.add(k)
|
||||||
|
yield k
|
||||||
|
if self.parent:
|
||||||
|
for k in self.parent._keys():
|
||||||
|
if k not in yielded:
|
||||||
|
yielded.add(k)
|
||||||
|
yield k
|
||||||
|
|
||||||
|
def get(self, k, d=None):
|
||||||
|
try:
|
||||||
|
return self[k]
|
||||||
|
except KeyError:
|
||||||
|
return d
|
||||||
|
|
||||||
|
def get_color(self, k, d=None):
|
||||||
|
"""
|
||||||
|
Strip leading # off colors if necessary
|
||||||
|
"""
|
||||||
|
color = self.get(k, d)
|
||||||
|
if hasattr(color, 'startswith') and color.startswith('#'):
|
||||||
|
color = color[1:]
|
||||||
|
if len(color) == 3: # Premailers reduces colors like #00ff00 to #0f0, openpyxl doesn't like that
|
||||||
|
color = ''.join(2 * c for c in color)
|
||||||
|
return color
|
||||||
|
|
||||||
|
|
||||||
|
class Element(object):
|
||||||
|
"""
|
||||||
|
Our base class for representing an html element along with a cascading style.
|
||||||
|
The element is created along with a parent so that the StyleDict that we store
|
||||||
|
can point to the parent's StyleDict.
|
||||||
|
"""
|
||||||
|
def __init__(self, element, parent=None):
|
||||||
|
self.element = element
|
||||||
|
self.number_format = None
|
||||||
|
parent_style = parent.style_dict if parent else None
|
||||||
|
self.style_dict = StyleDict(style_string_to_dict(element.get('style', '')), parent=parent_style)
|
||||||
|
self._style_cache = None
|
||||||
|
|
||||||
|
def style(self):
|
||||||
|
"""
|
||||||
|
Turn the css styles for this element into an openpyxl NamedStyle.
|
||||||
|
"""
|
||||||
|
if not self._style_cache:
|
||||||
|
self._style_cache = style_dict_to_named_style(self.style_dict, number_format=self.number_format)
|
||||||
|
return self._style_cache
|
||||||
|
|
||||||
|
def get_dimension(self, dimension_key):
|
||||||
|
"""
|
||||||
|
Extracts the dimension from the style dict of the Element and returns it as a float.
|
||||||
|
"""
|
||||||
|
dimension = self.style_dict.get(dimension_key)
|
||||||
|
if dimension:
|
||||||
|
if dimension[-2:] in ['px', 'em', 'pt', 'in', 'cm']:
|
||||||
|
dimension = dimension[:-2]
|
||||||
|
dimension = float(dimension)
|
||||||
|
return dimension
|
||||||
|
|
||||||
|
|
||||||
|
class Table(Element):
|
||||||
|
"""
|
||||||
|
The concrete implementations of Elements are semantically named for the types of elements we are interested in.
|
||||||
|
This defines a very concrete tree structure for html tables that we expect to deal with. I prefer this compared to
|
||||||
|
allowing Element to have an arbitrary number of children and dealing with an abstract element tree.
|
||||||
|
"""
|
||||||
|
def __init__(self, table):
|
||||||
|
"""
|
||||||
|
takes an html table object (from lxml)
|
||||||
|
"""
|
||||||
|
super(Table, self).__init__(table)
|
||||||
|
table_head = table.find('thead')
|
||||||
|
self.head = TableHead(table_head, parent=self) if table_head is not None else None
|
||||||
|
table_body = table.find('tbody')
|
||||||
|
self.body = TableBody(table_body if table_body is not None else table, parent=self)
|
||||||
|
|
||||||
|
|
||||||
|
class TableHead(Element):
|
||||||
|
"""
|
||||||
|
This class maps to the `<th>` element of the html table.
|
||||||
|
"""
|
||||||
|
def __init__(self, head, parent=None):
|
||||||
|
super(TableHead, self).__init__(head, parent=parent)
|
||||||
|
self.rows = [TableRow(tr, parent=self) for tr in head.findall('tr')]
|
||||||
|
|
||||||
|
|
||||||
|
class TableBody(Element):
|
||||||
|
"""
|
||||||
|
This class maps to the `<tbody>` element of the html table.
|
||||||
|
"""
|
||||||
|
def __init__(self, body, parent=None):
|
||||||
|
super(TableBody, self).__init__(body, parent=parent)
|
||||||
|
self.rows = [TableRow(tr, parent=self) for tr in body.findall('tr')]
|
||||||
|
|
||||||
|
|
||||||
|
class TableRow(Element):
|
||||||
|
"""
|
||||||
|
This class maps to the `<tr>` element of the html table.
|
||||||
|
"""
|
||||||
|
def __init__(self, tr, parent=None):
|
||||||
|
super(TableRow, self).__init__(tr, parent=parent)
|
||||||
|
self.cells = [TableCell(cell, parent=self) for cell in tr.findall('th') + tr.findall('td')]
|
||||||
|
|
||||||
|
|
||||||
|
def element_to_string(el):
|
||||||
|
return _element_to_string(el).strip()
|
||||||
|
|
||||||
|
|
||||||
|
def _element_to_string(el):
|
||||||
|
string = ''
|
||||||
|
|
||||||
|
for x in el.iterchildren():
|
||||||
|
string += '\n' + _element_to_string(x)
|
||||||
|
|
||||||
|
text = el.text.strip() if el.text else ''
|
||||||
|
tail = el.tail.strip() if el.tail else ''
|
||||||
|
|
||||||
|
return text + string + '\n' + tail
|
||||||
|
|
||||||
|
|
||||||
|
class TableCell(Element):
|
||||||
|
"""
|
||||||
|
This class maps to the `<td>` element of the html table.
|
||||||
|
"""
|
||||||
|
CELL_TYPES = {'TYPE_STRING', 'TYPE_FORMULA', 'TYPE_NUMERIC', 'TYPE_BOOL', 'TYPE_CURRENCY', 'TYPE_PERCENTAGE',
|
||||||
|
'TYPE_NULL', 'TYPE_INLINE', 'TYPE_ERROR', 'TYPE_FORMULA_CACHE_STRING', 'TYPE_INTEGER'}
|
||||||
|
|
||||||
|
def __init__(self, cell, parent=None):
|
||||||
|
super(TableCell, self).__init__(cell, parent=parent)
|
||||||
|
self.value = element_to_string(cell)
|
||||||
|
self.number_format = self.get_number_format()
|
||||||
|
|
||||||
|
def data_type(self):
|
||||||
|
cell_types = self.CELL_TYPES & set(self.element.get('class', '').split())
|
||||||
|
if cell_types:
|
||||||
|
if 'TYPE_FORMULA' in cell_types:
|
||||||
|
# Make sure TYPE_FORMULA takes precedence over the other classes in the set.
|
||||||
|
cell_type = 'TYPE_FORMULA'
|
||||||
|
elif cell_types & {'TYPE_CURRENCY', 'TYPE_INTEGER', 'TYPE_PERCENTAGE'}:
|
||||||
|
cell_type = 'TYPE_NUMERIC'
|
||||||
|
else:
|
||||||
|
cell_type = cell_types.pop()
|
||||||
|
else:
|
||||||
|
cell_type = 'TYPE_STRING'
|
||||||
|
return getattr(cell, cell_type)
|
||||||
|
|
||||||
|
def get_number_format(self):
|
||||||
|
if 'TYPE_CURRENCY' in self.element.get('class', '').split():
|
||||||
|
return FORMAT_CURRENCY_USD_SIMPLE
|
||||||
|
if 'TYPE_INTEGER' in self.element.get('class', '').split():
|
||||||
|
return '#,##0'
|
||||||
|
if 'TYPE_PERCENTAGE' in self.element.get('class', '').split():
|
||||||
|
return FORMAT_PERCENTAGE
|
||||||
|
if 'TYPE_DATE' in self.element.get('class', '').split():
|
||||||
|
return FORMAT_DATE_MMDDYYYY
|
||||||
|
if self.data_type() == cell.TYPE_NUMERIC:
|
||||||
|
try:
|
||||||
|
int(self.value)
|
||||||
|
except ValueError:
|
||||||
|
return '#,##0.##'
|
||||||
|
else:
|
||||||
|
return '#,##0'
|
||||||
|
|
||||||
|
def format(self, cell):
|
||||||
|
cell.style = self.style()
|
||||||
|
data_type = self.data_type()
|
||||||
|
if data_type:
|
||||||
|
cell.data_type = data_type
|
||||||
|
|
@ -0,0 +1,118 @@
|
||||||
|
# Do imports like python3 so our package works for 2 and 3
|
||||||
|
from __future__ import absolute_import
|
||||||
|
|
||||||
|
from lxml import html
|
||||||
|
from openpyxl import Workbook
|
||||||
|
from openpyxl.utils import get_column_letter
|
||||||
|
from premailer import Premailer
|
||||||
|
from tablepyxl.style import Table
|
||||||
|
|
||||||
|
|
||||||
|
def string_to_int(s):
|
||||||
|
if s.isdigit():
|
||||||
|
return int(s)
|
||||||
|
return 0
|
||||||
|
|
||||||
|
|
||||||
|
def get_Tables(doc):
|
||||||
|
tree = html.fromstring(doc)
|
||||||
|
comments = tree.xpath('//comment()')
|
||||||
|
for comment in comments:
|
||||||
|
comment.drop_tag()
|
||||||
|
return [Table(table) for table in tree.xpath('//table')]
|
||||||
|
|
||||||
|
|
||||||
|
def write_rows(worksheet, elem, row, column=1):
|
||||||
|
"""
|
||||||
|
Writes every tr child element of elem to a row in the worksheet
|
||||||
|
returns the next row after all rows are written
|
||||||
|
"""
|
||||||
|
from openpyxl.cell.cell import MergedCell
|
||||||
|
|
||||||
|
initial_column = column
|
||||||
|
for table_row in elem.rows:
|
||||||
|
for table_cell in table_row.cells:
|
||||||
|
cell = worksheet.cell(row=row, column=column)
|
||||||
|
while isinstance(cell, MergedCell):
|
||||||
|
column += 1
|
||||||
|
cell = worksheet.cell(row=row, column=column)
|
||||||
|
|
||||||
|
colspan = string_to_int(table_cell.element.get("colspan", "1"))
|
||||||
|
rowspan = string_to_int(table_cell.element.get("rowspan", "1"))
|
||||||
|
if rowspan > 1 or colspan > 1:
|
||||||
|
worksheet.merge_cells(start_row=row, start_column=column,
|
||||||
|
end_row=row + rowspan - 1, end_column=column + colspan - 1)
|
||||||
|
|
||||||
|
cell.value = table_cell.value
|
||||||
|
table_cell.format(cell)
|
||||||
|
min_width = table_cell.get_dimension('min-width')
|
||||||
|
max_width = table_cell.get_dimension('max-width')
|
||||||
|
|
||||||
|
if colspan == 1:
|
||||||
|
# Initially, when iterating for the first time through the loop, the width of all the cells is None.
|
||||||
|
# As we start filling in contents, the initial width of the cell (which can be retrieved by:
|
||||||
|
# worksheet.column_dimensions[get_column_letter(column)].width) is equal to the width of the previous
|
||||||
|
# cell in the same column (i.e. width of A2 = width of A1)
|
||||||
|
width = max(worksheet.column_dimensions[get_column_letter(column)].width or 0, len(table_cell.value) + 2)
|
||||||
|
if max_width and width > max_width:
|
||||||
|
width = max_width
|
||||||
|
elif min_width and width < min_width:
|
||||||
|
width = min_width
|
||||||
|
worksheet.column_dimensions[get_column_letter(column)].width = width
|
||||||
|
column += colspan
|
||||||
|
row += 1
|
||||||
|
column = initial_column
|
||||||
|
return row
|
||||||
|
|
||||||
|
|
||||||
|
def table_to_sheet(table, wb):
|
||||||
|
"""
|
||||||
|
Takes a table and workbook and writes the table to a new sheet.
|
||||||
|
The sheet title will be the same as the table attribute name.
|
||||||
|
"""
|
||||||
|
ws = wb.create_sheet(title=table.element.get('name'))
|
||||||
|
insert_table(table, ws, 1, 1)
|
||||||
|
|
||||||
|
|
||||||
|
def document_to_workbook(doc, wb=None, base_url=None):
|
||||||
|
"""
|
||||||
|
Takes a string representation of an html document and writes one sheet for
|
||||||
|
every table in the document.
|
||||||
|
The workbook is returned
|
||||||
|
"""
|
||||||
|
if not wb:
|
||||||
|
wb = Workbook()
|
||||||
|
wb.remove(wb.active)
|
||||||
|
|
||||||
|
inline_styles_doc = Premailer(doc, base_url=base_url, remove_classes=False).transform()
|
||||||
|
tables = get_Tables(inline_styles_doc)
|
||||||
|
|
||||||
|
for table in tables:
|
||||||
|
table_to_sheet(table, wb)
|
||||||
|
|
||||||
|
return wb
|
||||||
|
|
||||||
|
|
||||||
|
def document_to_xl(doc, filename, base_url=None):
|
||||||
|
"""
|
||||||
|
Takes a string representation of an html document and writes one sheet for
|
||||||
|
every table in the document. The workbook is written out to a file called filename
|
||||||
|
"""
|
||||||
|
wb = document_to_workbook(doc, base_url=base_url)
|
||||||
|
wb.save(filename)
|
||||||
|
|
||||||
|
|
||||||
|
def insert_table(table, worksheet, column, row):
|
||||||
|
if table.head:
|
||||||
|
row = write_rows(worksheet, table.head, row, column)
|
||||||
|
if table.body:
|
||||||
|
row = write_rows(worksheet, table.body, row, column)
|
||||||
|
|
||||||
|
|
||||||
|
def insert_table_at_cell(table, cell):
|
||||||
|
"""
|
||||||
|
Inserts a table at the location of an openpyxl Cell object.
|
||||||
|
"""
|
||||||
|
ws = cell.parent
|
||||||
|
column, row = cell.column, cell.row
|
||||||
|
insert_table(table, ws, column, row)
|
||||||
|
|
@ -0,0 +1,54 @@
|
||||||
|
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||||
|
#
|
||||||
|
# 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.
|
||||||
|
|
||||||
|
from PIL import Image
|
||||||
|
import numpy as np
|
||||||
|
from tools.infer.utility import draw_ocr_box_txt, init_args as infer_args
|
||||||
|
|
||||||
|
|
||||||
|
def init_args():
|
||||||
|
parser = infer_args()
|
||||||
|
|
||||||
|
# params for output
|
||||||
|
parser.add_argument("--output", type=str, default='./output/table')
|
||||||
|
# params for table structure
|
||||||
|
parser.add_argument("--structure_max_len", type=int, default=488)
|
||||||
|
parser.add_argument("--structure_model_dir", type=str)
|
||||||
|
parser.add_argument("--structure_char_type", type=str, default='en')
|
||||||
|
parser.add_argument("--structure_char_dict_path", type=str, default="../ppocr/utils/dict/table_structure_dict.txt")
|
||||||
|
|
||||||
|
return parser
|
||||||
|
|
||||||
|
|
||||||
|
def parse_args():
|
||||||
|
parser = init_args()
|
||||||
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
def draw_result(image, result, font_path):
|
||||||
|
if isinstance(image, np.ndarray):
|
||||||
|
image = Image.fromarray(image)
|
||||||
|
boxes, txts, scores = [], [], []
|
||||||
|
for region in result:
|
||||||
|
if region['type'] == 'Table':
|
||||||
|
pass
|
||||||
|
elif region['type'] == 'Figure':
|
||||||
|
pass
|
||||||
|
else:
|
||||||
|
for box, rec_res in zip(region['res'][0], region['res'][1]):
|
||||||
|
boxes.append(np.array(box).reshape(-1, 2))
|
||||||
|
txts.append(rec_res[0])
|
||||||
|
scores.append(rec_res[1])
|
||||||
|
im_show = draw_ocr_box_txt(image, boxes, txts, scores, font_path=font_path,drop_score=0)
|
||||||
|
return im_show
|
||||||
|
|
@ -43,7 +43,7 @@ class TextDetector(object):
|
||||||
pre_process_list = [{
|
pre_process_list = [{
|
||||||
'DetResizeForTest': {
|
'DetResizeForTest': {
|
||||||
'limit_side_len': args.det_limit_side_len,
|
'limit_side_len': args.det_limit_side_len,
|
||||||
'limit_type': args.det_limit_type
|
'limit_type': args.det_limit_type,
|
||||||
}
|
}
|
||||||
}, {
|
}, {
|
||||||
'NormalizeImage': {
|
'NormalizeImage': {
|
||||||
|
|
|
||||||
|
|
@ -24,6 +24,7 @@ import cv2
|
||||||
import copy
|
import copy
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import time
|
import time
|
||||||
|
import logging
|
||||||
from PIL import Image
|
from PIL import Image
|
||||||
import tools.infer.utility as utility
|
import tools.infer.utility as utility
|
||||||
import tools.infer.predict_rec as predict_rec
|
import tools.infer.predict_rec as predict_rec
|
||||||
|
|
@ -38,6 +39,9 @@ logger = get_logger()
|
||||||
|
|
||||||
class TextSystem(object):
|
class TextSystem(object):
|
||||||
def __init__(self, args):
|
def __init__(self, args):
|
||||||
|
if not args.show_log:
|
||||||
|
logger.setLevel(logging.INFO)
|
||||||
|
|
||||||
self.text_detector = predict_det.TextDetector(args)
|
self.text_detector = predict_det.TextDetector(args)
|
||||||
self.text_recognizer = predict_rec.TextRecognizer(args)
|
self.text_recognizer = predict_rec.TextRecognizer(args)
|
||||||
self.use_angle_cls = args.use_angle_cls
|
self.use_angle_cls = args.use_angle_cls
|
||||||
|
|
@ -88,7 +92,7 @@ class TextSystem(object):
|
||||||
ori_im = img.copy()
|
ori_im = img.copy()
|
||||||
dt_boxes, elapse = self.text_detector(img)
|
dt_boxes, elapse = self.text_detector(img)
|
||||||
|
|
||||||
logger.info("dt_boxes num : {}, elapse : {}".format(
|
logger.debug("dt_boxes num : {}, elapse : {}".format(
|
||||||
|
|
||||||
len(dt_boxes), elapse))
|
len(dt_boxes), elapse))
|
||||||
if dt_boxes is None:
|
if dt_boxes is None:
|
||||||
|
|
@ -104,11 +108,11 @@ class TextSystem(object):
|
||||||
if self.use_angle_cls and cls:
|
if self.use_angle_cls and cls:
|
||||||
img_crop_list, angle_list, elapse = self.text_classifier(
|
img_crop_list, angle_list, elapse = self.text_classifier(
|
||||||
img_crop_list)
|
img_crop_list)
|
||||||
logger.info("cls num : {}, elapse : {}".format(
|
logger.debug("cls num : {}, elapse : {}".format(
|
||||||
len(img_crop_list), elapse))
|
len(img_crop_list), elapse))
|
||||||
|
|
||||||
rec_res, elapse = self.text_recognizer(img_crop_list)
|
rec_res, elapse = self.text_recognizer(img_crop_list)
|
||||||
logger.info("rec_res num : {}, elapse : {}".format(
|
logger.debug("rec_res num : {}, elapse : {}".format(
|
||||||
len(rec_res), elapse))
|
len(rec_res), elapse))
|
||||||
# self.print_draw_crop_rec_res(img_crop_list, rec_res)
|
# self.print_draw_crop_rec_res(img_crop_list, rec_res)
|
||||||
filter_boxes, filter_rec_res = [], []
|
filter_boxes, filter_rec_res = [], []
|
||||||
|
|
|
||||||
|
|
@ -109,11 +109,12 @@ def init_args():
|
||||||
parser.add_argument("--use_mp", type=str2bool, default=False)
|
parser.add_argument("--use_mp", type=str2bool, default=False)
|
||||||
parser.add_argument("--total_process_num", type=int, default=1)
|
parser.add_argument("--total_process_num", type=int, default=1)
|
||||||
parser.add_argument("--process_id", type=int, default=0)
|
parser.add_argument("--process_id", type=int, default=0)
|
||||||
|
|
||||||
parser.add_argument("--benchmark", type=bool, default=False)
|
parser.add_argument("--benchmark", type=bool, default=False)
|
||||||
parser.add_argument("--save_log_path", type=str, default="./log_output/")
|
parser.add_argument("--save_log_path", type=str, default="./log_output/")
|
||||||
|
|
||||||
|
|
||||||
|
parser.add_argument("--show_log", type=str2bool, default=True)
|
||||||
return parser
|
return parser
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -199,6 +200,8 @@ def create_predictor(args, mode, logger):
|
||||||
model_dir = args.cls_model_dir
|
model_dir = args.cls_model_dir
|
||||||
elif mode == 'rec':
|
elif mode == 'rec':
|
||||||
model_dir = args.rec_model_dir
|
model_dir = args.rec_model_dir
|
||||||
|
elif mode == 'structure':
|
||||||
|
model_dir = args.structure_model_dir
|
||||||
else:
|
else:
|
||||||
model_dir = args.e2e_model_dir
|
model_dir = args.e2e_model_dir
|
||||||
|
|
||||||
|
|
@ -328,7 +331,9 @@ def create_predictor(args, mode, logger):
|
||||||
|
|
||||||
config.delete_pass("conv_transpose_eltwiseadd_bn_fuse_pass")
|
config.delete_pass("conv_transpose_eltwiseadd_bn_fuse_pass")
|
||||||
config.switch_use_feed_fetch_ops(False)
|
config.switch_use_feed_fetch_ops(False)
|
||||||
|
config.switch_ir_optim(True)
|
||||||
|
if mode == 'structure':
|
||||||
|
config.switch_ir_optim(False)
|
||||||
# create predictor
|
# create predictor
|
||||||
predictor = inference.create_predictor(config)
|
predictor = inference.create_predictor(config)
|
||||||
input_names = predictor.get_input_names()
|
input_names = predictor.get_input_names()
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue