summaryrefslogtreecommitdiffstats
path: root/tools/generate_gradients.py
blob: 290e75e4564a055f0fcefafdbf7b79bf268884f4 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
#!/usr/bin/env python
# vim:fileencoding=utf-8:noet

'''Gradients generator
'''

from __future__ import (unicode_literals, division, absolute_import, print_function)

import sys
import json
import argparse

from itertools import groupby

from colormath.color_objects import sRGBColor, LabColor
from colormath.color_conversions import convert_color
from colormath.color_diff import delta_e_cie2000

from powerline.colorscheme import cterm_to_hex


def num2(s):
	try:
		return (True, [int(v) for v in s.partition(' ')[::2]])
	except TypeError:
		return (False, [float(v) for v in s.partition(' ')[::2]])


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 cterm_to_lab[int(s)]
	else:
		return rgbint_to_lab(int(s, 16))


def nums(s):
	return [int(i) for i in s.split()]


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 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 lab_gradient(start, end, initial_offset, offset, y)
			initial_offset = offset
	return gradient_function


def get_upscaled_values(rgb):
	return [min(max(0, i), 255) for i in rgb.get_upscaled_value_tuple()]


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(ulab, colors, ctrans):
	cur_distance = float('inf')
	cur_color = None
	i = 0
	for clab in colors:
		dist = delta_e_cie2000(ulab, clab)
		if dist < cur_distance:
			cur_distance = dist
			cur_color = (ctrans(i), clab)
		i += 1
	return cur_color


def print_color(color):
	if type(color) is int:
		colstr = '5;' + str(color)
	else:
		rgb = convert_color(color, sRGBColor)
		colstr = '2;' + ';'.join((str(i) for i in get_upscaled_values(rgb)))
	sys.stdout.write('\033[48;' + colstr + 'm ')


def print_colors(colors, num):
	for i in range(num):
		color = colors[int(round(i * (len(colors) - 1) / num))]
		print_color(color)
	sys.stdout.write('\033[0m\n')


def dec_scale_generator(num):
	j = 0
	r = ''
	while num:
		r += '\033[{0}m'.format(j % 2)
		for i in range(10):
			r += str(i)
			num -= 1
			if not num:
				break
		j += 1
	r += '\033[0m\n'
	return r


def compute_steps(gradient, weights):
	maxweight = len(gradient) - 1
	if weights:
		weight_sum = sum(weights)
		norm_weights = [100.0 * weight / weight_sum for weight in weights]
		steps = [0]
		for weight in norm_weights:
			steps.append(steps[-1] + weight)
		steps.pop(0)
		steps.pop(0)
	else:
		step = m / maxweight
		steps = [i * step for i in range(1, maxweight + 1)]
	return steps


palettes = {
	'16': (cterm_to_lab[:16], lambda c: c),
	'256': (cterm_to_lab, lambda c: c),
	None: (cterm_to_lab[16:], lambda c: c + 16),
}


def show_scale(rng, num_output):
	if not rng and num_output >= 32 and (num_output - 1) // 10 >= 4 and (num_output - 1) % 10 == 0:
		sys.stdout.write('0')
		sys.stdout.write(''.join(('%*u' % (num_output // 10, i) for i in range(10, 101, 10))))
		sys.stdout.write('\n')
	else:
		if rng:
			vmin, vmax = rng[1]
			isint = rng[0]
		else:
			isint = True
			vmin = 0
			vmax = 100
		s = ''
		lasts = ' ' + str(vmax)
		while len(s) + len(lasts) < num_output:
			curpc = len(s) + 1 if s else 0
			curval = vmin + curpc * (vmax - vmin) / num_output
			if isint:
				curval = int(round(curval))
			s += str(curval) + ' '
		sys.stdout.write(s[:-1] + lasts + '\n')
	sys.stdout.write(dec_scale_generator(num_output) + '\n')


if __name__ == '__main__':
	p = argparse.ArgumentParser(description=__doc__)
	p.add_argument('gradient', nargs='*', metavar='COLOR', type=color, help='List of colors (either indexes from 8-bit palette or 24-bit RGB in hexadecimal notation)')
	p.add_argument('-n', '--num_items', metavar='INT', type=int, help='Number of items in resulting list', default=101)
	p.add_argument('-N', '--num_output', metavar='INT', type=int, help='Number of characters in sample', default=101)
	p.add_argument('-r', '--range', metavar='V1 V2', type=num2, help='Use this range when outputting scale')
	p.add_argument('-s', '--show', action='store_true', help='If present output gradient sample')
	p.add_argument('-p', '--palette', choices=('16', '256'), help='Use this palette. Defaults to 240-color palette (256 colors without first 16)')
	p.add_argument('-w', '--weights', metavar='INT INT ...', type=nums, help='Adjust weights of colors. Number of weights must be equal to number of colors')
	p.add_argument('-C', '--omit-terminal', action='store_true', help='If present do not compute values for terminal')

	args = p.parse_args()

	m = args.num_items

	steps = compute_steps(args.gradient, args.weights)

	data = [
		(weight, args.gradient[i - 1], args.gradient[i])
		for weight, i in zip(steps, range(1, len(args.gradient)))
	]
	gr_func = generate_gradient_function(data)
	gradient = [gr_func(y) for y in range(0, m)]

	r = [get_rgb(lab) for lab in gradient]
	if not args.omit_terminal:
		r2 = [find_color(lab, *palettes[args.palette])[0] for lab in gradient]
		r3 = [i[0] for i in groupby(r2)]

	if not args.omit_terminal:
		print(json.dumps(r3) + ',')
		print(json.dumps(r2) + ',')
	print(json.dumps(r))

	if args.show:
		print_colors(args.gradient, args.num_output)
		if not args.omit_terminal:
			print_colors(r3, args.num_output)
			print_colors(r2, args.num_output)
		print_colors(gradient, args.num_output)

		show_scale(args.range, args.num_output)