Update main_test_SRMNet.py

main_test_SRMNet.py

@@ -6,7 +6,7 @@ import os
 from skimage import img_as_ubyte
 from tqdm import tqdm
 from natsort import natsorted
-from glob import glob
+import glob
 import argparse
 from model_arch.SRMNet_SWFF import SRMNet_SWFF
 from model_arch.SRMNet import SRMNet
@@ -30,9 +30,9 @@ def main():
     args = parser.parse_args()
 
     # Prepare testing data
-    files = natsorted(glob(os.path.join(args.input_dir, '*.JPG')) + glob(os.path.join(args.input_dir, '*.PNG')))
+    files = natsorted(glob.glob(os.path.join(args.input_dir, '*')))
     if len(files) == 0:
-        raise Exception(f"No files found at {args.inp_dir}")
+        raise Exception(f"No files found at {args.input_dir}")
     os.makedirs(args.result_dir, exist_ok=True)
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     # Build model