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)

tools/infer/utility.py

@@ -112,36 +112,70 @@ def draw_text_det_res(dt_boxes, img_path):
     cv2.imwrite("./output/%s" % img_name_pure, src_im)
 
 
-def draw_ocr(image, boxes, txts, scores, draw_txt):
+def resize_img(img, input_size=600):
+    """
+    """
+    img = np.array(img)
+    im_shape = img.shape
+    im_size_min = np.min(im_shape[0:2])
+    im_size_max = np.max(im_shape[0:2])
+    im_scale = float(input_size) / float(im_size_max)
+    im = cv2.resize(img, None, None, fx=im_scale, fy=im_scale)
+    return im
+
+
+def draw_ocr(image, boxes, txts, scores, draw_txt=True, drop_score=0.5):
     from PIL import Image, ImageDraw, ImageFont
 
     w, h = image.size
     img = image.copy()
     draw = ImageDraw.Draw(img)
 
-    for (box, txt) in zip(boxes, txts):
-
+    for (box, score) in zip(boxes, scores):
+        if score < drop_score:
+            continue
         draw.line([(box[0][0], box[0][1]), (box[1][0], box[1][1])], fill='red')
         draw.line([(box[1][0], box[1][1]), (box[2][0], box[2][1])], fill='red')
         draw.line([(box[2][0], box[2][1]), (box[3][0], box[3][1])], fill='red')
         draw.line([(box[3][0], box[3][1]), (box[0][0], box[0][1])], fill='red')
+        draw.line(
+            [(box[0][0] - 1, box[0][1] + 1), (box[1][0] - 1, box[1][1] + 1)],
+            fill='red')
+        draw.line(
+            [(box[1][0] - 1, box[1][1] + 1), (box[2][0] - 1, box[2][1] + 1)],
+            fill='red')
+        draw.line(
+            [(box[2][0] - 1, box[2][1] + 1), (box[3][0] - 1, box[3][1] + 1)],
+            fill='red')
+        draw.line(
+            [(box[3][0] - 1, box[3][1] + 1), (box[0][0] - 1, box[0][1] + 1)],
+            fill='red')
 
     if draw_txt:
         txt_color = (0, 0, 0)
-
-        blank_img = np.ones(shape=[h, 800], dtype=np.int8) * 255
+        img = np.array(resize_img(img))
+        _h = img.shape[0]
+        blank_img = np.ones(shape=[_h, 600], dtype=np.int8) * 255
         blank_img = Image.fromarray(blank_img).convert("RGB")
         draw_txt = ImageDraw.Draw(blank_img)
 
-        font_size = 30
-        gap = 40 if h // len(txts) >= font_size else h // len(txts)
-
-        for i, txt in enumerate(txts):
+        font_size = 20
+        gap = 20
+        title = "index           text           score"
+        font = ImageFont.truetype(
+            "./doc/simfang.ttf", font_size, encoding="utf-8")
+
+        draw_txt.text((20, 0), title, txt_color, font=font)
+        count = 0
+        for idx, txt in enumerate(txts):
+            if scores[idx] < drop_score:
+                continue
             font = ImageFont.truetype(
-                "./doc/simfang.TTF", font_size, encoding="utf-8")
-            new_txt = str(i) + ':  ' + txt + '    ' + str(scores[i])
-            draw_txt.text((20, gap * (i + 1)), new_txt, txt_color, font=font)
-
+                "./doc/simfang.ttf", font_size, encoding="utf-8")
+            new_txt = str(count) + ':  ' + txt + '    ' + str(scores[count])
+            draw_txt.text(
+                (20, gap * (count + 1)), new_txt, txt_color, font=font)
+            count += 1
         img = np.concatenate([np.array(img), np.array(blank_img)], axis=1)
     return img