From 20466055b2dc463b715dcb5a5a3f7c4bcdcfdcb8 Mon Sep 17 00:00:00 2001
From: WenmuZhou <zjwenmu@gmail.com>
Date: Thu, 3 Jun 2021 20:09:05 +0800
Subject: [PATCH] add save_dir to args

---
 ppstructure/predict_system.py | 24 ++++++++++++++++--------
 1 file changed, 16 insertions(+), 8 deletions(-)

diff --git a/ppstructure/predict_system.py b/ppstructure/predict_system.py
index cd2ff0fb..f6852ab7 100644
--- a/ppstructure/predict_system.py
+++ b/ppstructure/predict_system.py
@@ -38,6 +38,8 @@ logger = get_logger()
 def parse_args():
     parser = utility.init_args()
 
+    # params for output
+    parser.add_argument("--table_output", type=str, default='output/table')
     # params for table structure
     parser.add_argument("--table_max_len", type=int, default=488)
     parser.add_argument("--table_max_text_length", type=int, default=100)
@@ -65,9 +67,9 @@ class OCRSystem():
         layout_res = self.table_layout(copy.deepcopy(img))
         for region in layout_res:
             x1, y1, x2, y2 = region['bbox']
-            roi_img = ori_im[y1:y2, x1:x2,:]
+            roi_img = ori_im[y1:y2, x1:x2, :]
             if region['label'] == 'table':
-                res = self.table_system(roi_img)
+                res = self.text_system(roi_img)
             else:
                 res = self.text_system(roi_img)
             region['res'] = res
@@ -77,15 +79,15 @@ class OCRSystem():
 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]
-    excel_save_folder = 'output/table'
-    os.makedirs(excel_save_folder, exist_ok=True)
+    save_folder = args.table_output
+    os.makedirs(save_folder, exist_ok=True)
 
     text_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)
-        imgname = os.path.basename(image_file).split('.')[0]
+        img_name = os.path.basename(image_file).split('.')[0]
         # excel_path = os.path.join(excel_save_folder, + '.xlsx')
         if not flag:
             img = cv2.imread(image_file)
@@ -95,11 +97,17 @@ def main(args):
         starttime = time.time()
         res = text_sys(img)
 
+        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['label'] == 'table':
-                # x1, y1, x2, y2 = region['bbox']
-                excel_path = os.path.join(excel_save_folder, '{}_{}.xlsx'.format(imgname,region['bbox']))
-                to_excel(region['res'],excel_path)
+                excel_path = os.path.join(excel_save_folder, '{}.xlsx'.format(region['bbox']))
+                to_excel(region['res'], excel_path)
+            else:
+                with open(os.path.join(excel_save_folder, 'res.txt'),'a',encoding='utf8') as f:
+                    for box, rec_res in zip(*region['res']):
+                        f.write('{}\t{}\n'.format(np.array(box).reshape(-1).tolist(), rec_res))
         logger.info(res)
         elapse = time.time() - starttime
         logger.info("Predict time : {:.3f}s".format(elapse))
-- 
GitLab