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