add visualized code for det results

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

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