soravux · November 23, 2017 00:39 · Nov 23, 2017
diff --git a/match_wb.py b/match_wb.py
@@ -0,0 +1,73 @@
+import numpy as np
+from scipy.misc import imread, imsave
+
+from matplotlib import pyplot as plt
+
+
+# Get a linear version of the images
+target_im = (imread("lighting_gt.png").astype('float32') / 255.)**2.2 # Poor man un-sRGB
+source_im = (imread("warped_0150.png").astype('float32') / 255.)**2.2
+
+# This will not work robustly
+# # Get a pixel that *seems* white
+# target_luminance = target_im.dot([0.299, 0.587, 0.114])
+# source_luminance = source_im.dot([0.299, 0.587, 0.114])
+
+# target_white_gray = np.percentile(target_luminance[target_luminance < 1.], 80)
+# src_white_gray = np.percentile(source_luminance[source_luminance < 1.], 80)
+
+# target_idx_h, target_idx_w = np.where(target_luminance == target_white_gray)
+# src_idx_h, src_idx_w = np.where(source_luminance == src_white_gray)
+
+
+mouse_x, mouse_y = -1, -1
+def onclick(event):
+    global mouse_x, mouse_y
+    mouse_x = event.xdata
+    mouse_y = event.ydata
+
+# Ask for point in target image
+fig = plt.figure()
+plt.imshow(target_im**(1./2.2))
+cid = fig.canvas.mpl_connect('button_press_event', onclick)
+plt.show()
+
+target_idx_h = int(mouse_y)
+target_idx_w = int(mouse_x)
+
+
+# Ask for point in source image
+fig = plt.figure()
+plt.imshow(source_im**(1./2.2))
+cid = fig.canvas.mpl_connect('button_press_event', onclick)
+plt.show()
+
+src_idx_h = int(mouse_y)
+src_idx_w = int(mouse_x)
+
+
+# Get area of 5x5 pixels around the selected positions
+target_white = np.mean(target_im[target_idx_h-2:target_idx_h+3, target_idx_w-2:target_idx_w+3, :], axis=(0, 1))
+src_white = np.mean(source_im[src_idx_h-2:src_idx_h+3, src_idx_w-2:src_idx_w+3, :], axis=(0, 1))
+
+# Normalize the colors (we don't want to change brightness, just color correction)
+target_white /= target_white.sum()
+src_white /= src_white.sum()
+
+print("target white:", target_white)
+print("source white:", src_white)
+
+# Apply white balance correction
+fix_matrix = np.diag(target_white/src_white)
+source_fixed = source_im.dot(fix_matrix)
+print("Correction matrix:\n", fix_matrix)
+
+
+plt.subplot(221); plt.imshow(target_im**(1./2.2)); plt.scatter(target_idx_w, target_idx_h, s=1); plt.axis('off')
+plt.subplot(222); plt.imshow(source_im**(1./2.2)); plt.scatter(src_idx_w, src_idx_h, s=1); plt.axis('off')
+plt.subplot(223); plt.imshow(np.clip(source_fixed, 0, 1)**(1./2.2)); plt.axis('off')
+plt.show()
+
+# Save back as sRGB
+source_fixed_sRGB = np.clip((255.*source_fixed**(1./2.2)), 0, 255).astype('uint8')
+imsave("warped_0150_fixed.png", source_fixed_sRGB)