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