From d27360a9b8b638112587f9d23d1054a8333a2c57 Mon Sep 17 00:00:00 2001 From: WenmuZhou Date: Tue, 8 Jun 2021 14:16:32 +0800 Subject: [PATCH] move draw_result to utilitu.py --- ppstructure/paddlestructure.py | 31 +++++++------------------------ ppstructure/predict_system.py | 7 +++++-- ppstructure/utility.py | 23 +++++++++++++++++++++-- 3 files changed, 33 insertions(+), 28 deletions(-) diff --git a/ppstructure/paddlestructure.py b/ppstructure/paddlestructure.py index cf49fd99..686267e1 100644 --- a/ppstructure/paddlestructure.py +++ b/ppstructure/paddlestructure.py @@ -25,27 +25,23 @@ from pathlib import Path from ppocr.utils.logging import get_logger from ppstructure.predict_system import OCRSystem, save_res -from utility import init_args +from ppstructure.table.predict_table import to_excel +from ppstructure.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 -__all__ = ['PaddleStructure'] +__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/ch/ch_ppocr_mobile_v2.0_det_infer.tar', - }, - 'rec': { - 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_det_infer.tar', - }, - 'structure': { - 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_det_infer.tar', - }, + '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' + } @@ -143,21 +139,8 @@ def main(): img_name = os.path.basename(img_path).split('.')[0] logger.info('{}{}{}'.format('*' * 10, img_path, '*' * 10)) result = structure_engine(img_path) - save_res(result, args.output, os.path.basename(img_path).split('.')[0]) 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))) - -if __name__ == '__main__': - table_engine = PaddleStructure( - det_model_dir='/Users/zhoujun20/Desktop/工作相关/table/table_pr/PaddleOCR/inference/table/ch_ppocr_mobile_v2.0_table_det_infer', - rec_model_dir='/Users/zhoujun20/Desktop/工作相关/table/table_pr/PaddleOCR/inference/table/ch_ppocr_mobile_v2.0_table_rec_infer', - structure_model_dir='/Users/zhoujun20/Desktop/工作相关/table/table_pr/PaddleOCR/inference/table/ch_ppocr_mobile_v2.0_table_structure_infer', - output='/Users/zhoujun20/Desktop/工作相关/table/table_pr/PaddleOCR/output/table', - show_log=True) - img = cv2.imread('/Users/zhoujun20/Desktop/工作相关/table/table_pr/PaddleOCR/ppstructure/test_imgs/table_1.png') - result = table_engine(img) - for line in result: - print(line) diff --git a/ppstructure/predict_system.py b/ppstructure/predict_system.py index 87306eae..907548e7 100644 --- a/ppstructure/predict_system.py +++ b/ppstructure/predict_system.py @@ -31,7 +31,7 @@ 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 ppstructure.table.predict_table import TableSystem, to_excel -from ppstructure.utility import parse_args +from ppstructure.utility import parse_args,draw_result logger = get_logger() @@ -39,7 +39,8 @@ logger = get_logger() class OCRSystem(object): def __init__(self, args): args.det_pad = True - args.det_pad_size = 640 + args.det_pad_size = 960 + args.drop_score = 0 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", @@ -107,6 +108,8 @@ def main(args): 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)) diff --git a/ppstructure/utility.py b/ppstructure/utility.py index 57659920..8112b9ef 100644 --- a/ppstructure/utility.py +++ b/ppstructure/utility.py @@ -11,9 +11,10 @@ # 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 -from tools.infer.utility import str2bool, init_args as infer_args +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(): @@ -38,3 +39,21 @@ def init_args(): 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 \ No newline at end of file -- GitLab