Update main_test_CMFNet.py

main_test_CMFNet.py

@@ -13,6 +13,7 @@ import torch.nn.functional as F
 from natsort import natsorted
 from model.CMFNet import CMFNet
 
+
 def main():
     parser = argparse.ArgumentParser(description='Demo Image Deraindrop')
     parser.add_argument('--input_dir', default='test/', type=str, help='Input images')
@@ -64,7 +65,7 @@ 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):
@@ -82,7 +83,17 @@ def load_checkpoint(model, weights):
             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()