Update main_test_SRMNet.py

main_test_SRMNet.py

@@ -23,7 +23,24 @@ def clean_folder(folder):
                 shutil.rmtree(file_path)
         except Exception as e:
             print('Failed to delete %s. Reason: %s' % (file_path, e))
-            
+
+def save_img(filepath, img):
+    cv2.imwrite(filepath, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
+
+
+def load_checkpoint(model, weights):
+    checkpoint = torch.load(weights, map_location=torch.device('cpu'))
+    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')
@@ -74,24 +91,8 @@ def main():
 
         f = os.path.splitext(os.path.split(file_)[-1])[0]
         save_img((os.path.join(out_dir, f + '.png')), restored)
-        clean_folder(inp_dir)
-
-def save_img(filepath, img):
-    cv2.imwrite(filepath, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
-
-
-def load_checkpoint(model, weights):
-    checkpoint = torch.load(weights, map_location=torch.device('cpu'))
-    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)
-
+    clean_folder(inp_dir)
+    
 
 if __name__ == '__main__':
     main()