Update main_test_SRMNet.py

main_test_SRMNet.py

@@ -66,6 +66,7 @@ def main():
         f = os.path.splitext(os.path.split(file_)[-1])[0]
         save_img((os.path.join(args.result_dir, f + '.png')), restored)
         print('{}'.format(os.path.join(args.result_dir, f + '.png')))
+    clean_folder(args.input_dir)
     print('finish !')
 
 
@@ -99,6 +100,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__':