yujunjun/PaddleOCR
1
0
Fork 0
mirror of https://github.com/PaddlePaddle/PaddleOCR.git synced 2025-11-03 19:29:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

add visualized code for det results

Browse Source
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
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))

5
tools/infer/utility.py
View File

@ -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):