augment script

6 files changed, 176 insertions, 13 deletions

diff --git a/augment.py b/augment.py
new file mode 100644
index 0000000..5edbc78
--- /dev/null
+++ b/augment.py
@@ -0,0 +1,139 @@
+### Copyright (C) 2017 NVIDIA Corporation. All rights reserved. 
+### Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
+import os
+import sys
+sys.path.append(os.path.join(os.path.dirname(os.path.realpath(__file__)), '../live-cortex/rpc/'))
+from collections import OrderedDict
+from options.test_options import TestOptions
+from options.dataset_options import DatasetOptions
+from data.data_loader import CreateDataLoader
+from models.models import create_model
+import util.util as util
+from util.visualizer import Visualizer
+from util import html
+import torch
+from run_engine import run_trt_engine, run_onnx
+from datetime import datetime
+from PIL import Image, ImageOps
+from shutil import copyfile, rmtree
+from random import randint
+
+from img_ops import read_sequence
+
+import torch.utils.data as data
+from PIL import Image
+import torchvision.transforms as transforms
+
+def get_transform(opt, method=Image.BICUBIC, normalize=True):
+  transform_list = []
+
+  base = float(2 ** opt.n_downsample_global)
+  if opt.netG == 'local':
+    base *= (2 ** opt.n_local_enhancers)
+  transform_list.append(transforms.Lambda(lambda img: __make_power_2(img, base, method)))
+
+  transform_list += [transforms.ToTensor()]
+
+  if normalize:
+    transform_list += [transforms.Normalize((0.5, 0.5, 0.5),
+                                              (0.5, 0.5, 0.5))]
+  return transforms.Compose(transform_list)
+
+def normalize():    
+  return transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+
+def __make_power_2(img, base, method=Image.BICUBIC):
+  ow, oh = img.size        
+  h = int(round(oh / base) * base)
+  w = int(round(ow / base) * base)
+  if (h == oh) and (w == ow):
+    return img
+  return img.resize((w, h), method)
+
+
+opt = TestOptions().parse(save=False)
+data_opt = DatasetOptions().parse(opt.unknown)
+opt.nThreads = 1   # test code only supports nThreads = 1
+opt.batchSize = 1  # test code only supports batchSize = 1
+opt.serial_batches = True  # no shuffle
+opt.no_flip = True  # no flip
+if data_opt.tag == '':
+  d = datetime.now()
+  tag = data_opt.tag = "{}_{}".format(
+    opt.name,
+    # opt.experiment,
+    d.strftime('%Y%m%d%H%M')
+  )
+else:
+  tag = data_opt.tag
+
+opt.render_dir = render_dir = opt.results_dir + opt.name + "/" + tag + "/"
+
+print('tag:', tag)
+print('render_dir:', render_dir)
+util.mkdir(render_dir)
+
+data_loader = CreateDataLoader(opt)
+dataset = data_loader.load_data()
+
+if not opt.engine and not opt.onnx:
+  model = create_model(opt)
+  if opt.data_type == 16:
+    model.half()
+  elif opt.data_type == 8:
+    model.type(torch.uint8)
+  if opt.verbose:
+    print(model)
+
+sequence = read_sequence(data_opt.sequence_name, '')
+print("Got sequence {}, {} images".format(data_opt.sequence, len(sequence)))
+_len = len(sequence) - data_opt.augment_take
+
+if _len <= 0:
+  print("Got empty sequence...")
+  data_opt.processing = False
+  rpc_client.send_status('processing', False)
+  sys.exit(1)
+
+transform = get_transform(opt)
+
+for m in range(data_opt.augment_take):
+  i = randint(0, _len)
+  index = i
+
+  for n in range(data_opt.augment_make):
+    index = i + n
+    if n == 0:
+      A_path = sequence[i]
+      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 = Image.open(A_path)
+      A_tensor = transform(A.convert('RGB'))
+    B_path = sequence[index+1]
+    inst_tensor = 0
+
+    input_dict = {'label': A_tensor, 'inst': inst_tensor}
+
+    if opt.data_type == 16:
+      data['label'] = data['label'].half()
+      data['inst']  = data['inst'].half()
+    elif opt.data_type == 8:
+      data['label'] = data['label'].uint8()
+      data['inst']  = data['inst'].uint8()
+    minibatch = 1
+    generated = model.inference(data['label'], data['inst'])
+    
+    tmp_path = os.path.join(opt.render_dir, "recur_{:05d}_{:05d}_tmp.png".format(m, index+1))
+    next_path = os.path.join(opt.render_dir, "recur_{:05d}_{:05d}.png".format(m, index+1))
+    print('process image... %i' % index)
+
+    im = util.tensor2im(generated.data[0])
+    image_pil = Image.fromarray(im, mode='RGB')
+    image_pil.save(tmp_path)
+    os.rename(tmp_path, next_path)
+
+    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)))
+
diff --git a/build_dataset.pl b/build_dataset.pl
index f007038..f88b818 100755
--- a/build_dataset.pl
+++ b/build_dataset.pl
@@ -39,12 +39,19 @@ my $dir;
 my $src;
 my $dst;
 for ($i = 0; $i < $count; $i++) {
-  if (($i % $test_split) != ($test_split-1)) {
-    $dir = $thumbs_dir . "train_";
-  } else {
+  if (($i % $test_split) = ($test_split-1)) {
     $dir = $thumbs_dir . "test_";
+
+    $src = $images_dir . $images[$i];
+    $dst = $dir . sprintf("A/frame_%05d.png", $i);
+    system("ln", "-s", $src, $dst);
+
+    $src = $images_dir . $images[$i+1];
+    $dst = $dir . sprintf("B/frame_%05d.png", $i);
+    system("ln", "-s", $src, $dst);
   }
 
+  $dir = $thumbs_dir . "train_";
   $src = $images_dir . $images[$i];
   $dst = $dir . sprintf("A/frame_%05d.png", $i);
   system("ln", "-s", $src, $dst);
diff --git a/data/recursive_dataset.py b/data/recursive_dataset.py
index 40b7ebf..dda9c95 100644
--- a/data/recursive_dataset.py
+++ b/data/recursive_dataset.py
@@ -13,7 +13,7 @@ class RecursiveDataset(BaseDataset):
 
         ### input A (label maps)
         self.dataset_size = 1000000
-      
+    
     def __getitem__(self, index):
         ### input A (label maps)
         A_path = os.path.join(self.opt.render_dir, "frame_{:05d}.png".format(index))
diff --git a/data/sequence_dataset.py b/data/sequence_dataset.py
index 3eaa12b..c3c7d44 100644
--- a/data/sequence_dataset.py
+++ b/data/sequence_dataset.py
@@ -8,18 +8,18 @@ from PIL import Image
 class SequenceDataset(BaseDataset):
     def initialize(self, opt):
         self.opt = opt
-        self.root = opt.dataroot    
+        self.root = opt.dataroot
 
         ### input A (label maps)
         self.dir_A = opt.dataroot
         self.A_paths = sorted(make_dataset(self.dir_A))
 
-        self.dataset_size = len(self.A_paths) 
-      
-    def __getitem__(self, index):        
+        self.dataset_size = len(self.A_paths)
+
+    def __getitem__(self, index):
         ### input A (label maps)
-        A_path = self.A_paths[index]              
-        A = Image.open(A_path)        
+        A_path = self.A_paths[index]
+        A = Image.open(A_path)
         params = get_params(self.opt, A.size)
         if self.opt.label_nc == 0:
             transform_A = get_transform(self.opt, params)
diff --git a/live.py b/live.py
index 859669b..394f764 100644
--- a/live.py
+++ b/live.py
@@ -24,8 +24,8 @@ import gevent
 from time import sleep
 from shutil import copyfile, rmtree
 
-from img_ops import process_image, mix_next_image
-from listener import Listener, read_sequence
+from img_ops import read_sequence, process_image, mix_next_image
+from listener import Listener
 
 module_name = 'pix2pixhd'
 
@@ -134,7 +134,8 @@ def process_live_input(opt, data_opt, rpc_client):
     if data_opt.pause:
       data_opt.pause = False
       break
-    gevent.sleep(data_opt.frame_delay)
+    if data_opt.frame_delay > 0:
+      gevent.sleep(data_opt.frame_delay)
 
   data_opt.processing = False
   rpc_client.send_status('processing', False)
diff --git a/options/dataset_options.py b/options/dataset_options.py
index d9c3bf7..29e07c6 100644
--- a/options/dataset_options.py
+++ b/options/dataset_options.py
@@ -293,6 +293,22 @@ class DatasetOptions(BaseOptions):
           help='amount of processed image (clahe, poster, etc) to use in feeder step'
         )
 
+        ### DATA AUGMENTATION
+
+        self.parser.add_argument(
+          '--augment-take',
+          default=100,
+          type=int,
+          help='number of random images to take'
+        )
+
+        self.parser.add_argument(
+          '--augment-make',
+          default=15,
+          type=int,
+          help='number of recursive images to generate'
+        )
+
         ### GRAYSCALE
 
         self.parser.add_argument(