Update main_test_SRMNet.py

main_test_SRMNet.py

@@ -13,6 +13,17 @@ import torch.nn.functional as F
 from natsort import natsorted
 from model.SRMNet import SRMNet
 
+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 Denoising')
     parser.add_argument('--input_dir', default='test/', type=str, help='Input images')
@@ -63,7 +74,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):
     cv2.imwrite(filepath, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))