mirrors/PaddleOCR
1
0
Fork 0
mirror of https://github.com/PaddlePaddle/PaddleOCR.git synced 2025-06-03 21:53:39 +08:00
Code Issues Packages Projects Releases Wiki Activity

add system pred save

Browse Source
This commit is contained in:
littletomatodonkey 2022-02-07 07:56:08 +00:00
parent aaae49584f
commit 01fd2d0f81
1 changed files with 19 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
tools/infer/predict_system.py
View File

@ -24,6 +24,7 @@ os.environ["FLAGS_allocator_strategy"] = 'auto_growth'
import cv2 import cv2
import copy import copy
import numpy as np import numpy as np
import json
import time import time
import logging import logging
from PIL import Image from PIL import Image
@ -92,11 +93,11 @@ class TextSystem(object):
self.draw_crop_rec_res(self.args.crop_res_save_dir, img_crop_list, self.draw_crop_rec_res(self.args.crop_res_save_dir, img_crop_list,
rec_res) rec_res)
filter_boxes, filter_rec_res = [], [] filter_boxes, filter_rec_res = [], []
for box, rec_result in zip(dt_boxes, rec_res): for box, rec_reuslt in zip(dt_boxes, rec_res):
text, score = rec_result text, score = rec_reuslt
if score >= self.drop_score: if score >= self.drop_score:
filter_boxes.append(box) filter_boxes.append(box)
filter_rec_res.append(rec_result) filter_rec_res.append(rec_reuslt)
return filter_boxes, filter_rec_res return filter_boxes, filter_rec_res
@ -128,6 +129,9 @@ def main(args):
is_visualize = True is_visualize = True
font_path = args.vis_font_path font_path = args.vis_font_path
drop_score = args.drop_score drop_score = args.drop_score
draw_img_save_dir = args.draw_img_save_dir
os.makedirs(draw_img_save_dir, exist_ok=True)
save_results = []
# warm up 10 times # warm up 10 times
if args.warmup: if args.warmup:
@ -157,6 +161,14 @@ def main(args):
for text, score in rec_res: for text, score in rec_res:
logger.debug("{}, {:.3f}".format(text, score)) logger.debug("{}, {:.3f}".format(text, score))
res = [{
"transcription": rec_res[idx][0],
"points": np.array(dt_boxes[idx]).astype(np.int32).tolist(),
} for idx in range(len(dt_boxes)) if rec_res[idx][1] >= drop_score]
save_pred = os.path.basename(image_file) + "\t" + json.dumps(
res, ensure_ascii=False) + "\n"
save_results.append(save_pred)
if is_visualize: if is_visualize:
image = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) image = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
boxes = dt_boxes boxes = dt_boxes
@ -170,8 +182,6 @@ def main(args):
scores, scores,
drop_score=drop_score, drop_score=drop_score,
font_path=font_path) font_path=font_path)
draw_img_save_dir = args.draw_img_save_dir
os.makedirs(draw_img_save_dir, exist_ok=True)
if flag: if flag:
image_file = image_file[:-3] + "png" image_file = image_file[:-3] + "png"
cv2.imwrite( cv2.imwrite(
@ -185,6 +195,10 @@ def main(args):
text_sys.text_detector.autolog.report() text_sys.text_detector.autolog.report()
text_sys.text_recognizer.autolog.report() text_sys.text_recognizer.autolog.report()
with open(os.path.join(draw_img_save_dir, "system_results.txt"), 'w') as f:
f.writelines(save_results)
f.close()
if __name__ == "__main__": if __name__ == "__main__":
args = utility.parse_args() args = utility.parse_args()