diff --git a/tools/infer/predict_system.py b/tools/infer/predict_system.py
index de7ee9d342063161f2e329c99d2428051c0ecf8c..ea622fac5e5870437907ee8c5b8068d77ecd3c0c 100755
--- a/tools/infer/predict_system.py
+++ b/tools/infer/predict_system.py
@@ -176,6 +176,8 @@ def main(args):
             draw_img_save = "./inference_results/"
             if not os.path.exists(draw_img_save):
                 os.makedirs(draw_img_save)
+            if flag:
+               image_file = image_file[:-3] + "png" 
             cv2.imwrite(
                 os.path.join(draw_img_save, os.path.basename(image_file)),
                 draw_img[:, :, ::-1])