1 files changed, 7 insertions, 1 deletions
diff --git a/uprez.py b/uprez.py
index 0be4643..249025d 100644
--- a/uprez.py
+++ b/uprez.py
@@ -1,18 +1,23 @@
 import os
 import time
 import subprocess
+import argparse
 
 start_time = time.time()
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--folder', required=True)
 parser.add_argument('--out_fn', required=True)
+parser.add_argument('--dataset', required=True)
+opt = parser.parse_args()
 
 os.makedirs('./uprez', exist_ok=True)
 
+dataset = opt.dataset
+
 print(">>> {}".format(dataset))
 
-in_dir = os.path.abspath('./recursive/{}'.format(dataset))
+in_dir = os.path.abspath(opt.folder)
 vsr_duf_inputs_dir = os.path.abspath('../vsr-duf/inputs')
 vsr_duf_results_dir = os.path.abspath('../vsr-duf/results/{}'.format(dataset))
 out_dir = os.path.abspath('./uprez')
@@ -45,3 +50,4 @@ subprocess.call([
 finish_time = time.time()
 
 print("took %02f s.".format(finish_time - start_time))
+