From 1e9683b42640efe87278a9c6bdef1dbc394d6c99 Mon Sep 17 00:00:00 2001 From: LDOUBLEV Date: Fri, 15 May 2020 17:06:43 +0800 Subject: [PATCH] add visualized code for det results --- tools/infer/predict_det.py | 10 +++++++--- tools/infer/utility.py | 5 ++--- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/tools/infer/predict_det.py b/tools/infer/predict_det.py index cbed5a4a..2a446cdf 100755 --- a/tools/infer/predict_det.py +++ b/tools/infer/predict_det.py @@ -153,6 +153,8 @@ class TextDetector(object): return dt_boxes, elapse +from tools.infer.utility import draw_text_det_res + if __name__ == "__main__": args = utility.parse_args() image_file_list = get_image_file_list(args.image_dir) @@ -169,7 +171,9 @@ if __name__ == "__main__": total_time += elapse count += 1 print("Predict time of %s:" % image_file, elapse) - """ - add visualized code - """ + img_draw = draw_text_det_res(dt_boxes, image_file, return_img=True) + save_path = os.path.join("./inference_det/", + os.path.basename(image_file)) + print("The visualized image saved in {}".format(save_path)) + print("Avg Time:", total_time / (count - 1)) diff --git a/tools/infer/utility.py b/tools/infer/utility.py index 95c23321..ff458787 100755 --- a/tools/infer/utility.py +++ b/tools/infer/utility.py @@ -103,13 +103,12 @@ def create_predictor(args, mode): return predictor, input_tensor, output_tensors -def draw_text_det_res(dt_boxes, img_path): +def draw_text_det_res(dt_boxes, img_path, return_img=True): src_im = cv2.imread(img_path) for box in dt_boxes: box = np.array(box).astype(np.int32).reshape(-1, 2) cv2.polylines(src_im, [box], True, color=(255, 255, 0), thickness=2) - img_name_pure = img_path.split("/")[-1] - cv2.imwrite("./output/%s" % img_name_pure, src_im) + return src_im def resize_img(img, input_size=600): -- GitLab