Update main_test_SRMNet.py

main_test_SRMNet.py

@@ -34,10 +34,10 @@ def load_checkpoint(model, weights):
 
 def main():
     parser = argparse.ArgumentParser(description='Demo Image Denoising')
-    parser.add_argument('--input_dir', default='./test/', type=str, help='Input images')
-    parser.add_argument('--result_dir', default='./result/', type=str, help='Directory for results')
+    parser.add_argument('--input_dir', default='test/', type=str, help='Input images')
+    parser.add_argument('--result_dir', default='result/', type=str, help='Directory for results')
     parser.add_argument('--weights',
-                        default='./checkpoints/SRMNet_real_denoise/models/model_bestPSNR.pth', type=str,
+                        default='experiments/pretrained_models/AWGN_denoising_SRMNet.pth', type=str,
                         help='Path to weights')
 
     args = parser.parse_args()
@@ -82,6 +82,11 @@ def main():
         save_img((os.path.join(out_dir, f + '.png')), restored)
 
 
+def setup(args):
+    save_dir = 'result/'
+    folder = 'test/'
+
+    return folder, save_dir
 
 if __name__ == '__main__':
     main()