Update main_test_SRMNet.py

main_test_SRMNet.py

@@ -30,9 +30,8 @@ def main():
     args = parser.parse_args()
 
     # Prepare testing data
-    files = natsorted(glob(os.path.join(inp_dir, '*.JPG')) + glob(os.path.join(inp_dir, '*.PNG')))
-    if len(files) == 0:
-        raise Exception("\nNo images in {} \nPlease enter the following tasks: \n\n{}".format(inp_dir, '\n'.join(tasks)))
+    files = natsorted(glob(os.path.join(args.input_dir, '*.JPG')) + glob(os.path.join(args.input_dir, '*.PNG')))
+
     os.makedirs(args.result_dir, exist_ok=True)
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     # Build model
@@ -63,7 +62,7 @@ def main():
         restored = img_as_ubyte(restored[0])
 
         f = os.path.splitext(os.path.split(file_)[-1])[0]
-        save_img((os.path.join(result_dir, '1.png')), restored)
+        save_img((os.path.join(args.result_dir, '1.png')), restored)
 
     print(f"Files saved at {out_dir}")
     print('finish !')