Added -seeds option to provide seeds from disk

Author: dribnet

enhance.py

@@ -38,6 +38,7 @@
 add_arg('--scales',             default=2, type=int,                help='How many times to perform 2x upsampling.')
 add_arg('--model',              default='small', type=str,          help='Name of the neural network to load/save.')
 add_arg('--train',              default=False, type=str,            help='File pattern to load for training.')
+add_arg('--seeds',              default=False, type=str,            help='File pattern to load for training seeds.')
 add_arg('--epochs',             default=10, type=int,               help='Total number of iterations in training.')
 add_arg('--epoch-size',         default=72, type=int,               help='Number of batches trained in an epoch.')
 add_arg('--save-every',         default=10, type=int,               help='Save generator after every training epoch.')
@@ -118,6 +119,28 @@ def extend(lst): return itertools.chain(lst, itertools.repeat(lst[-1]))
 
 print('{}  - Using the device `{}` for neural computation.{}\n'.format(ansi.CYAN, theano.config.device, ansi.ENDC))
 
+def confirm_pairs(list1, list2):
+    new_list1 = []
+    new_list2 = []
+    cur1 = 0
+    cur2 = 0
+    len1 = len(list1)
+    len2 = len(list2)
+    while(cur1 < len1 and cur2 < len2):
+        base1 = os.path.basename(list1[cur1])
+        base2 = os.path.basename(list2[cur2])
+        if base1 == base2:
+            new_list1.append(list1[cur1])
+            new_list2.append(list2[cur2])
+            cur1 = cur1 + 1
+            cur2 = cur2 + 1
+        elif base1 < base2:
+            # continue to look on list1, don't iterate list2
+            cur1 = cur1 + 1
+        else:
+            cur2 = cur2 + 1
+    print("List sizes went from {}, {} to {}, {}".format(len1, len2, len(new_list1), len(new_list2)))
+    return new_list1, new_list2
 
 #======================================================================================================================
 # Image Processing
@@ -133,7 +156,13 @@ def __init__(self):
 
         self.orig_buffer = np.zeros((args.buffer_size, 3, self.orig_shape, self.orig_shape), dtype=np.float32)
         self.seed_buffer = np.zeros((args.buffer_size, 3, self.seed_shape, self.seed_shape), dtype=np.float32)
-        self.files = glob.glob(args.train)
+        self.files = sorted(glob.glob(args.train))
+        if args.seeds:
+            self.seeds = sorted(glob.glob(args.seeds))
+            self.files, self.seeds = confirm_pairs(self.files, self.seeds)
+        else:
+            self.seeds = False
+
         if len(self.files) == 0:
             error("There were no files found to train from searching for `{}`".format(args.train),
                   "  - Try putting all your images in one folder and using `--train=data/*.jpg`")
@@ -146,22 +175,56 @@ def __init__(self):
 
     def run(self):
         while True:
-            random.shuffle(self.files)
+            indices = list(range(0, len(self.files)))
+            random.shuffle(indices)
 
-            for f in self.files:
+            for i in indices:
+                f = self.files[i]
                 filename = os.path.join(self.cwd, f)
                 try:
                     img = scipy.ndimage.imread(filename, mode='RGB')
                 except Exception as e:
                     warn('Could not load `{}` as image.'.format(filename),
                          '  - Try fixing or removing the file before next run.')
-                    files.remove(f)
+                    del self.files[i]
+                    if self.seeds:
+                        del self.seeds[i]
                     continue
-                
+
+                # determine seed
+                if self.seeds:
+                    f = self.seeds[i]
+                    filename = os.path.join(self.cwd, f)
+                    try:
+                        seed_img = scipy.ndimage.imread(filename, mode='RGB')
+                    except Exception as e:
+                        warn('Could not load `{}` as seed image.'.format(filename),
+                             '  - Try fixing or removing the file before next run.')
+                        del self.files[i]
+                        del self.seeds[i]
+                        continue
+                else:
+                    # synthetic seed
+                    seed_img = scipy.misc.imresize(img,
+                        float(self.seed_shape) / float(self.orig_shape),
+                        interp='bilinear')
+
                 for _ in range(args.buffer_similar):
-                    copy = img[:,::-1] if random.choice([True, False]) else img
-                    h = random.randint(0, copy.shape[0] - self.orig_shape)
-                    w = random.randint(0, copy.shape[1] - self.orig_shape)
+                    if random.choice([True, False]):
+                        copy = img[:,::-1]
+                        copy_seed = seed_img[:,::-1]
+                    else:
+                        copy = img
+                        copy_seed = seed_img
+
+                    # compute seed displacement
+                    h = random.randint(0, copy_seed.shape[0] - self.seed_shape)
+                    w = random.randint(0, copy_seed.shape[1] - self.seed_shape)
+                    copy_seed = copy_seed[h:h+self.seed_shape, w:w+self.seed_shape]
+
+                    # and matching image displacement
+                    h = int(h * self.orig_shape / self.seed_shape);
+                    w = int(w * self.orig_shape / self.seed_shape);
                     copy = copy[h:h+self.orig_shape, w:w+self.orig_shape]
 
                     while len(self.available) == 0:
@@ -170,8 +233,7 @@ def run(self):
 
                     i = self.available.pop()
                     self.orig_buffer[i] = np.transpose(copy / 255.0 - 0.5, (2, 0, 1))
-                    seed = scipy.misc.imresize(copy, size=(self.seed_shape, self.seed_shape), interp='bilinear')
-                    self.seed_buffer[i] = np.transpose(seed / 255.0 - 0.5, (2, 0, 1))
+                    self.seed_buffer[i] = np.transpose(copy_seed / 255.0 - 0.5, (2, 0, 1))
                     self.ready.add(i)
 
                     if len(self.ready) >= args.batch_size: