Merge pull request #210 from ZhangXinNan/zxdev

增加一个可视化函数
PaddlePaddle · Jun 17, 2020 · fea9454 · fea9454
2 parents 63221f8 + 5d63991
commit fea9454
Show file tree

Hide file tree

Showing 19 changed files with 39 additions and 0 deletions.
diff --git a/doc/imgs_results_vis2/1.jpg b/doc/imgs_results_vis2/1.jpg
diff --git a/doc/imgs_results_vis2/10.jpg b/doc/imgs_results_vis2/10.jpg
diff --git a/doc/imgs_results_vis2/11.jpg b/doc/imgs_results_vis2/11.jpg
diff --git a/doc/imgs_results_vis2/12.jpg b/doc/imgs_results_vis2/12.jpg
diff --git a/doc/imgs_results_vis2/13.png b/doc/imgs_results_vis2/13.png
diff --git a/doc/imgs_results_vis2/15.jpg b/doc/imgs_results_vis2/15.jpg
diff --git a/doc/imgs_results_vis2/16.png b/doc/imgs_results_vis2/16.png
diff --git a/doc/imgs_results_vis2/17.png b/doc/imgs_results_vis2/17.png
diff --git a/doc/imgs_results_vis2/2.jpg b/doc/imgs_results_vis2/2.jpg
diff --git a/doc/imgs_results_vis2/22.jpg b/doc/imgs_results_vis2/22.jpg
diff --git a/doc/imgs_results_vis2/3.jpg b/doc/imgs_results_vis2/3.jpg
diff --git a/doc/imgs_results_vis2/4.jpg b/doc/imgs_results_vis2/4.jpg
diff --git a/doc/imgs_results_vis2/5.jpg b/doc/imgs_results_vis2/5.jpg
diff --git a/doc/imgs_results_vis2/6.jpg b/doc/imgs_results_vis2/6.jpg
diff --git a/doc/imgs_results_vis2/7.jpg b/doc/imgs_results_vis2/7.jpg
diff --git a/doc/imgs_results_vis2/8.jpg b/doc/imgs_results_vis2/8.jpg
diff --git a/doc/imgs_results_vis2/9.jpg b/doc/imgs_results_vis2/9.jpg
diff --git a/tools/infer/predict_system.py b/tools/infer/predict_system.py
@@ -30,6 +30,7 @@
 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):

diff --git a/tools/infer/utility.py b/tools/infer/utility.py
@@ -157,6 +157,44 @@ def draw_ocr(image, boxes, txts, scores, draw_txt=True, drop_score=0.5):
  return image
 
 
+def draw_ocr_box_txt(image, boxes, txts):
+ h, w = image.height, image.width
+ img_left = image.copy()
+ img_right = Image.new('RGB', (w, h), (255, 255, 255))
+
+ import random
+ # 每次使用相同的随机种子 ，可以保证两次颜色一致
+ random.seed(0)
+ draw_left = ImageDraw.Draw(img_left)
+ draw_right = ImageDraw.Draw(img_right)
+ for (box, txt) in zip(boxes, txts):
+ color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
+ draw_left.polygon(box, fill=color)
+ draw_right.polygon([box[0][0], box[0][1],
+ box[1][0], box[1][1],
+ box[2][0], box[2][1],
+ box[3][0], box[3][1]], outline=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_right.text((box[0][0] + 3, cur_y), c, fill=(0, 0, 0), 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_right.text([box[0][0], box[0][1]], txt, fill=(0, 0, 0), font=font)
+ img_left = Image.blend(image, img_left, 0.5)
+ img_show = Image.new('RGB', (w * 2, h), (255, 255, 255))
+ img_show.paste(img_left, (0, 0, w, h))
+ img_show.paste(img_right, (w, 0, w * 2, h))
+ return np.array(img_show)
+
+
 def str_count(s):
  """
  Count the number of Chinese characters,