提交 55065c04 编写于 作者: T tink2123

Merge branch 'develop' of https://github.com/PaddlePaddle/PaddleOCR into develop

...@@ -11,7 +11,7 @@ Global: ...@@ -11,7 +11,7 @@ Global:
test_batch_size_per_card: 16 test_batch_size_per_card: 16
image_shape: [3, 640, 640] image_shape: [3, 640, 640]
reader_yml: ./configs/det/det_db_icdar15_reader.yml reader_yml: ./configs/det/det_db_icdar15_reader.yml
pretrain_weights: ./pretrain_models/MobileNetV3_pretrained/MobileNetV3_large_x0_5_pretrained/ pretrain_weights: ./pretrain_models/MobileNetV3_large_x0_5_pretrained/
checkpoints: checkpoints:
save_res_path: ./output/det_db/predicts_db.txt save_res_path: ./output/det_db/predicts_db.txt
save_inference_dir: save_inference_dir:
......
...@@ -89,13 +89,13 @@ class EvalTestReader(object): ...@@ -89,13 +89,13 @@ class EvalTestReader(object):
def batch_iter_reader(): def batch_iter_reader():
batch_outs = [] batch_outs = []
for img_path, img_name in img_list: for img_path in img_list:
img = cv2.imread(img_path) img = cv2.imread(img_path)
if img is None: if img is None:
logger.info("load image error:" + img_path) logger.info("load image error:" + img_path)
continue continue
outs = process_function(img) outs = process_function(img)
outs.append(img_name) outs.append(img_path)
batch_outs.append(outs) batch_outs.append(outs)
if len(batch_outs) == batch_size: if len(batch_outs) == batch_size:
yield batch_outs yield batch_outs
......
...@@ -20,11 +20,14 @@ from ppocr.data.det.east_process import EASTProcessTest ...@@ -20,11 +20,14 @@ from ppocr.data.det.east_process import EASTProcessTest
from ppocr.data.det.db_process import DBProcessTest from ppocr.data.det.db_process import DBProcessTest
from ppocr.postprocess.db_postprocess import DBPostProcess from ppocr.postprocess.db_postprocess import DBPostProcess
from ppocr.postprocess.east_postprocess import EASTPostPocess from ppocr.postprocess.east_postprocess import EASTPostPocess
from ppocr.utils.utility import get_image_file_list
from tools.infer.utility import draw_ocr
import copy import copy
import numpy as np import numpy as np
import math import math
import time import time
import sys import sys
import os
class TextDetector(object): class TextDetector(object):
...@@ -152,7 +155,7 @@ class TextDetector(object): ...@@ -152,7 +155,7 @@ class TextDetector(object):
if __name__ == "__main__": if __name__ == "__main__":
args = utility.parse_args() 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_detector = TextDetector(args) text_detector = TextDetector(args)
count = 0 count = 0
total_time = 0 total_time = 0
...@@ -166,5 +169,14 @@ if __name__ == "__main__": ...@@ -166,5 +169,14 @@ if __name__ == "__main__":
total_time += elapse total_time += elapse
count += 1 count += 1
print("Predict time of %s:" % image_file, elapse) print("Predict time of %s:" % image_file, elapse)
utility.draw_text_det_res(dt_boxes, image_file) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
draw_img = draw_ocr(img, dt_boxes, None, None, False)
draw_img_save = "./inference_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[:, :, ::-1])
print("The visualized image saved in {}".format(
os.path.join(draw_img_save, os.path.basename(image_file))))
print("Avg Time:", total_time / (count - 1)) print("Avg Time:", total_time / (count - 1))
...@@ -127,10 +127,10 @@ def resize_img(img, input_size=600): ...@@ -127,10 +127,10 @@ def resize_img(img, input_size=600):
def draw_ocr(image, boxes, txts, scores, draw_txt=True, drop_score=0.5): def draw_ocr(image, boxes, txts, scores, draw_txt=True, drop_score=0.5):
from PIL import Image, ImageDraw, ImageFont from PIL import Image, ImageDraw, ImageFont
w, h = image.size
img = image.copy() img = image.copy()
draw = ImageDraw.Draw(img) draw = ImageDraw.Draw(img)
if scores is None:
scores = [1] * len(boxes)
for (box, score) in zip(boxes, scores): for (box, score) in zip(boxes, scores):
if score < drop_score: if score < drop_score:
continue continue
......
...@@ -40,7 +40,7 @@ set_paddle_flags( ...@@ -40,7 +40,7 @@ set_paddle_flags(
) )
from paddle import fluid from paddle import fluid
from ppocr.utils.utility import create_module from ppocr.utils.utility import create_module, get_image_file_list
import program import program
from ppocr.utils.save_load import init_model from ppocr.utils.save_load import init_model
from ppocr.data.reader_main import reader_main from ppocr.data.reader_main import reader_main
...@@ -50,20 +50,18 @@ from ppocr.utils.utility import initial_logger ...@@ -50,20 +50,18 @@ from ppocr.utils.utility import initial_logger
logger = initial_logger() logger = initial_logger()
def draw_det_res(dt_boxes, config, img_name, ino): def draw_det_res(dt_boxes, config, img, img_name):
if len(dt_boxes) > 0: if len(dt_boxes) > 0:
img_set_path = config['TestReader']['img_set_dir']
img_path = img_set_path + img_name
import cv2 import cv2
src_im = cv2.imread(img_path) src_im = img
for box in dt_boxes: for box in dt_boxes:
box = box.astype(np.int32).reshape((-1, 1, 2)) box = box.astype(np.int32).reshape((-1, 1, 2))
cv2.polylines(src_im, [box], True, color=(255, 255, 0), thickness=2) cv2.polylines(src_im, [box], True, color=(255, 255, 0), thickness=2)
save_det_path = os.path.basename(config['Global'][ save_det_path = os.path.dirname(config['Global'][
'save_res_path']) + "/det_results/" 'save_res_path']) + "/det_results/"
if not os.path.exists(save_det_path): if not os.path.exists(save_det_path):
os.makedirs(save_det_path) os.makedirs(save_det_path)
save_path = os.path.join(save_det_path, "det_{}.jpg".format(img_name)) save_path = os.path.join(save_det_path, os.path.basename(img_name))
cv2.imwrite(save_path, src_im) cv2.imwrite(save_path, src_im)
logger.info("The detected Image saved in {}".format(save_path)) logger.info("The detected Image saved in {}".format(save_path))
...@@ -103,8 +101,12 @@ def main(): ...@@ -103,8 +101,12 @@ def main():
raise Exception("{} not exists!".format(checkpoints)) raise Exception("{} not exists!".format(checkpoints))
save_res_path = config['Global']['save_res_path'] save_res_path = config['Global']['save_res_path']
if not os.path.exists(os.path.dirname(save_res_path)):
os.makedirs(os.path.dirname(save_res_path))
with open(save_res_path, "wb") as fout: with open(save_res_path, "wb") as fout:
test_reader = reader_main(config=config, mode='test') test_reader = reader_main(config=config, mode='test')
# image_file_list = get_image_file_list(args.image_dir)
tackling_num = 0 tackling_num = 0
for data in test_reader(): for data in test_reader():
img_num = len(data) img_num = len(data)
...@@ -128,7 +130,13 @@ def main(): ...@@ -128,7 +130,13 @@ def main():
postprocess_params.update(global_params) postprocess_params.update(global_params)
postprocess = create_module(postprocess_params['function'])\ postprocess = create_module(postprocess_params['function'])\
(params=postprocess_params) (params=postprocess_params)
dt_boxes_list = postprocess({"maps": outs[0]}, ratio_list) if config['Global']['algorithm'] == 'EAST':
dic = {'f_score': outs[0], 'f_geo': outs[1]}
elif config['Global']['algorithm'] == 'DB':
dic = {'maps': outs[0]}
else:
raise Exception("only support algorithm: ['EAST', 'BD']")
dt_boxes_list = postprocess(dic, ratio_list)
for ino in range(img_num): for ino in range(img_num):
dt_boxes = dt_boxes_list[ino] dt_boxes = dt_boxes_list[ino]
img_name = img_name_list[ino] img_name = img_name_list[ino]
...@@ -139,7 +147,8 @@ def main(): ...@@ -139,7 +147,8 @@ def main():
dt_boxes_json.append(tmp_json) dt_boxes_json.append(tmp_json)
otstr = img_name + "\t" + json.dumps(dt_boxes_json) + "\n" otstr = img_name + "\t" + json.dumps(dt_boxes_json) + "\n"
fout.write(otstr.encode()) fout.write(otstr.encode())
draw_det_res(dt_boxes, config, img_name, ino) src_img = cv2.imread(img_name)
draw_det_res(dt_boxes, config, src_img, img_name)
logger.info("success!") logger.info("success!")
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册