utils: raspberrypi: ctt: Improved color matrix fitting

Added code which optimises the color matrices based off
delta E values for the calibration images. Working in LAB
color space.

Signed-off-by: Ben Benson <ben.benson@raspberrypi.com>
Reviewed-by: David Plowman <david.plowman@raspberrypi.com>
Reviewed-by: Naushir Patuck <naush@raspberrypi.com>
Signed-off-by: Naushir Patuck <naush@raspberrypi.com>

3 files changed, 300 insertions, 38 deletions

diff --git a/utils/raspberrypi/ctt/colors.py b/utils/raspberrypi/ctt/colors.py
new file mode 100644
index 00000000..1ab986d6
--- /dev/null
+++ b/utils/raspberrypi/ctt/colors.py
@@ -0,0 +1,30 @@
+# colors.py - Program to convert from RGB to LAB color space
+def RGB_to_LAB(RGB):  # where RGB is a 1x3 array.   e.g RGB = [100, 255, 230]
+    num = 0
+    XYZ = [0, 0, 0]
+    # converted all the three R, G, B to X, Y, Z
+    X = RGB[0] * 0.4124 + RGB[1] * 0.3576 + RGB[2] * 0.1805
+    Y = RGB[0] * 0.2126 + RGB[1] * 0.7152 + RGB[2] * 0.0722
+    Z = RGB[0] * 0.0193 + RGB[1] * 0.1192 + RGB[2] * 0.9505
+
+    XYZ[0] = X / 255 * 100
+    XYZ[1] = Y / 255 * 100  # XYZ Must be in range 0 -> 100, so scale down from 255
+    XYZ[2] = Z / 255 * 100
+    XYZ[0] = XYZ[0] / 95.047  # ref_X =  95.047   Observer= 2°, Illuminant= D65
+    XYZ[1] = XYZ[1] / 100.0  # ref_Y = 100.000
+    XYZ[2] = XYZ[2] / 108.883  # ref_Z = 108.883
+    num = 0
+    for value in XYZ:
+        if value > 0.008856:
+            value = value ** (0.3333333333333333)
+        else:
+            value = (7.787 * value) + (16 / 116)
+        XYZ[num] = value
+        num = num + 1
+
+    # L, A, B, values calculated below
+    L = (116 * XYZ[1]) - 16
+    a = 500 * (XYZ[0] - XYZ[1])
+    b = 200 * (XYZ[1] - XYZ[2])
+
+    return [L, a, b]
diff --git a/utils/raspberrypi/ctt/ctt_ccm.py b/utils/raspberrypi/ctt/ctt_ccm.py
index 376cc712..49159535 100644
--- a/utils/raspberrypi/ctt/ctt_ccm.py
+++ b/utils/raspberrypi/ctt/ctt_ccm.py
@@ -6,27 +6,66 @@
 
 from ctt_image_load import *
 from ctt_awb import get_alsc_patches
-
-
+import colors
+from scipy.optimize import minimize
+from ctt_visualise import visualise_macbeth_chart
+import numpy as np
 """
 takes 8-bit macbeth chart values, degammas and returns 16 bit
 """
+
+'''
+This program has many options from which to derive the color matrix from.
+The first is average. This minimises the average delta E across all patches of
+the macbeth chart. Testing across all cameras yeilded this as the most color
+accurate and vivid. Other options are avalible however.
+Maximum minimises the maximum Delta E of the patches. It iterates through till
+a minimum maximum is found (so that there is
+not one patch that deviates wildly.)
+This yields generally good results but overall the colors are less accurate
+Have a fiddle with maximum and see what you think.
+The final option allows you to select the patches for which to average across.
+This means that you can bias certain patches, for instance if you want the
+reds to be more accurate.
+'''
+
+matrix_selection_types = ["average", "maximum", "patches"]
+typenum = 0  # select from array above, 0 = average, 1 = maximum, 2 = patches
+test_patches = [1, 2, 5, 8, 9, 12, 14]
+
+'''
+Enter patches to test for. Can also be entered twice if you
+would like twice as much bias on one patch.
+'''
+
+
 def degamma(x):
-    x = x / ((2**8)-1)
-    x = np.where(x < 0.04045, x/12.92, ((x+0.055)/1.055)**2.4)
-    x = x * ((2**16)-1)
+    x = x / ((2 ** 8) - 1)  # takes 255 and scales it down to one
+    x = np.where(x < 0.04045, x / 12.92, ((x + 0.055) / 1.055) ** 2.4)
+    x = x * ((2 ** 16) - 1)  # takes one and scales up to 65535, 16 bit color
     return x
 
 
