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

add visualized code for det results

This commit is contained in:
LDOUBLEV 2020-05-15 17:06:43 +08:00
parent 5c043e5b83
commit 1e9683b426
2 changed files with 9 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
tools/infer/predict_det.py
View File

@ -153,6 +153,8 @@ class TextDetector(object):
return dt_boxes, elapse return dt_boxes, elapse
from tools.infer.utility import draw_text_det_res
if __name__ == "__main__": if __name__ == "__main__":
args = utility.parse_args() args = utility.parse_args()
image_file_list = get_image_file_list(args.image_dir) image_file_list = get_image_file_list(args.image_dir)
@ -169,7 +171,9 @@ 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)
""" img_draw = draw_text_det_res(dt_boxes, image_file, return_img=True)
add visualized code 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)) print("Avg Time:", total_time / (count - 1))

5
tools/infer/utility.py
View File

@ -103,13 +103,12 @@ def create_predictor(args, mode):
return predictor, input_tensor, output_tensors 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) src_im = cv2.imread(img_path)
for box in dt_boxes: for box in dt_boxes:
box = np.array(box).astype(np.int32).reshape(-1, 2) box = np.array(box).astype(np.int32).reshape(-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)
img_name_pure = img_path.split("/")[-1] return src_im
cv2.imwrite("./output/%s" % img_name_pure, src_im)
def resize_img(img, input_size=600): def resize_img(img, input_size=600):