Merge pull request #47 from alexjc/v0.2

Release 0.2: Tiled Rendering, 1:1 Noise/Blur Fixing

docker-cpu.df

@@ -26,9 +26,8 @@ RUN /opt/conda/bin/python3.5 -m pip install -q -r "requirements.txt"
 COPY enhance.py .
 
 # Get a pre-trained neural networks, non-commercial & attribution.
-RUN wget -q "https://github.com/alexjc/neural-enhance/releases/download/v0.1/ne4x-small-0.1.pkl.bz2"
+RUN wget -q "https://github.com/alexjc/neural-enhance/releases/download/v0.2/ne1x-small-0.2.pkl.bz2"
-RUN wget -q "https://github.com/alexjc/neural-enhance/releases/download/v0.1/ne4x-medium-0.1.pkl.bz2"
+RUN wget -q "https://github.com/alexjc/neural-enhance/releases/download/v0.2/ne2x-small-0.2.pkl.bz2"
-RUN wget -q "https://github.com/alexjc/neural-enhance/releases/download/v0.1/ne4x-large-0.1.pkl.bz2"
 
 # Set an entrypoint to the main enhance.py script
 ENTRYPOINT ["/opt/conda/bin/python3.5", "enhance.py", "--device=cpu"]

docker-gpu.df

@@ -24,9 +24,8 @@ RUN /opt/conda/bin/python3.5 -m pip install -q -r "requirements.txt"
 COPY enhance.py .
 
 # Get a pre-trained neural networks, non-commercial & attribution.
-RUN wget -q "https://github.com/alexjc/neural-enhance/releases/download/v0.1/ne4x-small-0.1.pkl.bz2"
+RUN wget -q "https://github.com/alexjc/neural-enhance/releases/download/v0.2/ne1x-small-0.2.pkl.bz2"
-RUN wget -q "https://github.com/alexjc/neural-enhance/releases/download/v0.1/ne4x-medium-0.1.pkl.bz2"
+RUN wget -q "https://github.com/alexjc/neural-enhance/releases/download/v0.2/ne2x-small-0.2.pkl.bz2"
-RUN wget -q "https://github.com/alexjc/neural-enhance/releases/download/v0.1/ne4x-large-0.1.pkl.bz2"
 
 # Set an entrypoint to the main enhance.py script
 ENTRYPOINT ["/opt/conda/bin/python3.5", "enhance.py", "--device=gpu"]

enhance.py

@@ -14,8 +14,9 @@
 # without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 #
 
-__version__ = '0.1'
+__version__ = '0.2'
 
+import io
 import os
 import sys
 import bz2
@@ -35,19 +36,26 @@ parser = argparse.ArgumentParser(description='Generate a new image by applying s
                                  formatter_class=argparse.ArgumentDefaultsHelpFormatter)
 add_arg = parser.add_argument
 add_arg('files',                nargs='*', default=[])
-add_arg('--scales',             default=2, type=int,                help='How many times to perform 2x upsampling.')
+add_arg('--zoom',               default=1, type=int,                help='Resolution increase factor for inference.')
+add_arg('--rendering-tile',     default=128, type=int,              help='Size of tiles used for rendering images.')
+add_arg('--rendering-overlap',  default=32, type=int,               help='Number of pixels padding around each tile.')
 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('--train-blur',         default=None, type=int,             help='Sigma value for gaussian blur preprocess.')
+add_arg('--train-noise',        default=None, type=float,           help='Radius for preprocessing gaussian blur.')
+add_arg('--train-jpeg',         default=None, type=int,             help='JPEG compression level in preprocessing.')
 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.')
 add_arg('--batch-shape',        default=192, type=int,              help='Resolution of images in training batch.')
-add_arg('--batch-size',         default=15, type=int,               help='Number of images per training batch.')
+add_arg('--batch-size',         default=10, type=int,               help='Number of images per training batch.')
 add_arg('--buffer-size',        default=1500, type=int,             help='Total image fragments kept in cache.')
 add_arg('--buffer-similar',     default=5, type=int,                help='Fragments cached for each image loaded.')
