diff options
Diffstat (limited to 'augment.py')
| -rw-r--r-- | augment.py | 12 |
1 files changed, 7 insertions, 5 deletions
@@ -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))) - |