add visualized code for det results

tools/infer/predict_det.py

@@ -153,6 +153,8 @@ def __call__(self, img):
  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 @@ def __call__(self, img):
  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))

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