move draw_result to utilitu.py

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)

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))

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