From 3f402caea7510dd022d3d9726823e37368b8a81c Mon Sep 17 00:00:00 2001
From: zhangxin <zhangxin19870504@163.com>
Date: Mon, 15 Jun 2020 18:40:01 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E4=B8=80=E4=B8=AA=E7=94=BB?=
 =?UTF-8?q?=E5=9B=BE=E5=87=BD=E6=95=B0=EF=BC=8C=E5=8F=AF=E4=BB=A5=E6=9B=B4?=
 =?UTF-8?q?=E6=96=B9=E4=BE=BF=E6=9B=B4=E7=9B=B4=E8=A7=82=E7=9A=84=E7=9C=8B?=
 =?UTF-8?q?=E6=A3=80=E6=B5=8B=E5=92=8C=E8=AF=86=E5=88=AB=E7=BB=93=E6=9E=9C?=
 =?UTF-8?q?=EF=BC=8C=E5=AF=B9=E6=A8=AA=E6=8E=92=E3=80=81=E7=AB=96=E6=8E=92?=
 =?UTF-8?q?=E6=96=87=E6=9C=AC=E9=83=BD=E5=8F=AF=E4=BB=A5=E6=98=BE=E7=A4=BA?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 tools/infer/predict_system.py |  6 ++++--
 tools/infer/utility.py        | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/tools/infer/predict_system.py b/tools/infer/predict_system.py
index 1019279a..ba8aaa7e 100755
--- a/tools/infer/predict_system.py
+++ b/tools/infer/predict_system.py
@@ -29,6 +29,7 @@ import time
 from ppocr.utils.utility import get_image_file_list
 from PIL import Image
 from tools.infer.utility import draw_ocr
+from tools.infer.utility import draw_ocr_box_txt
 
 
 class TextSystem(object):
@@ -133,8 +134,9 @@ if __name__ == "__main__":
             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 = draw_ocr(
+            #     image, boxes, txts, scores, draw_txt=True, drop_score=0.5)
+            draw_img = draw_ocr_box_txt(image, boxes, txts)
             draw_img_save = "./inference_results/"
             if not os.path.exists(draw_img_save):
                 os.makedirs(draw_img_save)
diff --git a/tools/infer/utility.py b/tools/infer/utility.py
index c3d3b445..9fbf3513 100755
--- a/tools/infer/utility.py
+++ b/tools/infer/utility.py
@@ -157,6 +157,38 @@ def draw_ocr(image, boxes, txts, scores, draw_txt=True, drop_score=0.5):
     return image
 
 
+def draw_ocr_box_txt(image, boxes, txts, drop_score=0.5):
+    img_show = Image.new('RGB', (image.width * 2, image.height), (255, 255, 255))
+    img_show.paste(image, (image.width, 0, image.width * 2, image.height))
+    img_show_bak = img_show.copy()
+    height, width = image.height, image.width
+
+    import random
+    random.seed(12345)
+    draw = ImageDraw.Draw(img_show)
+    for (box, txt) in zip(boxes, txts):
+        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
+        draw.polygon(box, outline=color)
+        draw.polygon([width + box[0][0], box[0][1], width + box[1][0], box[1][1], width + box[2][0], box[2][1], width + box[3][0], box[3][1]],
+                     fill=color)
+        box_height = math.sqrt((box[0][0] - box[3][0]) ** 2 + (box[0][1] - box[3][1]) ** 2)
+        box_width = math.sqrt((box[0][0] - box[1][0]) ** 2 + (box[0][1] - box[1][1]) ** 2)
+        if box_height > 2 * box_width:
+            font_size = max(int(box_width * 0.9), 10)
+            font = ImageFont.truetype("./doc/simfang.ttf", font_size, encoding="utf-8")
+            cur_y = box[0][1]
+            for c in txt:
+                char_size = font.getsize(c)
+                draw.text((box[0][0] + 3, cur_y), c, fill=(0, 0, 255), font=font)
+                cur_y += char_size[1]
+        else:
+            font_size = max(int(box_height * 0.8), 10)
+            font = ImageFont.truetype("./doc/simfang.ttf", font_size, encoding="utf-8")
+            draw.text(box[0], txt, fill=(255, 0, 0), font=font)
+    img_show = Image.blend(img_show_bak, img_show, 0.75)
+    return np.array(img_show)
+
+
 def str_count(s):
     """
     Count the number of Chinese characters,
-- 
GitLab