1 #!/usr/bin/env python 2 3 from random import choice, random, randrange 4 from os.path import splitext 5 import EXIF 6 import PIL.Image 7 import itertools 8 import math 9 import sys 10 11 corners = [ 12 (0, 0, 0), (255, 0, 0), (0, 255, 0), (255, 255, 0), 13 (0, 0, 255), (255, 0, 255), (0, 255, 255), (255, 255, 255) 14 ] 15 16 def distance(rgb1, rgb2): 17 r1, g1, b1 = rgb1 18 r2, g2, b2 = rgb2 19 return math.sqrt(pow(r1 - r2, 2) + pow(g1 - g2, 2) + pow(b1 - b2, 2)) 20 21 def nearest(rgb, values): 22 l = [(distance(rgb, value), value) for value in values] 23 total = sum([d for d, value in l]) 24 l = [(d / total, value) for d, value in l] 25 l.sort(reverse=True) 26 lc = [(d, value) for d, value in l if common(rgb, value)] 27 return get_from_distribution(rgb, lc or l) 28 29 def common(rgb1, rgb2): 30 return sum(map(lambda x: x[0] * x[1], zip(rgb1, rgb2))) != 0 31 32 def restore(srgb): 33 return tuple(map(lambda x: int(x * 255.0), srgb)) 34 35 def scale(rgb): 36 return tuple(map(lambda x: x / 255.0, rgb)) 37 38 def square(srgb): 39 return tuple(map(lambda x: pow(x, 2), srgb)) 40 41 def invert(srgb): 42 return tuple(map(lambda x: 1.0 - x, srgb)) 43 44 def combination(rgb): 45 rgb = square(scale(rgb)) 46 rgbi = invert(rgb) 47 pairs = zip(rgbi, rgb) 48 d = [] 49 for corner in corners: 50 rs, gs, bs = scale(corner) 51 d.append((pairs[0][int(rs)] * pairs[1][int(gs)] * pairs[2][int(bs)], corner)) 52 return balance(d) 53 54 def complements(rgb): 55 r, g, b = rgb 56 return rgb, restore(invert(scale(rgb))) 57 58 def balance(d): 59 d = dict([(value, f) for f, value in d]) 60 for primary, secondary in map(complements, [(255, 0, 0), (0, 255, 0), (0, 0, 255)]): 61 common = min(d[primary], d[secondary]) 62 d[primary] -= common 63 d[secondary] -= common 64 d[(0, 0, 0)] += common 65 d[(255, 255, 255)] += common 66 return [(f, value) for value, f in d.items()] 67 68 def combine(d): 69 out = [0, 0, 0] 70 for v, rgb in d: 71 out[0] += v * rgb[0] 72 out[1] += v * rgb[1] 73 out[2] += v * rgb[2] 74 return out 75 76 def pattern(rgb): 77 l = combination(rgb) 78 l.sort(reverse=True) 79 return l 80 81 def get_from_distribution(rgb, dist): 82 choose = random() 83 threshold = 0 84 for f, c in dist: 85 threshold += f 86 if choose < threshold: 87 return c 88 return c 89 90 def get_value(rgb): 91 return get_from_distribution(rgb, pattern(rgb)) 92 93 def sign(x): 94 return x >= 0 and 1 or -1 95 96 def saturate_rgb(rgb, exp): 97 return tuple([saturate_value(x, exp) for x in rgb]) 98 99 def saturate_value(x, exp): 100 return int(127.5 + sign(x - 127.5) * 127.5 * pow(abs(x - 127.5) / 127.5, exp)) 101 102 def replace(value, values): 103 if value not in values: 104 for i, v in list(enumerate(values))[::-1]: 105 if v != value: 106 values[i] = value 107 return 108 109 def test(): 110 size = 512 111 for r in (0, 63, 127, 191, 255): 112 im = PIL.Image.new("RGB", (size, size)) 113 for g in range(0, size): 114 for b in range(0, size): 115 value = get_value((r, (g * 256) / size, (b * 256 / size))) 116 im.putpixel((g, b), value) 117 im.save("rgb%d.png" % r) 118 119 def test_flat(rgb): 120 size = 64 121 im = PIL.Image.new("RGB", (size, size)) 122 for y in range(0, size): 123 for x in range(0, size): 124 im.putpixel((x, y), get_value(rgb)) 125 im.save("rgb%02d%02d%02d.png" % rgb) 126 127 def rotate_and_scale(im, width, height, rotate): 128 if rotate or x and x["Image Orientation"].values == [6L]: 129 im = im.rotate(270) 130 131 w, h = im.size 132 if w > h: 133 height = (width * h) / w 134 else: 135 width = (height * w) / h 136 137 return im.resize((width, height)) 138 139 if __name__ == "__main__": 140 if "--test" in sys.argv: 141 test() 142 sys.exit(0) 143 elif "--test-flat" in sys.argv: 144 test_flat((120, 40, 60)) 145 sys.exit(0) 146 147 width = 320 148 height = 256 149 150 input_filename, output_filename = sys.argv[1:3] 151 basename, ext = splitext(output_filename) 152 preview_filename = "".join([basename + "_preview", ext]) 153 154 rotate = "-r" in sys.argv[3:] 155 saturate = sys.argv[3:].count("-s") 156 desaturate = sys.argv[3:].count("-d") 157 preview = "-p" in sys.argv[3:] 158 square = "-2" in sys.argv[3:] and square or (lambda x: x) 159 160 x = EXIF.process_file(open(input_filename)) 161 im = PIL.Image.open(input_filename).convert("RGB") 162 im = rotate_and_scale(im, width, height, rotate) 163 164 width, height = im.size 165 166 colours = [] 167 168 for y in range(0, height): 169 c = {} 170 for x in range(0, width): 171 rgb = im.getpixel((x, y)) 172 173 # Saturate if requested. 174 175 if saturate or desaturate: 176 rgb = saturate_rgb(rgb, saturate and math.pow(0.5, saturate) or math.pow(2, desaturate)) 177 im.putpixel((x, y), rgb) 178 179 # Sum the colour probabilities. 180 181 for f, value in combination(rgb): 182 if not c.has_key(value): 183 c[value] = f 184 else: 185 c[value] += f 186 187 c = [(n/width, value) for value, n in c.items()] 188 c.sort(reverse=True) 189 colours.append(c) 190 191 if preview: 192 imp = im.copy() 193 for y in range(0, height): 194 for x in range(0, width): 195 rgb = imp.getpixel((x, y)) 196 value = get_value(rgb) 197 imp.putpixel((x, y), value) 198 imp.save(preview_filename) 199 200 for y, c in enumerate(colours): 201 most = [value for n, value in c[:4]] 202 least = [value for n, value in c[4:]] 203 204 if least: 205 switched = [] 206 for j in 1, 2: 207 i = randrange(0, 4) 208 n, value = c[i] 209 if n < 0.1: 210 switched.append(c[i]) 211 del c[i] 212 c += switched 213 most = [value for n, value in c[:4]] 214 least = [value for n, value in c[4:]] 215 216 for x in range(0, width): 217 rgb = im.getpixel((x, y)) 218 219 # Get the requested colours and choose the closest alternative for 220 # less common colours. 221 222 value = get_value(rgb) 223 if value in least: 224 value = nearest(rgb, most) 225 226 im.putpixel((x, y), value) 227 228 im.save(output_filename) 229 230 # vim: tabstop=4 expandtab shiftwidth=4