+def gamma(x):
+    # return (x *  * (1 / 2.4)  *  1.055 - 0.055)
+    e = []
+    for i in range(len(x)):
+        e.append(((x[i] / 255) ** (1 / 2.4) * 1.055 - 0.055) * 255)
+    return e
+
+
 """
 FInds colour correction matrices for list of images
 """
+
+
 def ccm(Cam, cal_cr_list, cal_cb_list):
+    global matrix_selection_types, typenum
     imgs = Cam.imgs
     """
     standard macbeth chart colour values
     """
-    m_rgb = np.array([  # these are in sRGB
+    m_rgb = np.array([  # these are in RGB
         [116, 81, 67],    # dark skin
         [199, 147, 129],  # light skin
         [91, 122, 156],   # blue sky
@@ -34,7 +73,7 @@ def ccm(Cam, cal_cr_list, cal_cb_list):
         [130, 128, 176],  # blue flower
         [92, 190, 172],   # bluish green
         [224, 124, 47],   # orange
-        [68, 91, 170],     # purplish blue
+        [68, 91, 170],    # purplish blue
         [198, 82, 97],    # moderate red
         [94, 58, 106],    # purple
         [159, 189, 63],   # yellow green
@@ -52,16 +91,22 @@ def ccm(Cam, cal_cr_list, cal_cb_list):
         [82, 84, 86],     # neutral 3.5
         [49, 49, 51]      # black 2
     ])
-
     """
     convert reference colours from srgb to rgb
     """
-    m_srgb = degamma(m_rgb)
+    m_srgb = degamma(m_rgb)  # now in 16 bit color.
+
+    m_lab = []
+    for col in m_srgb:
+        m_lab.append(colors.RGB_to_LAB(col / 256))
+    # This produces matrix of LAB values for ideal color chart)
+
     """
     reorder reference values to match how patches are ordered
     """
     m_srgb = np.array([m_srgb[i::6] for i in range(6)]).reshape((24, 3))
-
+    m_lab = np.array([m_lab[i::6] for i in range(6)]).reshape((24, 3))
+    m_rgb = np.array([m_rgb[i::6] for i in range(6)]).reshape((24, 3))
     """
     reformat alsc correction tables or set colour_cals to None if alsc is
     deactivated
@@ -76,8 +121,8 @@ def ccm(Cam, cal_cr_list, cal_cb_list):
             """
             normalise tables so min value is 1
             """
-            cr_tab = cr_tab/np.min(cr_tab)
-            cb_tab = cb_tab/np.min(cb_tab)
+            cr_tab = cr_tab / np.min(cr_tab)
+            cb_tab = cb_tab / np.min(cb_tab)
             colour_cals[cr['ct']] = [cr_tab, cb_tab]
 
     """
@@ -94,6 +139,8 @@ def ccm(Cam, cal_cr_list, cal_cb_list):
         the function will simply return the macbeth patches
         """
         r, b, g = get_alsc_patches(Img, colour_cals, grey=False)
+        # 256 values for each patch of sRGB values
+
         """
         do awb
         Note: awb is done by measuring the macbeth chart in the image, rather
@@ -101,34 +148,123 @@ def ccm(Cam, cal_cr_list, cal_cb_list):
         and the ccm matrices will be more accurate.
         """
         r_greys, b_greys, g_greys = r[3::4], b[3::4], g[3::4]
-        r_g = np.mean(r_greys/g_greys)
-        b_g = np.mean(b_greys/g_greys)
+        r_g = np.mean(r_greys / g_greys)
+        b_g = np.mean(b_greys / g_greys)
         r = r / r_g
         b = b / b_g
-
         """
         normalise brightness wrt reference macbeth colours and then average
         each channel for each patch
         """
-        gain = np.mean(m_srgb)/np.mean((r, g, b))
+        gain = np.mean(m_srgb) / np.mean((r, g, b))
         Cam.log += '\nGain with respect to standard colours: {:.3f}'.format(gain)
-        r = np.mean(gain*r, axis=1)
-        b = np.mean(gain*b, axis=1)
-        g = np.mean(gain*g, axis=1)
-
+        r = np.mean(gain * r, axis=1)
+        b = np.mean(gain * b, axis=1)
+        g = np.mean(gain * g, axis=1)
         """
         calculate ccm matrix
         """
+        # ==== All of below should in sRGB ===##
+        sumde = 0
         ccm = do_ccm(r, g, b, m_srgb)
+        # This is the initial guess that our optimisation code works with.
+
+        r1 = ccm[0]
+        r2 = ccm[1]
+        g1 = ccm[3]
+        g2 = ccm[4]
+        b1 = ccm[6]
+        b2 = ccm[7]
+        '''
+        COLOR MATRIX LOOKS AS BELOW
+        R1 R2 R3   Rval     Outr
+        G1 G2 G3  *  Gval  =  G
+        B1 B2 B3   Bval     B
+        Will be optimising 6 elements and working out the third element using 1-r1-r2 = r3
+        '''
+
+        x0 = [r1, r2, g1, g2, b1, b2]
+        '''
+        We use our old CCM as the initial guess for the program to find the
+        optimised matrix
+        '''
+        result = minimize(guess, x0, args=(r, g, b, m_lab), tol=0.0000000001)
+        '''
+        This produces a color matrix which has the lowest delta E possible,
+        based off the input data. Note it is impossible for this to reach
+        zero since the input data is imperfect
+        '''
+
+        Cam.log += ("\n \n Optimised Matrix Below: \n \n")
+        [r1, r2, g1, g2, b1, b2] = result.x
+        # The new, optimised color correction matrix values
+        optimised_ccm = [r1, r2, (1 - r1 - r2), g1, g2, (1 - g1 - g2), b1, b2, (1 - b1 - b2)]
+        # This is the optimised Color Matrix (preserving greys by summing rows up to 1)
+        Cam.log += str(optimised_ccm)
+        Cam.log += "\n Old Color Correction Matrix Below \n"
+        Cam.log += str(ccm)
+
+        formatted_ccm = np.array(ccm).reshape((3, 3))
+
+        '''
+        below is a whole load of code that then applies the latest color
+        matrix, and returns LAB values for color. This can then be used
+        to calculate the final delta E
+        '''
+        optimised_ccm_rgb = []  # Original Color Corrected Matrix RGB / LAB
+        optimised_ccm_lab = []
+        for w in range(24):
+            RGB = np.array([r[w], g[w], b[w]])
+            ccm_applied_rgb = np.dot(formatted_ccm, (RGB / 256))
+            optimised_ccm_rgb.append(gamma(ccm_applied_rgb))
+            optimised_ccm_lab.append(colors.RGB_to_LAB(ccm_applied_rgb))
+
+        formatted_optimised_ccm = np.array(ccm).reshape((3, 3))
+        after_gamma_rgb = []
+        after_gamma_lab = []
+        for w in range(24):
+            RGB = np.array([r[w], g[w], b[w]])
+            optimised_ccm_applied_rgb = np.dot(formatted_optimised_ccm, RGB / 256)
+            after_gamma_rgb.append(gamma(optimised_ccm_applied_rgb))
+            after_gamma_lab.append(colors.RGB_to_LAB(optimised_ccm_applied_rgb))
+        '''
+        Gamma After RGB / LAB
+        We now want to spit out some data that shows
+        how the optimisation has improved the color matrices
+        '''
+        Cam.log += "Here are the Improvements"
+
+        # CALCULATE WORST CASE delta e
+        old_worst_delta_e = 0
+        before_average = transform_and_evaluate(formatted_ccm, r, g, b, m_lab)
+        new_worst_delta_e = 0
+        after_average = transform_and_evaluate(formatted_optimised_ccm, r, g, b, m_lab)
+        for i in range(24):
+            old_delta_e = deltae(optimised_ccm_lab[i], m_lab[i])  # Current Old Delta E
+            new_delta_e = deltae(after_gamma_lab[i], m_lab[i])  # Current New Delta E
+            if old_delta_e > old_worst_delta_e:
+                old_worst_delta_e = old_delta_e
+            if new_delta_e > new_worst_delta_e:
+                new_worst_delta_e = new_delta_e
+
+        Cam.log += "Before color correction matrix was optimised, we got an average delta E of " + str(before_average) + " and a maximum delta E of " + str(old_worst_delta_e)
+        Cam.log += "After color correction matrix was optimised, we got an average delta E of " + str(after_average) + " and a maximum delta E of " + str(new_worst_delta_e)
+
+        visualise_macbeth_chart(m_rgb, optimised_ccm_rgb, after_gamma_rgb, str(Img.col) + str(matrix_selection_types[typenum]))
+        '''
+        The program will also save some visualisations of improvements.
+        Very pretty to look at. Top rectangle is ideal, Left square is
+        before optimisation, right square is after.
+        '''
 
         """
         if a ccm has already been calculated for that temperature then don't
         overwrite but save both. They will then be averaged later on
-        """
+        """  # Now going to use optimised color matrix, optimised_ccm
         if Img.col in ccm_tab.keys():
-            ccm_tab[Img.col].append(ccm)
+            ccm_tab[Img.col].append(optimised_ccm)
         else:
-            ccm_tab[Img.col] = [ccm]
+            ccm_tab[Img.col] = [optimised_ccm]
         Cam.log += '\n'
 
     Cam.log += '\nFinished processing images'
@@ -137,8 +273,8 @@ def ccm(Cam, cal_cr_list, cal_cb_list):
     """
     for k, v in ccm_tab.items():
         tab = np.mean(v, axis=0)
-        tab = np.where((10000*tab) % 1 <= 0.05, tab+0.00001, tab)
-        tab = np.where((10000*tab) % 1 >= 0.95, tab-0.00001, tab)
+        tab = np.where((10000 * tab) % 1 <= 0.05, tab + 0.00001, tab)
+        tab = np.where((10000 * tab) % 1 >= 0.95, tab - 0.00001, tab)
         ccm_tab[k] = list(np.round(tab, 5))
         Cam.log += '\nMatrix calculated for colour temperature of {} K'.format(k)
 
@@ -156,20 +292,67 @@ def ccm(Cam, cal_cr_list, cal_cb_list):
     return ccms
 
 
+def guess(x0, r, g, b, m_lab):       # provides a method of numerical feedback for the optimisation code
+    [r1, r2, g1, g2, b1, b2] = x0
+    ccm = np.array([r1, r2, (1 - r1 - r2),
+                    g1, g2, (1 - g1 - g2),
+                    b1, b2, (1 - b1 - b2)]).reshape((3, 3))  # format the matrix correctly
+    return transform_and_evaluate(ccm, r, g, b, m_lab)
+
+
+def transform_and_evaluate(ccm, r, g, b, m_lab):  # Transforms colors to LAB and applies the correction matrix
+    # create list of matrix changed colors
+    realrgb = []
+    for i in range(len(r)):
+        RGB = np.array([r[i], g[i], b[i]])
+        rgb_post_ccm = np.dot(ccm, RGB)  # This is RGB values after the color correction matrix has been applied
+        realrgb.append(colors.RGB_to_LAB(rgb_post_ccm))
+    # now compare that with m_lab and return numeric result, averaged for each patch
+    return (sumde(realrgb, m_lab) / 24)  # returns an average result of delta E
+
+
+def sumde(listA, listB):
+    global typenum, test_patches
+    sumde = 0
+    maxde = 0
+    patchde = []
+    for i in range(len(listA)):
+        if maxde < (deltae(listA[i], listB[i])):
+            maxde = deltae(listA[i], listB[i])
+        patchde.append(deltae(listA[i], listB[i]))
+        sumde += deltae(listA[i], listB[i])
+    '''
+    The different options specified at the start allow for
+    the maximum to be returned, average or specific patches
+    '''
+    if typenum == 0:
+        return sumde
+    if typenum == 1:
+        return maxde
+    if typenum == 2:
+        output = 0
+        for y in range(len(test_patches)):
+            output += patchde[test_patches[y]]   # grabs the specific patches (no need for averaging here)
+        return output
+
+
 """
 calculates the ccm for an individual image.
-ccms are calculate in rgb space, and are fit by hand. Although it is a 3x3
+ccms are calculated in rgb space, and are fit by hand. Although it is a 3x3
 matrix, each row must add up to 1 in order to conserve greyness, simplifying
 calculation.
-Should you want to fit them in another space (e.g. LAB) we wish you the best of
-luck and send us the code when you are done! :-)
+The initial CCM is calculated in RGB, and then optimised in LAB color space
+This simplifies the initial calculation but then gets us the accuracy of
+using LAB color space.
 """
+
+
 def do_ccm(r, g, b, m_srgb):
     rb = r-b
     gb = g-b
-    rb_2s = (rb*rb)
-    rb_gbs = (rb*gb)
-    gb_2s = (gb*gb)
+    rb_2s = (rb * rb)
+    rb_gbs = (rb * gb)
+    gb_2s = (gb * gb)
 
     r_rbs = rb * (m_srgb[..., 0] - b)
     r_gbs = gb * (m_srgb[..., 0] - b)
@@ -191,7 +374,7 @@ def do_ccm(r, g, b, m_srgb):
     b_rb = np.sum(b_rbs)
     b_gb = np.sum(b_gbs)
 
-    det = rb_2*gb_2 - rb_gb*rb_gb
+    det = rb_2 * gb_2 - rb_gb * rb_gb
 
     """
     Raise error if matrix is singular...
@@ -201,19 +384,19 @@ def do_ccm(r, g, b, m_srgb):
     if det < 0.001:
         raise ArithmeticError
 
-    r_a = (gb_2*r_rb - rb_gb*r_gb)/det
-    r_b = (rb_2*r_gb - rb_gb*r_rb)/det
+    r_a = (gb_2 * r_rb - rb_gb * r_gb) / det
+    r_b = (rb_2 * r_gb - rb_gb * r_rb) / det
     """
     Last row can be calculated by knowing the sum must be 1
     """
     r_c = 1 - r_a - r_b
 
-    g_a = (gb_2*g_rb - rb_gb*g_gb)/det
-    g_b = (rb_2*g_gb - rb_gb*g_rb)/det
+    g_a = (gb_2 * g_rb - rb_gb * g_gb) / det
+    g_b = (rb_2 * g_gb - rb_gb * g_rb) / det
     g_c = 1 - g_a - g_b
 
-    b_a = (gb_2*b_rb - rb_gb*b_gb)/det
-    b_b = (rb_2*b_gb - rb_gb*b_rb)/det
+    b_a = (gb_2 * b_rb - rb_gb * b_gb) / det
+    b_b = (rb_2 * b_gb - rb_gb * b_rb) / det
     b_c = 1 - b_a - b_b
 
     """
@@ -222,3 +405,9 @@ def do_ccm(r, g, b, m_srgb):
     ccm = [r_a, r_b, r_c, g_a, g_b, g_c, b_a, b_b, b_c]
 
     return ccm
+
+
+def deltae(colorA, colorB):
+    return ((colorA[0] - colorB[0]) ** 2 + (colorA[1] - colorB[1]) ** 2 + (colorA[2] - colorB[2]) ** 2) ** 0.5
+    # return ((colorA[1]-colorB[1]) *  * 2 + (colorA[2]-colorB[2]) *  * 2) *  * 0.5
+    # UNCOMMENT IF YOU WANT TO NEGLECT LUMINANCE FROM CALCULATION OF DELTA E
diff --git a/utils/raspberrypi/ctt/ctt_visualise.py b/utils/raspberrypi/ctt/ctt_visualise.py
new file mode 100644
index 00000000..ed2339fd
--- /dev/null
+++ b/utils/raspberrypi/ctt/ctt_visualise.py
@@ -0,0 +1,43 @@
+"""
+Some code that will save virtual macbeth charts that show the difference between optimised matrices and non optimised matrices
+
+The function creates an image that is 1550 by 1050 pixels wide, and fills it with patches which are 200x200 pixels in size
+Each patch contains the ideal color, the color from the original matrix, and the color from the final matrix
+_________________
+|               |
+|  Ideal Color  |
+|_______________|
+|  Old  |  new  |
+| Color | Color |
+|_______|_______|
+
+Nice way of showing how the optimisation helps change the colors and the color matricies
+"""
+import numpy as np
+from PIL import Image
+
+
+def visualise_macbeth_chart(macbeth_rgb, original_rgb, new_rgb, output_filename):
+    image = np.zeros((1050, 1550, 3), dtype=np.uint8)
+    colorindex = -1
+    for y in range(6):
+        for x in range(4):  # Creates 6 x 4 grid of macbeth chart
+            colorindex += 1
+            xlocation = 50 + 250 * x  # Means there is 50px of black gap between each square, more like the real macbeth chart.
+            ylocation = 50 + 250 * y
+            for g in range(200):
+                for i in range(100):
+                    image[xlocation + i, ylocation + g] = macbeth_rgb[colorindex]
+            xlocation = 150 + 250 * x
+            ylocation = 50 + 250 * y
+            for i in range(100):
+                for g in range(100):
+                    image[xlocation + i, ylocation + g] = original_rgb[colorindex]  # Smaller squares below to compare the old colors with the new ones
+            xlocation = 150 + 250 * x
+            ylocation = 150 + 250 * y
+            for i in range(100):
+                for g in range(100):
+                    image[xlocation + i, ylocation + g] = new_rgb[colorindex]
+
+    img = Image.fromarray(image, 'RGB')
+    img.save(str(output_filename) + 'Generated Macbeth Chart.png')