Update main_test_SRMNet.py

main_test_SRMNet.py

@@ -14,24 +14,6 @@ from natsort import natsorted
 from model.SRMNet import SRMNet
 from utils import util_calculate_psnr_ssim as util
 
-
-def save_img(filepath, img):
-    cv2.imwrite(filepath, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
-
-
-def load_checkpoint(model, weights):
-    checkpoint = torch.load(weights)
-    try:
-        model.load_state_dict(checkpoint["state_dict"])
-    except:
-        state_dict = checkpoint["state_dict"]
-        new_state_dict = OrderedDict()
-        for k, v in state_dict.items():
-            name = k[7:]  # remove `module.`
-            new_state_dict[name] = v
-        model.load_state_dict(new_state_dict)
-
-
 def main():
     parser = argparse.ArgumentParser(description='Demo Image Denoising')
     parser.add_argument('--input_dir', default='test/', type=str, help='Input images')
@@ -47,7 +29,7 @@ def main():
 
     os.makedirs(out_dir, exist_ok=True)
 
-    files = natsorted(glob(os.path.join(inp_dir, '*')))
+    files = natsorted(glob.glob(os.path.join(inp_dir, '*')))
 
     if len(files) == 0:
         raise Exception(f"No files found at {inp_dir}")
@@ -82,6 +64,22 @@ def main():
         save_img((os.path.join(out_dir, f + '.png')), restored)
 
 
+def save_img(filepath, img):
+    cv2.imwrite(filepath, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
+
+
+def load_checkpoint(model, weights):
+    checkpoint = torch.load(weights)
+    try:
+        model.load_state_dict(checkpoint["state_dict"])
+    except:
+        state_dict = checkpoint["state_dict"]
+        new_state_dict = OrderedDict()
+        for k, v in state_dict.items():
+            name = k[7:]  # remove `module.`
+            new_state_dict[name] = v
+        model.load_state_dict(new_state_dict)
+
 def setup(args):
     save_dir = 'result/'
     folder = 'test/'