Merge branch 'master' of asdf.us:pix2pixhd

7 files changed, 309 insertions, 13 deletions

diff --git a/augment.py b/augment.py
new file mode 100644
index 0000000..7f448d8
--- /dev/null
+++ b/augment.py
@@ -0,0 +1,141 @@
+### 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
+
+if opt.render_dir == '':
+  opt.render_dir = os.path.join(opt.results_dir, opt.name, opt.which_epoch)
+
+print('tag:', tag)
+print('render_dir:', opt.render_dir)
+util.mkdir(opt.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)
+
+# add augment name
+
+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(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..93fbb6d 100644
--- a/options/dataset_options.py
+++ b/options/dataset_options.py
@@ -293,6 +293,28 @@ 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'
+        )
+
+        self.parser.add_argument(
+          '--augment-name',
+          type=str,
+          help='name to tag epoch'
+        )
+
         ### GRAYSCALE
 
         self.parser.add_argument(
diff --git a/train_epoch.py b/train_epoch.py
new file mode 100644
index 0000000..223d150
--- /dev/null
+++ b/train_epoch.py
@@ -0,0 +1,125 @@
+### 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 time
+from collections import OrderedDict
+from options.train_options import TrainOptions
+from data.data_loader import CreateDataLoader
+from models.models import create_model
+import util.util as util
+from util.visualizer import Visualizer
+import os
+import numpy as np
+import torch
+from torch.autograd import Variable
+
+opt = TrainOptions().parse()
+iter_path = os.path.join(opt.checkpoints_dir, opt.name, 'iter.txt')
+if opt.continue_train:
+    try:
+        start_epoch, epoch_iter = np.loadtxt(iter_path, delimiter=',', dtype=int)
+    except:
+        start_epoch, epoch_iter = 1, 0
+    print('Resuming from epoch %d at iteration %d' % (start_epoch, epoch_iter))        
+else:    
+    start_epoch, epoch_iter = 1, 0
+
+if opt.debug:
+    opt.display_freq = 1
+    opt.print_freq = 1
+    opt.niter = 1
+    opt.niter_decay = 0
+    opt.max_dataset_size = 10
+
+data_loader = CreateDataLoader(opt)
+dataset = data_loader.load_data()
+dataset_size = len(data_loader)
+print('#training images = %d' % dataset_size)
+
+model = create_model(opt)
+visualizer = Visualizer(opt)
+
+total_steps = (start_epoch-1) * dataset_size + epoch_iter
+
+display_delta = total_steps % opt.display_freq
+print_delta = total_steps % opt.print_freq
+save_delta = total_steps % opt.save_latest_freq
+
+print("{} {}".format(start_epoch, start_epoch + opt.niter + opt.niter_decay + 1))
+
+for epoch in range(start_epoch, start_epoch + opt.niter + opt.niter_decay + 1):
+    epoch_start_time = time.time()
+    if epoch != start_epoch:
+        epoch_iter = epoch_iter % dataset_size
+    for i, data in enumerate(dataset, start=epoch_iter):
+        iter_start_time = time.time()
+        total_steps += opt.batchSize
+        epoch_iter += opt.batchSize
+
+        # whether to collect output images
+        save_fake = total_steps % opt.display_freq == display_delta
+
+        ############## Forward Pass ######################
+        losses, generated = model(Variable(data['label']), Variable(data['inst']), 
+            Variable(data['image']), Variable(data['feat']), infer=save_fake)
+
+        # sum per device losses
+        losses = [ torch.mean(x) if not isinstance(x, int) else x for x in losses ]
+        loss_dict = dict(zip(model.module.loss_names, losses))
+
+        # calculate final loss scalar
+        loss_D = (loss_dict['D_fake'] + loss_dict['D_real']) * 0.5
+        loss_G = loss_dict['G_GAN'] + loss_dict.get('G_GAN_Feat',0) + loss_dict.get('G_VGG',0)
+
+        ############### Backward Pass ####################
+        # update generator weights
+        model.module.optimizer_G.zero_grad()
+        loss_G.backward()
+        model.module.optimizer_G.step()
+
+        # update discriminator weights
+        model.module.optimizer_D.zero_grad()
+        loss_D.backward()
+        model.module.optimizer_D.step()
+
+        #call(["nvidia-smi", "--format=csv", "--query-gpu=memory.used,memory.free"]) 
+
+        ############## Display results and errors ##########
+        ### print out errors
+        if total_steps % opt.print_freq == print_delta:
+            errors = {k: v.data[0] if not isinstance(v, int) else v for k, v in loss_dict.items()}
+            t = (time.time() - iter_start_time) / opt.batchSize
+            visualizer.print_current_errors(epoch, epoch_iter, errors, t)
+            visualizer.plot_current_errors(errors, total_steps)
+
+        ### display output images
+        if save_fake:
+            visuals = OrderedDict([('input_label', util.tensor2label(data['label'][0], opt.label_nc)),
+                                   ('synthesized_image', util.tensor2im(generated.data[0])),
+                                   ('real_image', util.tensor2im(data['image'][0]))])
+            visualizer.display_current_results(visuals, epoch, total_steps)
+
+        ### save latest model
+        if total_steps % opt.save_latest_freq == save_delta:
+            print('saving the latest model (epoch %d, total_steps %d)' % (epoch, total_steps))
+            model.module.save('latest')            
+            np.savetxt(iter_path, (epoch, epoch_iter), delimiter=',', fmt='%d')
+       
+    # end of epoch 
+    iter_end_time = time.time()
+    print('End of epoch %d / %d \t Time Taken: %d sec' %
+          (epoch, opt.niter + opt.niter_decay, time.time() - epoch_start_time))
+
+    ### save model for this epoch
+    if epoch % opt.save_epoch_freq == 0:
+        print('saving the model at the end of epoch %d, iters %d' % (epoch, total_steps))        
+        model.module.save('latest')
+        model.module.save(epoch)
+        np.savetxt(iter_path, (epoch+1, 0), delimiter=',', fmt='%d')
+
+    ### instead of only training the local enhancer, train the entire network after certain iterations
+    if (opt.niter_fix_global != 0) and (epoch == opt.niter_fix_global):
+        model.module.update_fixed_params()
+
+    ### linearly decay learning rate after certain iterations
+    if epoch > opt.niter:
+        model.module.update_learning_rate()