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

Merge pull request #36 from LDOUBLEV/fixocr 

valid det inference
This commit is contained in:
dyning 2020-05-15 14:26:52 +08:00 committed by GitHub
parent 3b40c32a2f d539508e76
commit 81d8d19038
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 37 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
configs/det/det_db_mv3.yml
View File

@ -11,7 +11,7 @@ Global:
test_batch_size_per_card: 16
image_shape: [3, 640, 640]
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:
save_res_path: ./output/det_db/predicts_db.txt
save_inference_dir:

4
ppocr/data/det/dataset_traversal.py
View File

@ -89,13 +89,13 @@ class EvalTestReader(object):
def batch_iter_reader():
batch_outs = []
for img_path, img_name in img_list:
for img_path in img_list:
img = cv2.imread(img_path)
if img is None:
logger.info("load image error:" + img_path)
continue
outs = process_function(img)
outs.append(img_name)
outs.append(img_path)
batch_outs.append(outs)
if len(batch_outs) == batch_size:
yield batch_outs

16
tools/infer/predict_det.py
View File

@ -20,11 +20,14 @@ from ppocr.data.det.east_process import EASTProcessTest
from ppocr.data.det.db_process import DBProcessTest
from ppocr.postprocess.db_postprocess import DBPostProcess
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 numpy as np
import math
import time
import sys
import os
class TextDetector(object):
@ -152,7 +155,7 @@ class TextDetector(object):
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_detector = TextDetector(args)
count = 0
total_time = 0
@ -166,5 +169,14 @@ if __name__ == "__main__":
total_time += elapse
count += 1
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))

4
tools/infer/utility.py
View File

@ -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):
from PIL import Image, ImageDraw, ImageFont
w, h = image.size
img = image.copy()
draw = ImageDraw.Draw(img)
if scores is None:
scores = [1] * len(boxes)
for (box, score) in zip(boxes, scores):
if score < drop_score:
continue

27
tools/infer_det.py
View File

@ -40,7 +40,7 @@ set_paddle_flags(
)
from paddle import fluid
from ppocr.utils.utility import create_module
from ppocr.utils.utility import create_module, get_image_file_list
import program
from ppocr.utils.save_load import init_model
from ppocr.data.reader_main import reader_main
@ -50,20 +50,18 @@ from ppocr.utils.utility import 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:
img_set_path = config['TestReader']['img_set_dir']
img_path = img_set_path + img_name
import cv2
src_im = cv2.imread(img_path)
src_im = img
for box in dt_boxes:
box = box.astype(np.int32).reshape((-1, 1, 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/"
if not os.path.exists(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)
logger.info("The detected Image saved in {}".format(save_path))
@ -103,8 +101,12 @@ def main():
raise Exception("{} not exists!".format(checkpoints))
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:
test_reader = reader_main(config=config, mode='test')
# image_file_list = get_image_file_list(args.image_dir)
tackling_num = 0
for data in test_reader():
img_num = len(data)
@ -128,7 +130,13 @@ def main():
postprocess_params.update(global_params)
postprocess = create_module(postprocess_params['function'])\
(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):
dt_boxes = dt_boxes_list[ino]
img_name = img_name_list[ino]
@ -139,7 +147,8 @@ def main():
dt_boxes_json.append(tmp_json)
otstr = img_name + "\t" + json.dumps(dt_boxes_json) + "\n"
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!")