1 files changed, 19 insertions, 15 deletions
diff --git a/app/relay/modules/pix2pixhd.js b/app/relay/modules/pix2pixhd.js
index 73b49ca..a90fc15 100644
--- a/app/relay/modules/pix2pixhd.js
+++ b/app/relay/modules/pix2pixhd.js
@@ -110,31 +110,35 @@ const generate = {
 }
 const augment = {
   type: 'pytorch',
-  script: 'test.py',
+  script: 'augment.py',
   params: (task) => {
     let epoch = 0
     const dataset = task.dataset.toLowerCase()
     const datasets_path = path.join(cwd, 'datasets', dataset)
     const checkpoints_path = path.join(cwd, 'checkpoints', dataset)
-    const iter_txt = path.join(checkpoints_path, 'iter.txt')
-    console.log(dataset, iter_txt)
-    if (fs.existsSync(iter_txt)) {
-      const iter = fs.readFileSync(iter_txt).toString().split('\n');
-      console.log(iter)
-      epoch = iter[0] || 0
-      console.log(task.module, dataset, '=>', epoch, task.epochs)
-    } else {
-      console.log(task.module, dataset, '=>', 'starting new training')
-    }
+    // supply render_dir
     return [
       '--dataroot', datasets_path,
+      '--results_dir', './recursive',
       '--module_name', task.module,
       '--name', dataset,
       '--model', 'pix2pixHD',
       '--label_nc', 0, '--no_instance',
-      '--niter', task.epochs,
-      '--niter_decay', 0,
-      '--save_epoch_freq', 1,
+      '--augment-take', task.opt.augment_take,
+      '--augment-make', task.opt.augment_make,
+      '--augment-name', task.opt.augment_name,
+      '--which_epoch', task.opt.epoch,
+    ]
+  },
+}
+const clear_recursive = {
+  type: 'pytorch',
+  script: 'clear_recursive.py',
+  params: (task) => {
+    const dataset = task.dataset.toLowerCase()
+    return [
+      '--name', dataset,
+      '--epoch', epoch,
     ]
   },
 }
@@ -153,7 +157,6 @@ const live = {
       '--name', task.checkpoint,
       '--module_name', 'pix2pixHD',
       '--sequence-name', task.dataset,
-      '--recursive', '--recursive-frac', 0.1,
       '--sequence', '--sequence-frac', 0.3,
       '--process-frac', 0.5,
       '--label_nc', '0', '--no_instance',
@@ -200,6 +203,7 @@ export default {
     train,
     generate,
     augment,
+    clear_recursive,
     live,
     render,
   }