-add_arg('--learning-rate',      default=1E-4, type=float,           help='Parameter for the ADAM optimizer.')
+add_arg('--learning-rate',      default=5E-4, type=float,           help='Parameter for the ADAM optimizer.')
-add_arg('--learning-period',    default=50, type=int,               help='How often to decay the learning rate.')
+add_arg('--learning-period',    default=100, type=int,              help='How often to decay the learning rate.')
 add_arg('--learning-decay',     default=0.5, type=float,            help='How much to decay the learning rate.')
+add_arg('--generator-upscale',  default=2, type=int,                help='Steps of 2x up-sampling as post-process.')
+add_arg('--generator-downscale',default=0, type=int,                help='Steps of 2x down-sampling as preprocess.')
 add_arg('--generator-filters',  default=[64], nargs='+', type=int,  help='Number of convolution units in network.')
 add_arg('--generator-blocks',   default=4, type=int,                help='Number of residual blocks per iteration.')
 add_arg('--generator-residual', default=2, type=int,                help='Number of layers in a residual block.')
@@ -55,7 +63,7 @@ add_arg('--perceptual-layer',   default='conv2_2', type=str,        help='Which
 add_arg('--perceptual-weight',  default=1e0, type=float,            help='Weight for VGG-layer perceptual loss.')
 add_arg('--discriminator-size', default=32, type=int,               help='Multiplier for number of filters in D.')
 add_arg('--smoothness-weight',  default=2e5, type=float,            help='Weight of the total-variation loss.')
-add_arg('--adversary-weight',   default=1e2, type=float,            help='Weight of adversarial loss compoment.')
+add_arg('--adversary-weight',   default=5e2, type=float,            help='Weight of adversarial loss compoment.')
 add_arg('--generator-start',    default=0, type=int,                help='Epoch count to start training generator.')
 add_arg('--discriminator-start',default=1, type=int,                help='Epoch count to update the discriminator.')
 add_arg('--adversarial-start',  default=2, type=int,                help='Epoch for generator to use discriminator.')
@@ -100,11 +108,10 @@ os.environ.setdefault('THEANO_FLAGS', 'floatX=float32,device={},force_device=Tru
 
 # Scientific & Imaging Libraries
 import numpy as np
-import scipy.optimize, scipy.ndimage, scipy.misc
+import scipy.ndimage, scipy.misc, PIL.Image
 
 # Numeric Computing (GPU)
-import theano
+import theano, theano.tensor as T
-import theano.tensor as T
 T.nnet.softminus = lambda x: x - T.nnet.softplus(x)
 
 # Support ansi colors in Windows too.
@@ -129,7 +136,7 @@ class DataLoader(threading.Thread):
         self.data_ready = threading.Event()
         self.data_copied = threading.Event()
 
-        self.orig_shape, self.seed_shape = args.batch_shape, int(args.batch_shape / 2**args.scales)
+        self.orig_shape, self.seed_shape = args.batch_shape, int(args.batch_shape / args.zoom)
 
         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)
@@ -147,31 +154,51 @@ class DataLoader(threading.Thread):
     def run(self):
         while True:
             random.shuffle(self.files)
-
             for f in self.files:
+                self.add_to_buffer(f)
+
+    def add_to_buffer(self, f):
         filename = os.path.join(self.cwd, f)
         try:
-                    img = scipy.ndimage.imread(filename, mode='RGB')
+            orig = PIL.Image.open(filename).convert('RGB')
+            # if all(s > args.batch_shape * 2 for s in orig.size): 
+            #     orig = orig.resize((orig.size[0]//2, orig.size[1]//2), resample=PIL.Image.LANCZOS)
+            if any(s < args.batch_shape for s in orig.size):
+                raise ValueError('Image is too small for training with size {}'.format(orig.size))
         except Exception as e:
             warn('Could not load `{}` as image.'.format(filename),
                  '  - Try fixing or removing the file before next run.')
-                    files.remove(f)
+            self.files.remove(f)
-                    continue
+            return
+
+        seed = orig.filter(PIL.ImageFilter.GaussianBlur(radius=args.train_blur)) if args.train_blur else orig
+        seed = seed.resize((orig.size[0]//args.zoom, orig.size[1]//args.zoom), resample=PIL.Image.LANCZOS)
+
+        if args.train_jpeg:
+            buffer = io.BytesIO()
+            seed.save(buffer, format='jpeg', quality=args.train_jpeg+random.randrange(-15,+15))
+            seed = PIL.Image.open(buffer)
+
+        seed = scipy.misc.fromimage(seed, mode='RGB').astype(np.float32)
+        seed += scipy.random.normal(scale=args.train_noise, size=(seed.shape[0], seed.shape[1], 1))\
+                                                            if args.train_noise else 0.0
+
+        orig = scipy.misc.fromimage(orig).astype(np.float32)
 
         for _ in range(args.buffer_similar):
-                    copy = img[:,::-1] if random.choice([True, False]) else img
+            h = random.randint(0, seed.shape[0] - self.seed_shape)
-                    h = random.randint(0, copy.shape[0] - self.orig_shape)
+            w = random.randint(0, seed.shape[1] - self.seed_shape)
-                    w = random.randint(0, copy.shape[1] - self.orig_shape)
+            seed_chunk = seed[h:h+self.seed_shape, w:w+self.seed_shape]
-                    copy = copy[h:h+self.orig_shape, w:w+self.orig_shape]
+            h, w = h * args.zoom, w * args.zoom
+            orig_chunk = orig[h:h+self.orig_shape, w:w+self.orig_shape]
 
             while len(self.available) == 0:
                 self.data_copied.wait()
                 self.data_copied.clear()
 
             i = self.available.pop()
-                    self.orig_buffer[i] = np.transpose(copy / 255.0 - 0.5, (2, 0, 1))
+            self.orig_buffer[i] = np.transpose(orig_chunk.astype(np.float32) / 127.5 - 1.0, (2, 0, 1))
-                    seed = scipy.misc.imresize(copy, size=(self.seed_shape, self.seed_shape), interp='bilinear')
+            self.seed_buffer[i] = np.transpose(seed_chunk.astype(np.float32) / 127.5 - 1.0, (2, 0, 1))
-                    self.seed_buffer[i] = np.transpose(seed / 255.0 - 0.5, (2, 0, 1))
             self.ready.add(i)
 
             if len(self.ready) >= args.batch_size:
@@ -184,9 +211,7 @@ class DataLoader(threading.Thread):
         for i, j in enumerate(random.sample(self.ready, args.batch_size)):
             origs_out[i] = self.orig_buffer[j]
             seeds_out[i] = self.seed_buffer[j]
-
             self.available.add(j)
-
         self.data_copied.set()
 
 
@@ -214,6 +239,34 @@ class SubpixelReshuffleLayer(lasagne.layers.Layer):
         return out
 
 
+class ReflectLayer(lasagne.layers.Layer):
+    """Based on more code by ajbrock: https://gist.github.com/ajbrock/a3858c26282d9731191901b397b3ce9f
+    """
+
+    def __init__(self, incoming, pad, batch_ndim=2, **kwargs):
+        super(ReflectLayer, self).__init__(incoming, **kwargs)
+        self.pad = (pad, pad)
+        self.batch_ndim = batch_ndim
+
+    def get_output_shape_for(self, input_shape):
+        output_shape = list(input_shape)
+        for k, p in enumerate(self.pad):
+            if output_shape[k + self.batch_ndim] is None: continue
+            l, r = p, p
+            output_shape[k + self.batch_ndim] += l + r
+        return tuple(output_shape)
+
+    def get_output_for(self, x, **kwargs):
+        out = T.zeros(self.get_output_shape_for(x.shape))
+        p0, p1 = self.pad
+        out = T.set_subtensor(out[:,:,:p0,p1:-p1],    x[:,:,p0:0:-1,:])
+        out = T.set_subtensor(out[:,:,-p0:,p1:-p1],   x[:,:,-2:-(2+p0):-1,:])
+        out = T.set_subtensor(out[:,:,p0:-p0,p1:-p1], x)
+        out = T.set_subtensor(out[:,:,:,:p1],         out[:,:,:,(2*p1):p1:-1])
+        out = T.set_subtensor(out[:,:,:,-p1:],        out[:,:,:,-(p1+2):-(2*p1+2):-1])
+        return out
+
+
 class Model(object):
 
     def __init__(self):
@@ -230,7 +283,6 @@ class Model(object):
             self.load_perceptual()
             self.setup_discriminator()
         self.load_generator(params)
-
         self.compile()
 
     #------------------------------------------------------------------------------------------------------------------
@@ -241,7 +293,8 @@ class Model(object):
         return list(self.network.values())[-1]
 
     def make_layer(self, name, input, units, filter_size=(3,3), stride=(1,1), pad=(1,1), alpha=0.25):
-        conv = ConvLayer(input, units, filter_size=filter_size, stride=stride, pad=pad, nonlinearity=None)
+        reflected = ReflectLayer(input, pad=pad[0]) if pad[0] > 0 else input
+        conv = ConvLayer(reflected, units, filter_size, stride=stride, pad=(0,0), nonlinearity=None)
         prelu = lasagne.layers.ParametricRectifierLayer(conv, alpha=lasagne.init.Constant(alpha))
         self.network[name+'x'] = conv
         self.network[name+'>'] = prelu
@@ -254,30 +307,35 @@ class Model(object):
 
     def setup_generator(self, input, config):
         for k, v in config.items(): setattr(args, k, v)
+        args.zoom = 2**(args.generator_upscale - args.generator_downscale)
+
         units_iter = extend(args.generator_filters)
         units = next(units_iter)
         self.make_layer('iter.0-A', input, units, filter_size=(5,5), pad=(2,2))
         self.make_layer('iter.0-B', self.last_layer(), units, filter_size=(5,5), pad=(2,2))
         self.network['iter.0'] = self.last_layer()
 
+        for i in range(0, args.generator_downscale):
+            self.make_layer('downscale%i'%i, self.last_layer(), next(units_iter), filter_size=(4,4), stride=(2,2))
+
+        units = next(units_iter)
         for i in range(0, args.generator_blocks):
             self.make_block('iter.%i'%(i+1), self.last_layer(), units)
 
-        for i in range(0, args.scales):
+        for i in range(0, args.generator_upscale):
             u = next(units_iter)
-            self.make_layer('scale%i.3'%i, self.last_layer(), u*4)
+            self.make_layer('upscale%i.3'%i, self.last_layer(), u*4)
-            self.network['scale%i.2'%i] = SubpixelReshuffleLayer(self.last_layer(), u, 2)
+            self.network['upscale%i.2'%i] = SubpixelReshuffleLayer(self.last_layer(), u, 2)
-            self.make_layer('scale%i.1'%i, self.last_layer(), u)
+            self.make_layer('upscale%i.1'%i, self.last_layer(), u)
 
-        self.network['out'] = ConvLayer(self.last_layer(), 3, filter_size=(5,5), stride=(1,1), pad=(2,2),
+        self.network['out'] = ConvLayer(self.last_layer(), 3, filter_size=(3,3), stride=(1,1), pad=(1,1),
                                                               nonlinearity=lasagne.nonlinearities.tanh)
 
     def setup_perceptual(self, input):
         """Use lasagne to create a network of convolution layers using pre-trained VGG19 weights.
         """
-
         offset = np.array([103.939, 116.779, 123.680], dtype=np.float32).reshape((1,3,1,1))
-        self.network['percept'] = lasagne.layers.NonlinearityLayer(input, lambda x: ((x+0.5).clip(0.0, 1.0)*255.0) - offset)
+        self.network['percept'] = lasagne.layers.NonlinearityLayer(input, lambda x: ((x+1.0)*127.5) - offset)
 
         self.network['mse'] = self.network['percept']
         self.network['conv1_1'] = ConvLayer(self.network['percept'], 64, 3, pad=1)
@@ -340,13 +398,14 @@ class Model(object):
     def save_generator(self):
         def cast(p): return p.get_value().astype(np.float16)
         params = {k: [cast(p) for p in l.get_params()] for (k, l) in self.list_generator_layers()}
-        config = {k: getattr(args, k) for k in ['generator_blocks', 'generator_residual', 'generator_filters']}
+        config = {k: getattr(args, k) for k in ['generator_blocks', 'generator_residual', 'generator_filters'] + \
-        filename = 'ne%ix-%s-%s.pkl.bz2' % (2**args.scales, args.model, __version__)
+                                               ['generator_upscale', 'generator_downscale']}
+        filename = 'ne%ix-%s-%s.pkl.bz2' % (args.zoom, args.model, __version__)
         pickle.dump((config, params), bz2.open(filename, 'wb'))
         print('  - Saved model as `{}` after training.'.format(filename))
 
     def load_model(self):
-        filename = 'ne%ix-%s-%s.pkl.bz2' % (2**args.scales, args.model, __version__)
+        filename = 'ne%ix-%s-%s.pkl.bz2' % (args.zoom, args.model, __version__)
         if not os.path.exists(filename):
             if args.train: return {}, {}
             error("Model file with pre-trained convolution layers not found. Download it here...",
@@ -380,7 +439,7 @@ class Model(object):
         return T.mean(T.nnet.softminus(d[args.batch_size:]) - T.nnet.softplus(d[:args.batch_size]))
 
     def compile(self):
-        # Helper function for rendering test images during training, or standalone non-training mode.
+        # Helper function for rendering test images during training, or standalone inference mode.
         input_tensor, seed_tensor = T.tensor4(), T.tensor4()
         input_layers = {self.network['img']: input_tensor, self.network['seed']: seed_tensor}
         output = lasagne.layers.get_output([self.network[k] for k in ['seed','out']], input_layers, deterministic=True)
@@ -406,7 +465,8 @@ class Model(object):
         disc_losses = [self.loss_discriminator(disc_out)]
         disc_params = list(itertools.chain(*[l.get_params() for k, l in self.network.items() if 'disc' in k]))
         print('  - {} tensors learned for discriminator.'.format(len(disc_params)))
-        disc_updates = lasagne.updates.adam(sum(disc_losses, 0.0), disc_params, learning_rate=self.disc_lr)
+        grads = [g.clip(-5.0, +5.0) for g in T.grad(sum(disc_losses, 0.0), disc_params)]
+        disc_updates = lasagne.updates.adam(grads, disc_params, learning_rate=self.disc_lr)
 
         # Combined Theano function for updating both generator and discriminator at the same time.
         updates = collections.OrderedDict(list(gen_updates.items()) + list(disc_updates.items()))
@@ -416,7 +476,7 @@ class Model(object):
 
 class NeuralEnhancer(object):
 
-    def __init__(self):
+    def __init__(self, loader):
         if args.train:
             print('{}Training {} epochs on random image sections with batch size {}.{}'\
                   .format(ansi.BLUE_B, args.epochs, args.batch_size, ansi.BLUE))
@@ -425,33 +485,31 @@ class NeuralEnhancer(object):
             print('{}Enhancing {} image(s) specified on the command-line.{}'\
                   .format(ansi.BLUE_B, len(args.files), ansi.BLUE))
 
-        self.thread = DataLoader() if args.train else None
+        self.thread = DataLoader() if loader else None
         self.model = Model()
 
         print('{}'.format(ansi.ENDC))
 
     def imsave(self, fn, img):
-        img = np.transpose(img + 0.5, (1, 2, 0)).clip(0.0, 1.0)
+        scipy.misc.toimage(np.transpose(img + 1.0, (1, 2, 0)) * 127.5, cmin=0, cmax=255).save(fn)
-        image = scipy.misc.toimage(img * 255.0, cmin=0, cmax=255)
-        image.save(fn)
 
     def show_progress(self, orign, scald, repro):
         os.makedirs('valid', exist_ok=True)
         for i in range(args.batch_size):
-            self.imsave('valid/%03i_origin.png' % i, orign[i])
+            self.imsave('valid/%s_%03i_origin.png' % (args.model, i), orign[i])
-            self.imsave('valid/%03i_pixels.png' % i, scald[i])
+            self.imsave('valid/%s_%03i_pixels.png' % (args.model, i), scald[i])
-            self.imsave('valid/%03i_reprod.png' % i, repro[i])
+            self.imsave('valid/%s_%03i_reprod.png' % (args.model, i), repro[i])
 
     def decay_learning_rate(self):
         l_r, t_cur = args.learning_rate, 0
 
         while True:
-            yield l_r if t_cur > 0 else l_r * 0.1
+            yield l_r
             t_cur += 1
             if t_cur % args.learning_period == 0: l_r *= args.learning_decay
 
     def train(self):
-        seed_size = int(args.batch_shape / 2**args.scales)
+        seed_size = args.batch_shape // args.zoom
         images = np.zeros((args.batch_size, 3, args.batch_shape, args.batch_shape), dtype=np.float32)
         seeds = np.zeros((args.batch_size, 3, seed_size, seed_size), dtype=np.float32)
         learning_rate = self.decay_learning_rate()
@@ -480,11 +538,12 @@ class NeuralEnhancer(object):
                 stats /= args.epoch_size
                 totals, labels = [sum(total)] + list(total), ['total', 'prcpt', 'smthn', 'advrs']
                 gen_info = ['{}{}{}={:4.2e}'.format(ansi.WHITE_B, k, ansi.ENDC, v) for k, v in zip(labels, totals)]
-                print('\rEpoch #{} at {:4.1f}s, lr={:4.2e}    {}'.format(epoch+1, time.time()-start, l_r, ' '*args.epoch_size))
+                print('\rEpoch #{} at {:4.1f}s, lr={:4.2e}{}'.format(epoch+1, time.time()-start, l_r, ' '*(args.epoch_size-30)))
                 print('  - generator {}'.format(' '.join(gen_info)))
 
                 real, fake = stats[:args.batch_size], stats[args.batch_size:]
-                print('  - discriminator', real.mean(), len(np.where(real > 0.5)[0]), fake.mean(), len(np.where(fake < -0.5)[0]))
+                print('  - discriminator', real.mean(), len(np.where(real > 0.5)[0]),
+                                           fake.mean(), len(np.where(fake < -0.5)[0]))
                 if epoch == args.adversarial_start-1:
                     print('  - generator now optimizing against discriminator.')
                     self.model.adversary_weight.set_value(args.adversary_weight)
@@ -497,31 +556,33 @@ class NeuralEnhancer(object):
             pass
 
         print('\n{}Trained {}x super-resolution for {} epochs.{}'\
-                .format(ansi.CYAN_B, 2**args.scales, epoch+1, ansi.CYAN))
+                .format(ansi.CYAN_B, args.zoom, epoch+1, ansi.CYAN))
         self.model.save_generator()
         print(ansi.ENDC)
 
-    def process(self, image):
+    def process(self, original):
-        img = np.transpose(image / 255.0 - 0.5, (2, 0, 1))[np.newaxis].astype(np.float32)
+        s, p, z = args.rendering_tile, args.rendering_overlap, args.zoom
+        image = np.pad(original, ((p*z, p*z), (p*z, p*z), (0, 0)), mode='reflect')
+        output = np.zeros((original.shape[0] * z, original.shape[1] * z, 3), dtype=np.float32)
+        for y, x in itertools.product(range(0, original.shape[0], s), range(0, original.shape[1], s)):
+            img = np.transpose(image[y:y+p*2+s,x:x+p*2+s,:] / 127.5 - 1.0, (2, 0, 1))[np.newaxis].astype(np.float32)
             *_, repro = self.model.predict(img)
-        repro = np.transpose(repro[0] + 0.5, (1, 2, 0)).clip(0.0, 1.0)
+            output[y*z:(y+s)*z,x*z:(x+s)*z,:] = np.transpose(repro[0] + 1.0, (1, 2, 0))[p*z:-p*z,p*z:-p*z,:]
-        return scipy.misc.toimage(repro * 255.0, cmin=0, cmax=255)
+            print('.', end='', flush=True)
+        return scipy.misc.toimage(output * 127.5, cmin=0, cmax=255)
 
 
 if __name__ == "__main__":
-    enhancer = NeuralEnhancer()
-
     if args.train:
+        args.zoom = 2**(args.generator_upscale - args.generator_downscale)
+        enhancer = NeuralEnhancer(loader=True)
         enhancer.train()
     else:
+        enhancer = NeuralEnhancer(loader=False)
         for filename in args.files:
-            print(filename)
+            print(filename, end=' ')
             img = scipy.ndimage.imread(filename, mode='RGB')
-            if img.shape[0] * img.shape[1] > 256 ** 2 and args.scales >= 2:
-                error('This file is (probably) too large to process in one shot and was ignored.',
-                      '  - Until tiled rendering is added, edit this code at your own peril!')
-                continue
-
             out = enhancer.process(img)
-            out.save(os.path.splitext(filename)[0]+'_ne%ix.png'%(2**args.scales))
+            out.save(os.path.splitext(filename)[0]+'_ne%ix.png' % args.zoom)
+            print(flush=True)
         print(ansi.ENDC)

scripts/small-1x.sh

@@ -0,0 +1,18 @@
+#!/bin/sh
+
+python3.4 enhance.py \
+    --train "$OPEN_IMAGES_PATH/*/*.jpg" --model small \
+    --epochs=50 --batch-shape=192 --device=gpu0 \
+    --generator-downscale=2 --generator-upscale=2 \
+    --generator-blocks=8 --generator-filters=64 \
+    --perceptual-layer=conv2_2 --smoothness-weight=1e7 --adversary-weight=0.0 \
+    --train-blur=3 --train-noise=5.0
+
+python3.4 enhance.py \
+    --train "$OPEN_IMAGES_PATH/*/*.jpg" --model small \
+    --epochs=500 --batch-shape=192 --device=gpu0 \
+    --generator-downscale=2 --generator-upscale=2 \
+    --perceptual-layer=conv5_2 --smoothness-weight=2e4 --adversary-weight=2e2 \
+    --generator-start=5 --discriminator-start=0 --adversarial-start=5 \
+    --discriminator-size=32 \
+    --train-blur=3 --train-noise=5.0

scripts/small-2x.sh

@@ -0,0 +1,18 @@
+#!/bin/sh
+
+python3.4 enhance.py \
+    --train "$OPEN_IMAGES_PATH/*/*.jpg" --model small \
+    --epochs=50 --batch-shape=192 --device=gpu0 \
+    --generator-downscale=1 --generator-upscale=2 \
+    --generator-blocks=8 --generator-filters=64 \
+    --perceptual-layer=conv2_2 --smoothness-weight=1e7 --adversary-weight=0.0 \
+    --train-blur=2 --train-noise=4.0
+
+python3.4 enhance.py \
+    --train "$OPEN_IMAGES_PATH/*/*.jpg" --model small \
+    --epochs=500 --batch-shape=192 --device=gpu0 \
+    --generator-downscale=1 --generator-upscale=2 \
+    --perceptual-layer=conv5_2 --smoothness-weight=2e4 --adversary-weight=2e2 \
+    --generator-start=5 --discriminator-start=0 --adversarial-start=5 \
+    --discriminator-size=32 \
+    --train-blur=2 --train-noise=4.0

scripts/small-4x.sh

@@ -0,0 +1,18 @@
+#!/bin/sh
+
+python3.4 enhance.py \
+    --train "$OPEN_IMAGES_PATH/*/*.jpg" --model small \
+    --epochs=50 --batch-shape=192 --device=gpu0 \
+    --generator-downscale=0 --generator-upscale=2 \
+    --generator-blocks=8 --generator-filters=64 \
+    --perceptual-layer=conv2_2 --smoothness-weight=1e7 --adversary-weight=0.0 \
+    --train-blur=1 --train-noise=3.0
+
+python3.4 enhance.py \
+    --train "$OPEN_IMAGES_PATH/*/*.jpg" --model small \
+    --epochs=500 --batch-shape=192 --device=gpu0 \
+    --generator-downscale=0 --generator-upscale=2 \
+    --perceptual-layer=conv5_2 --smoothness-weight=2e4 --adversary-weight=2e2 \
+    --generator-start=5 --discriminator-start=0 --adversarial-start=5 \
+    --discriminator-size=32 \
+    --train-blur=1 --train-noise=3.0