Use CIE2000 to determine color distances, compute gradients in Lab

Slow

tools/generate_gradients.py

@@ -8,6 +8,9 @@ import json
 from powerline.colorscheme import cterm_to_hex
 from itertools import groupby
 import argparse
+from colormath.color_objects import sRGBColor, LabColor
+from colormath.color_conversions import convert_color
+from colormath.color_diff import delta_e_cie2000
 
 try:
 	from __builtin__ import unicode
@@ -22,15 +25,20 @@ def num2(s):
 		return (False, [float(v) for v in s.partition(' ')[::2]])
 
 
-def rgbint_to_rgb(rgbint):
-	return ((rgbint >> 16) & 0xFF, (rgbint >> 8) & 0xFF, rgbint & 0xFF)
+def rgbint_to_lab(rgbint):
+	rgb = sRGBColor((rgbint >> 16) & 0xFF, (rgbint >> 8) & 0xFF, rgbint & 0xFF,
+			is_upscaled=True)
+	return convert_color(rgb, LabColor)
+
+
+cterm_to_lab = tuple((rgbint_to_lab(v) for v in cterm_to_hex))
 
 
 def color(s):
 	if len(s) <= 3:
-		return rgbint_to_rgb(cterm_to_hex[int(s)])
+		return cterm_to_lab[int(s)]
 	else:
-		return rgbint_to_rgb(int(s, 16))
+		return rgbint_to_lab(int(s, 16))
 
 
 def nums(s):
@@ -53,34 +61,42 @@ def linear_gradient(start_value, stop_value, start_offset, stop_offset, offset):
 	return start_value + ((offset - start_offset) * (stop_value - start_value) / (stop_offset - start_offset))
 
 
-def gradient(DATA):
+def lab_gradient(slab, elab, soff, eoff, off):
+	svals = slab.get_value_tuple()
+	evals = elab.get_value_tuple()
+	return LabColor(*[linear_gradient(start_value, end_value, soff, eoff, off)
+						for start_value, end_value in zip(svals, evals)])
+
+
+def generate_gradient_function(DATA):
 	def gradient_function(y):
 		initial_offset = 0
 		for offset, start, end in DATA:
 			if y <= offset:
-				return [linear_gradient(start[i], end[i], initial_offset, offset, y) for i in range(3)]
+				return lab_gradient(start, end, initial_offset, offset, y)
 			initial_offset = offset
 	return gradient_function
 
 
-def get_rgb(*args):
-	return "%02x%02x%02x" % args
+def get_upscaled_values(rgb):
+	return [min(max(0, i), 255) for i in rgb.get_upscaled_value_tuple()]
 
 
-def col_distance(rgb1, rgb2):
-	return sum(((rgb1[i] - rgb2[i]) ** 2 for i in range(3)))
+def get_rgb(lab):
+	rgb = convert_color(lab, sRGBColor)
+	rgb = sRGBColor(*get_upscaled_values(rgb), is_upscaled=True)
+	return rgb.get_rgb_hex()[1:]
 
 
-def find_color(urgb, colors, ctrans):
-	cur_distance = 3 * (255 ** 2 + 1)
+def find_color(ulab, colors, ctrans):
+	cur_distance = float('inf')
 	cur_color = None
 	i = 0
-	for crgbint in colors:
-		crgb = rgbint_to_rgb(crgbint)
-		dist = col_distance(urgb, crgb)
+	for clab in colors:
+		dist = delta_e_cie2000(ulab, clab)
 		if dist < cur_distance:
 			cur_distance = dist
-			cur_color = (ctrans(i), crgb)
+			cur_color = (ctrans(i), clab)
 		i += 1
 	return cur_color
 
@@ -89,7 +105,8 @@ def print_color(color):
 	if type(color) is int:
 		colstr = '5;' + str(color)
 	else:
-		colstr = '2;' + ';'.join((str(int(round(i))) for i in color))
+		rgb = convert_color(color, sRGBColor)
+		colstr = '2;' + ';'.join((str(i) for i in get_upscaled_values(rgb)))
 	sys.stdout.write('\033[48;' + colstr + 'm ')
 
 
@@ -131,15 +148,15 @@ else:
 	steps = [i * step for i in range(1, maxweight + 1)]
 
 data = [(weight, args.gradient[i - 1], args.gradient[i]) for weight, i in zip(steps, range(1, len(args.gradient)))]
-gr_func = gradient(data)
+gr_func = generate_gradient_function(data)
 gradient = [gr_func(y) for y in range(0, m)]
 palettes = {
-	'16': (cterm_to_hex[:16], lambda c: c),
-	'256': (cterm_to_hex, lambda c: c),
-	None: (cterm_to_hex[16:], lambda c: c + 16),
+	'16': (cterm_to_lab[:16], lambda c: c),
+	'256': (cterm_to_lab, lambda c: c),
+	None: (cterm_to_lab[16:], lambda c: c + 16),
 }
-r = [get_rgb(*col) for col in gradient]
-r2 = [find_color(col, *palettes[args.palette])[0] for col in gradient]
+r = [get_rgb(lab) for lab in gradient]
+r2 = [find_color(lab, *palettes[args.palette])[0] for lab in gradient]
 r3 = [i[0] for i in groupby(r2)]
 print(json.dumps(r))
 print(json.dumps(r2))