diff --git a/enhance.py b/enhance.py index bd1e587..ac54d36 100755 --- a/enhance.py +++ b/enhance.py @@ -41,6 +41,7 @@ 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('--seeds', default=False, type=str, help='File pattern to load for training seeds.') 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.') @@ -125,6 +126,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 @@ -140,7 +163,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`") @@ -153,11 +182,14 @@ def __init__(self): def run(self): while True: - random.shuffle(self.files) - for f in self.files: - self.add_to_buffer(f) + indices = list(range(0, len(self.files))) + random.shuffle(indices) + + for file_index in indices: + self.add_to_buffer(file_index) - def add_to_buffer(self, f): + def add_to_buffer(self, file_index): + f = self.files[file_index] filename = os.path.join(self.cwd, f) try: orig = PIL.Image.open(filename).convert('RGB') @@ -168,27 +200,55 @@ def add_to_buffer(self, f): except Exception as e: warn('Could not load `{}` as image.'.format(filename), ' - Try fixing or removing the file before next run.') - self.files.remove(f) + del self.files[file_index] + if self.seeds: + del self.seeds[file_index] 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) + # determine seed + if self.seeds: + # file based seed + f = self.seeds[file_index] + filename = os.path.join(self.cwd, f) + try: + seed = PIL.Image.open(filename).convert('RGB') + if any(s < self.seed_shape for s in seed.size): + raise ValueError('Image is too small for seed size (found {}, expected {})'.format(seed.size, self.seed_shape)) + except Exception as e: + warn('Could not load `{}` as seed image.'.format(filename), + ' - Try fixing or removing the file before next run. ({})'.format(e)) + del self.files[file_index] + del self.seeds[file_index] + return + else: + # synthetic seed + # optionally blur before scaling + seed = orig.filter(PIL.ImageFilter.GaussianBlur(radius=args.train_blur)) if args.train_blur else orig + # seed is scaled down version of original + seed = seed.resize((orig.size[0]//args.zoom, orig.size[1]//args.zoom), resample=PIL.Image.LANCZOS) + + # convert original to numpy float32 + orig = scipy.misc.fromimage(orig).astype(np.float32) - if args.train_jpeg: + # optionally convert to jpeg + if args.train_jpeg is not None: buffer = io.BytesIO() seed.save(buffer, format='jpeg', quality=args.train_jpeg+random.randrange(-15,+15)) seed = PIL.Image.open(buffer) + # convert seed to numpy float32 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) + # optionally add training noise to seed + if args.train_noise is not None: + seed += scipy.random.normal(scale=args.train_noise, size=(seed.shape[0], seed.shape[1], 1)) for _ in range(args.buffer_similar): + # compute seed displacement h = random.randint(0, seed.shape[0] - self.seed_shape) w = random.randint(0, seed.shape[1] - self.seed_shape) seed_chunk = seed[h:h+self.seed_shape, w:w+self.seed_shape] + # and matching image displacement h, w = h * args.zoom, w * args.zoom orig_chunk = orig[h:h+self.orig_shape, w:w+self.orig_shape]