Update main_test_CMFNet.py

main_test_CMFNet.py

@@ -14,6 +14,34 @@ from natsort import natsorted
 from model.CMFNet import CMFNet
 
 
+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 clean_folder(folder):
+    for filename in os.listdir(folder):
+        file_path = os.path.join(folder, filename)
+        try:
+            if os.path.isfile(file_path) or os.path.islink(file_path):
+                os.unlink(file_path)
+            elif os.path.isdir(file_path):
+                shutil.rmtree(file_path)
+        except Exception as e:
+            print('Failed to delete %s. Reason: %s' % (file_path, e))
+            
+
 def main():
     parser = argparse.ArgumentParser(description='Demo Image Deraindrop')
     parser.add_argument('--input_dir', default='test/', type=str, help='Input images')
@@ -66,33 +94,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)
-        
-def clean_folder(folder):
-    for filename in os.listdir(folder):
-        file_path = os.path.join(folder, filename)
-        try:
-            if os.path.isfile(file_path) or os.path.islink(file_path):
-                os.unlink(file_path)
-            elif os.path.isdir(file_path):
-                shutil.rmtree(file_path)
-        except Exception as e:
-            print('Failed to delete %s. Reason: %s' % (file_path, e))
 
 if __name__ == '__main__':
     main()