augment script

1 files changed, 7 insertions, 5 deletions

diff --git a/augment.py b/augment.py
index 5edbc78..7f448d8 100644
--- a/augment.py
+++ b/augment.py
@@ -67,11 +67,12 @@ if data_opt.tag == '':
 else:
   tag = data_opt.tag
 
-opt.render_dir = render_dir = opt.results_dir + opt.name + "/" + tag + "/"
+if opt.render_dir == '':
+  opt.render_dir = os.path.join(opt.results_dir, opt.name, opt.which_epoch)
 
 print('tag:', tag)
-print('render_dir:', render_dir)
-util.mkdir(render_dir)
+print('render_dir:', opt.render_dir)
+util.mkdir(opt.render_dir)
 
 data_loader = CreateDataLoader(opt)
 dataset = data_loader.load_data()
@@ -97,6 +98,8 @@ if _len <= 0:
 
 transform = get_transform(opt)
 
+# add augment name
+
 for m in range(data_opt.augment_take):
   i = randint(0, _len)
   index = i
@@ -108,7 +111,7 @@ for m in range(data_opt.augment_take):
       A = Image.open(A_path)
       A_tensor = transform(A.convert('RGB'))
     else:
-      A_path = os.path.join(self.opt.render_dir, "recur_{:05d}_{:05d}.png".format(m, index))
+      A_path = os.path.join(opt.render_dir, "recur_{:05d}_{:05d}.png".format(m, index))
       A = Image.open(A_path)
       A_tensor = transform(A.convert('RGB'))
     B_path = sequence[index+1]
@@ -136,4 +139,3 @@ for m in range(data_opt.augment_take):
 
     os.symlink(next_path, os.path.join("./datasets/", data_opt.sequence, "train_A", "recur_{:05d}_{:05d}.png".format(m, index+1)))
     os.symlink(sequence[i+1], os.path.join("./datasets/", data_opt.sequence, "train_B", "recur_{:05d}_{:05d}.png".format(m, index+1)))
-