Update main_test_SRMNet.py

main_test_SRMNet.py

@@ -10,7 +10,7 @@ import requests
 from PIL import Image
 import torchvision.transforms.functional as TF
 import torch.nn.functional as F
-
+from natsort import natsorted
 from model.SRMNet import SRMNet
 from utils import util_calculate_psnr_ssim as util
 
@@ -47,7 +47,8 @@ def main():
 
     os.makedirs(out_dir, exist_ok=True)
 
-    files = sorted(glob.glob(os.path.join(inp_dir, '*.PNG')))
+    files = natsorted(glob(os.path.join(inp_dir, '*.JPG'))
+                  + glob(os.path.join(inp_dir, '*.PNG')))
 
     if len(files) == 0:
         raise Exception(f"No files found at {inp_dir}")