From 22e8cf8c6b6d93c3ae016d82603124eba27a1784 Mon Sep 17 00:00:00 2001 From: zhangxin Date: Mon, 15 Jun 2020 20:05:34 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E6=98=BE=E7=A4=BA=E6=96=B9?= =?UTF-8?q?=E5=BC=8F=EF=BC=8C=E5=BE=AE=E8=B0=83=E6=98=BE=E7=A4=BA=E5=8F=82?= =?UTF-8?q?=E6=95=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tools/infer/predict_system.py | 5 ++--- tools/infer/utility.py | 32 +++++++++++++++++++------------- 2 files changed, 21 insertions(+), 16 deletions(-) diff --git a/tools/infer/predict_system.py b/tools/infer/predict_system.py index ba8aaa7e..1991f355 100755 --- a/tools/infer/predict_system.py +++ b/tools/infer/predict_system.py @@ -134,9 +134,8 @@ 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_box_txt(image, boxes, txts) + draw_img = draw_ocr( + image, boxes, txts, scores, draw_txt=True, drop_score=0.5) 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 9fbf3513..ba5a3f38 100755 --- a/tools/infer/utility.py +++ b/tools/infer/utility.py @@ -157,20 +157,23 @@ 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 +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(12345) - draw = ImageDraw.Draw(img_show) + # 每次使用相同的随机种子 ，可以保证两次颜色一致 + 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.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) + 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: @@ -179,13 +182,16 @@ def draw_ocr_box_txt(image, boxes, txts, drop_score=0.5): 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) + 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.text(box[0], txt, fill=(255, 0, 0), font=font) - img_show = Image.blend(img_show_bak, img_show, 0.75) + 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) -- GitLab