From a69d2ae0604b254ddef7e7bd7cc5a2ca002504cb Mon Sep 17 00:00:00 2001 From: "Alex J. Champandard" Date: Thu, 10 Nov 2016 16:45:30 +0100 Subject: [PATCH] Add histogram color matching, enabled with `--rendering-histogram`. --- README.rst | 2 +- enhance.py | 18 +++++++++++++++++- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/README.rst b/README.rst index 4e25eb7..e51e2b3 100644 --- a/README.rst +++ b/README.rst @@ -51,7 +51,7 @@ A list of example command lines you can use with the pre-trained models provided Here's a list of currently supported models, image types, and zoom levels in one table. ================== ===================== ==================== ===================== ==================== - ``--model=default`` ``--model=repair`` ``--model=denoise`` ``--model=deblur`` + FEATURES ``--model=default`` ``--model=repair`` ``--model=denoise`` ``--model=deblur`` ================== ===================== ==================== ===================== ==================== ``--type=photo`` 2x 1x … … ================== ===================== ==================== ===================== ==================== diff --git a/enhance.py b/enhance.py index 0290617..4f23967 100755 --- a/enhance.py +++ b/enhance.py @@ -39,6 +39,7 @@ 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('--rendering-histogram',default=False, action='store_true', help='Match color histogram of output to input.') add_arg('--type', default='photo', type=str, help='Name of the neural network to load/save.') add_arg('--model', default='default', type=str, help='Specific trained version of the model.') add_arg('--train', default=False, type=str, help='File pattern to load for training.') @@ -532,6 +533,14 @@ def train(self): self.model.save_generator() print(ansi.ENDC) + def match_histograms(self, A, B, rng=(0.0, 255.0), bins=64): + (Ha, Xa), (Hb, Xb) = [np.histogram(i, bins=bins, range=rng, density=True) for i in [A, B]] + X = np.linspace(rng[0], rng[1], bins, endpoint=True) + Hpa, Hpb = [np.cumsum(i) * (rng[1] - rng[0]) ** 2 / float(bins) for i in [Ha, Hb]] + inv_Ha = scipy.interpolate.interp1d(X, Hpa, bounds_error=False) + map_Hb = scipy.interpolate.interp1d(Hpb, X, bounds_error=False) + return map_Hb(inv_Ha(A)) + def process(self, original): # Snap the image to a shape that's compatible with the generator (2x, 4x) s = 2 ** max(args.generator_upscale, args.generator_downscale) @@ -549,7 +558,14 @@ def process(self, original): *_, repro = self.model.predict(img) output[y*z:(y+s)*z,x*z:(x+s)*z,:] = np.transpose(repro[0] + 0.5, (1, 2, 0))[p*z:-p*z,p*z:-p*z,:] print('.', end='', flush=True) - return scipy.misc.toimage(output.clip(0.0, 1.0) * 255.0, cmin=0, cmax=255) + output = output.clip(0.0, 1.0) * 255.0 + + # Match color histograms if the user specified this option. + if args.rendering_histogram: + for i in range(3): + output[:,:,i] = self.match_histograms(output[:,:,i], original[:,:,i]) + + return scipy.misc.toimage(output, cmin=0, cmax=255) if __name__ == "__main__":