Update main_test_SUNet.py

main_test_SUNet.py

@@ -18,7 +18,16 @@ from model.SUNet import SUNet_model
 with open('training.yaml', 'r') as config:
     opt = yaml.safe_load(config)
 
-
+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 Restoration')
@@ -90,6 +99,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))