From 01fd2d0f81d2f0079ea9df2f5b5190d1421b5caa Mon Sep 17 00:00:00 2001 From: littletomatodonkey Date: Mon, 7 Feb 2022 07:56:08 +0000 Subject: [PATCH 1/3] add system pred save --- tools/infer/predict_system.py | 24 +++++++++++++++++++----- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/tools/infer/predict_system.py b/tools/infer/predict_system.py index 16789b81cd..64df865a16 100755 --- a/tools/infer/predict_system.py +++ b/tools/infer/predict_system.py @@ -24,6 +24,7 @@ os.environ["FLAGS_allocator_strategy"] = 'auto_growth' import cv2 import copy import numpy as np +import json import time import logging 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, rec_res) filter_boxes, filter_rec_res = [], [] - for box, rec_result in zip(dt_boxes, rec_res): - text, score = rec_result + for box, rec_reuslt in zip(dt_boxes, rec_res): + text, score = rec_reuslt if score >= self.drop_score: filter_boxes.append(box) - filter_rec_res.append(rec_result) + filter_rec_res.append(rec_reuslt) return filter_boxes, filter_rec_res @@ -128,6 +129,9 @@ def main(args): is_visualize = True font_path = args.vis_font_path 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 if args.warmup: @@ -157,6 +161,14 @@ def main(args): for text, score in rec_res: 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: image = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) boxes = dt_boxes @@ -170,8 +182,6 @@ def main(args): scores, drop_score=drop_score, font_path=font_path) - draw_img_save_dir = args.draw_img_save_dir - os.makedirs(draw_img_save_dir, exist_ok=True) if flag: image_file = image_file[:-3] + "png" cv2.imwrite( @@ -185,6 +195,10 @@ def main(args): text_sys.text_detector.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__": args = utility.parse_args() From b5f861300744dca6cd84fe75f2f5902527b44041 Mon Sep 17 00:00:00 2001 From: littletomatodonkey Date: Mon, 7 Feb 2022 08:00:14 +0000 Subject: [PATCH 2/3] fix typo --- tools/infer/predict_system.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tools/infer/predict_system.py b/tools/infer/predict_system.py index 64df865a16..0265e2a70a 100755 --- a/tools/infer/predict_system.py +++ b/tools/infer/predict_system.py @@ -93,11 +93,11 @@ class TextSystem(object): self.draw_crop_rec_res(self.args.crop_res_save_dir, img_crop_list, rec_res) filter_boxes, filter_rec_res = [], [] - for box, rec_reuslt in zip(dt_boxes, rec_res): - text, score = rec_reuslt + for box, rec_result in zip(dt_boxes, rec_res): + text, score = rec_result if score >= self.drop_score: filter_boxes.append(box) - filter_rec_res.append(rec_reuslt) + filter_rec_res.append(rec_result) return filter_boxes, filter_rec_res @@ -197,7 +197,6 @@ def main(args): with open(os.path.join(draw_img_save_dir, "system_results.txt"), 'w') as f: f.writelines(save_results) - f.close() if __name__ == "__main__": From 417f76b06081e15ddc60bc7076a0e4d04051421b Mon Sep 17 00:00:00 2001 From: littletomatodonkey Date: Mon, 7 Feb 2022 08:23:40 +0000 Subject: [PATCH 3/3] fix duplicate thres --- tools/infer/predict_system.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/infer/predict_system.py b/tools/infer/predict_system.py index 0265e2a70a..b4e316d6a5 100755 --- a/tools/infer/predict_system.py +++ b/tools/infer/predict_system.py @@ -164,7 +164,7 @@ def main(args): 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] + } for idx in range(len(dt_boxes))] save_pred = os.path.basename(image_file) + "\t" + json.dumps( res, ensure_ascii=False) + "\n" save_results.append(save_pred)