forgetest5
/
PaddleOCR
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

增加一个画图函数，可以更方便更直观的看检测和识别结果，对横排、竖排文本都可以显示

This commit is contained in:
zhangxin 2020-06-15 18:40:01 +08:00
parent b47b552124
commit 3f402caea7
2 changed files with 36 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tools/infer/predict_system.py
View File

@ -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)

32
tools/infer/utility.py
View File

@ -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,