Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

增加一个可视化函数 #210

Merged
merged 4 commits into from
Jun 17, 2020
Merged

增加一个可视化函数 #210

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
3f402ca
增加一个画图函数，可以更方便更直观的看检测和识别结果，对横排、竖排文本都可以显示
ZhangXinNan Jun 15, 2020
22e8cf8
修改显示方式，微调显示参数
ZhangXinNan Jun 15, 2020
af6b5fe
增加可视化效果图
ZhangXinNan Jun 16, 2020
5d63991
不增加args.out_dir
ZhangXinNan Jun 16, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions tools/infer/predict_system.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,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):
Expand Down Expand Up @@ -133,8 +134,9 @@ def sorted_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 = draw_ocr(
ZhangXinNan marked this conversation as resolved.
Show resolved Hide resolved
# 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)
Expand Down
32 changes: 32 additions & 0 deletions tools/infer/utility.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
ZhangXinNan marked this conversation as resolved.
Show resolved Hide resolved
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,
Expand Down