提交 c15d3bb0 编写于 作者: L LDOUBLEV

add visualize code to predict_eval

上级 073b8517
......@@ -11,7 +11,7 @@ PaddleOCR旨在打造一套丰富、领先、且实用的OCR工具库,助力
## **超轻量级中文OCR体验**
![](./doc/imgs_draw/11.jpg)
![](doc/imgs_results/11.jpg)
上图是超轻量级中文OCR模型效果展示,更多效果图请见文末[效果展示](#效果展示)
......@@ -97,14 +97,14 @@ PaddleOCR文本识别算法的训练和使用请参考文档教程中[文本识
<a name="效果展示"></a>
## 效果展示
![](./doc/imgs_draw/1.jpg)
![](./doc/imgs_draw/7.jpg)
![](./doc/imgs_draw/12.jpg)
![](./doc/imgs_draw/4.jpg)
![](./doc/imgs_draw/6.jpg)
![](./doc/imgs_draw/9.jpg)
![](./doc/imgs_draw/16.png)
![](./doc/imgs_draw/22.jpg)
![](doc/imgs_results/1.jpg)
![](doc/imgs_results/7.jpg)
![](doc/imgs_results/12.jpg)
![](doc/imgs_results/4.jpg)
![](doc/imgs_results/6.jpg)
![](doc/imgs_results/9.jpg)
![](doc/imgs_results/16.png)
![](doc/imgs_results/22.jpg)
## 参考文献
......
......@@ -22,6 +22,10 @@ import copy
import numpy as np
import math
import time
from ppocr.utils.utility import get_image_file_list
from PIL import Image
from tools.infer.utility import draw_ocr
import os
class TextSystem(object):
......@@ -99,8 +103,9 @@ def sorted_boxes(dt_boxes):
if __name__ == "__main__":
args = utility.parse_args()
image_file_list = utility.get_image_file_list(args.image_dir)
image_file_list = get_image_file_list(args.image_dir)
text_sys = TextSystem(args)
is_visualize = True
for image_file in image_file_list:
img = cv2.imread(image_file)
if img is None:
......@@ -114,8 +119,22 @@ if __name__ == "__main__":
dt_boxes_final = []
for dno in range(dt_num):
text, score = rec_res[dno]
if score >= 0:
if score >= 0.5:
text_str = "%s, %.3f" % (text, score)
print(text_str)
dt_boxes_final.append(dt_boxes[dno])
utility.draw_text_det_res(dt_boxes_final, image_file)
if is_visualize:
image = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
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_save = "./doc/imgs_results/"
if not os.path.exists(draw_img_save):
os.makedirs(draw_img_save)
cv2.imwrite(
os.path.join(draw_img_save, os.path.basename(image_file)),
draw_img)
......@@ -112,36 +112,70 @@ def draw_text_det_res(dt_boxes, img_path):
cv2.imwrite("./output/%s" % img_name_pure, src_im)
def draw_ocr(image, boxes, txts, scores, draw_txt):
def resize_img(img, input_size=600):
"""
"""
img = np.array(img)
im_shape = img.shape
im_size_min = np.min(im_shape[0:2])
im_size_max = np.max(im_shape[0:2])
im_scale = float(input_size) / float(im_size_max)
im = cv2.resize(img, None, None, fx=im_scale, fy=im_scale)
return im
def draw_ocr(image, boxes, txts, scores, draw_txt=True, drop_score=0.5):
from PIL import Image, ImageDraw, ImageFont
w, h = image.size
img = image.copy()
draw = ImageDraw.Draw(img)
for (box, txt) in zip(boxes, txts):
for (box, score) in zip(boxes, scores):
if score < drop_score:
continue
draw.line([(box[0][0], box[0][1]), (box[1][0], box[1][1])], fill='red')
draw.line([(box[1][0], box[1][1]), (box[2][0], box[2][1])], fill='red')
draw.line([(box[2][0], box[2][1]), (box[3][0], box[3][1])], fill='red')
draw.line([(box[3][0], box[3][1]), (box[0][0], box[0][1])], fill='red')
draw.line(
[(box[0][0] - 1, box[0][1] + 1), (box[1][0] - 1, box[1][1] + 1)],
fill='red')
draw.line(
[(box[1][0] - 1, box[1][1] + 1), (box[2][0] - 1, box[2][1] + 1)],
fill='red')
draw.line(
[(box[2][0] - 1, box[2][1] + 1), (box[3][0] - 1, box[3][1] + 1)],
fill='red')
draw.line(
[(box[3][0] - 1, box[3][1] + 1), (box[0][0] - 1, box[0][1] + 1)],
fill='red')
if draw_txt:
txt_color = (0, 0, 0)
blank_img = np.ones(shape=[h, 800], dtype=np.int8) * 255
img = np.array(resize_img(img))
_h = img.shape[0]
blank_img = np.ones(shape=[_h, 600], dtype=np.int8) * 255
blank_img = Image.fromarray(blank_img).convert("RGB")
draw_txt = ImageDraw.Draw(blank_img)
font_size = 30
gap = 40 if h // len(txts) >= font_size else h // len(txts)
for i, txt in enumerate(txts):
font_size = 20
gap = 20
title = "index text score"
font = ImageFont.truetype(
"./doc/simfang.TTF", font_size, encoding="utf-8")
new_txt = str(i) + ': ' + txt + ' ' + str(scores[i])
draw_txt.text((20, gap * (i + 1)), new_txt, txt_color, font=font)
"./doc/simfang.ttf", font_size, encoding="utf-8")
draw_txt.text((20, 0), title, txt_color, font=font)
count = 0
for idx, txt in enumerate(txts):
if scores[idx] < drop_score:
continue
font = ImageFont.truetype(
"./doc/simfang.ttf", font_size, encoding="utf-8")
new_txt = str(count) + ': ' + txt + ' ' + str(scores[count])
draw_txt.text(
(20, gap * (count + 1)), new_txt, txt_color, font=font)
count += 1
img = np.concatenate([np.array(img), np.array(blank_img)], axis=1)
return img
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册