Merge pull request #15 from LDOUBLEV/fixocr

add visualize code to predict_eval

README.md

@@ -11,7 +11,7 @@ PaddleOCR旨在打造一套丰富、领先、且实用的OCR工具库，助力
 
 ## **超轻量级中文OCR体验**
 
-![](./doc/imgs_draw/11.jpg)
+![](doc/imgs_results/11.jpg)
 
 上图是超轻量级中文OCR模型效果展示，更多效果图请见文末[效果展示](#效果展示)。
 
@@ -97,14 +97,14 @@ PaddleOCR文本识别算法的训练和使用请参考文档教程中[文本识
 
 <a name="效果展示"></a>
 ## 效果展示
-![](./doc/imgs_draw/1.jpg)
-![](./doc/imgs_draw/7.jpg)
-![](./doc/imgs_draw/12.jpg)
-![](./doc/imgs_draw/4.jpg)
-![](./doc/imgs_draw/6.jpg)
-![](./doc/imgs_draw/9.jpg)
-![](./doc/imgs_draw/16.png)
-![](./doc/imgs_draw/22.jpg)
+![](doc/imgs_results/1.jpg)
+![](doc/imgs_results/7.jpg)
+![](doc/imgs_results/12.jpg)
+![](doc/imgs_results/4.jpg)
+![](doc/imgs_results/6.jpg)
+![](doc/imgs_results/9.jpg)
+![](doc/imgs_results/16.png)
+![](doc/imgs_results/22.jpg)
 
 
 ## 参考文献

tools/infer/predict_system.py

@@ -22,6 +22,10 @@ import copy
 import numpy as np
 import math
 import time
+from ppocr.utils.utility import get_image_file_list
+from PIL import Image
+from tools.infer.utility import draw_ocr
+import os
 
 
 class TextSystem(object):
@@ -99,8 +103,9 @@ def sorted_boxes(dt_boxes):
 
 if __name__ == "__main__":
     args = utility.parse_args()
-    image_file_list = utility.get_image_file_list(args.image_dir)
+    image_file_list = get_image_file_list(args.image_dir)
     text_sys = TextSystem(args)
+    is_visualize = True
     for image_file in image_file_list:
         img = cv2.imread(image_file)
         if img is None:
@@ -114,8 +119,22 @@ if __name__ == "__main__":
         dt_boxes_final = []
         for dno in range(dt_num):
             text, score = rec_res[dno]
-            if score >= 0:
+            if score >= 0.5:
                 text_str = "%s, %.3f" % (text, score)
                 print(text_str)
                 dt_boxes_final.append(dt_boxes[dno])
-        utility.draw_text_det_res(dt_boxes_final, image_file)
+
+        if is_visualize:
+            image = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
+            boxes = dt_boxes
+            txts = [rec_res[i][0] for i in range(len(rec_res))]
+            scores = [rec_res[i][1] for i in range(len(rec_res))]
+
+            draw_img = draw_ocr(
+                image, boxes, txts, scores, draw_txt=True, drop_score=0.5)
+            draw_img_save = "./doc/imgs_results/"
+            if not os.path.exists(draw_img_save):
+                os.makedirs(draw_img_save)
+            cv2.imwrite(
+                os.path.join(draw_img_save, os.path.basename(image_file)),
+                draw_img)