From 6ae1ca4c2cc4537b86cc660762fb87a828f10002 Mon Sep 17 00:00:00 2001 From: Shengwei An Date: Thu, 4 Jan 2024 09:12:08 -0500 Subject: [PATCH] Skip non-image files in the input_dir_path The script will crash when the image_path is a non-image file. --- mist_v3.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/mist_v3.py b/mist_v3.py index f597628..64d4f91 100644 --- a/mist_v3.py +++ b/mist_v3.py @@ -291,6 +291,10 @@ def infer(img: PIL.Image.Image, config, tar_img: PIL.Image.Image = None, mask: P for img_id in os.listdir(image_dir_path): image_path = os.path.join(image_dir_path, img_id) + if os.path.splitext(image_path)[1].lower().strip('.') not in ['png', 'jpg', 'jpeg', ]: + print(f'skip non-image file {image_path}') + continue + if resize: img, target_size = closing_resize(image_path, input_size, block_num) bls_h = target_size[0]//block_num @@ -363,4 +367,4 @@ def infer(img: PIL.Image.Image, config, tar_img: PIL.Image.Image = None, mask: P save_parameter = '_' + str(epsilon) + '_' + str(steps) + '_' + str(input_size) + '_' + str(block_num) + '_' + str(mode) + '_' + str(args.rate) + '_' + str(int(mask)) + '_' + str(int(resize)) output_name += save_parameter + '.png' print("Output image saved in path {}".format(output_name)) - output.save(output_name) \ No newline at end of file + output.save(output_name)