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 within(v, lower, upper): 17 return min(max(v, lower), upper) 18 19 def clip(v): 20 return int(within(v, 0, 255)) 21 22 def distance(rgb1, rgb2): 23 r1, g1, b1 = rgb1 24 r2, g2, b2 = rgb2 25 return math.sqrt(pow(r1 - r2, 2) + pow(g1 - g2, 2) + pow(b1 - b2, 2)) 26 27 def nearest(rgb, values): 28 l = [(distance(rgb, value), value) for value in values] 29 l.sort() 30 return l[0][1] 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 cache = {} 45 46 def combination(rgb): 47 if not cache.has_key(rgb): 48 rgb = square(scale(rgb)) 49 rgbi = invert(rgb) 50 pairs = zip(rgbi, rgb) 51 d = [] 52 for corner in corners: 53 rs, gs, bs = scale(corner) 54 d.append((pairs[0][int(rs)] * pairs[1][int(gs)] * pairs[2][int(bs)], corner)) 55 cache[rgb] = balance(d) 56 return cache[rgb] 57 58 def complements(rgb): 59 r, g, b = rgb 60 return rgb, restore(invert(scale(rgb))) 61 62 def balance(d): 63 d = dict([(value, f) for f, value in d]) 64 for primary, secondary in map(complements, [(255, 0, 0), (0, 255, 0), (0, 0, 255)]): 65 common = min(d[primary], d[secondary]) 66 d[primary] -= common 67 d[secondary] -= common 68 d[(0, 0, 0)] += common 69 d[(255, 255, 255)] += common 70 return [(f, value) for value, f in d.items()] 71 72 def compensate(d, chosen): 73 dd = dict([(value, f) for f, value in d]) 74 for f, value in d: 75 if value in chosen: 76 _value, complement = complements(value) 77 if complement not in chosen: 78 f = max(0, f - dd[complement]) 79 dd[value] = f 80 return [(f, value) for value, f in dd.items() if value in chosen] 81 82 def combine(d): 83 out = [0, 0, 0] 84 for v, rgb in d: 85 out[0] += v * rgb[0] 86 out[1] += v * rgb[1] 87 out[2] += v * rgb[2] 88 return out 89 90 def pattern(rgb, chosen=None): 91 l = combination(rgb) 92 if chosen: 93 l = compensate(l, chosen) 94 l.sort(reverse=True) 95 return l 96 97 def get_value(rgb, chosen=None): 98 l = pattern(rgb, chosen) 99 limit = sum([f for f, c in l]) 100 choose = random() * limit 101 threshold = 0 102 for f, c in l: 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 amplify_rgb(rgb, exp): 118 return tuple([amplify_value(x, exp) for x in rgb]) 119 120 def amplify_value(x, exp): 121 return int(pow(x / 255.0, exp) * 255.0) 122 123 def get_colours(im, y): 124 width, height = im.size 125 c = {} 126 for x in range(0, width): 127 rgb = im.getpixel((x, y)) 128 129 # Sum the colour probabilities. 130 131 for f, value in combination(rgb): 132 if not c.has_key(value): 133 c[value] = f 134 else: 135 c[value] += f 136 137 c = [(n/width, value) for value, n in c.items()] 138 c.sort(reverse=True) 139 return c 140 141 def test(): 142 size = 512 143 for r in (0, 63, 127, 191, 255): 144 im = PIL.Image.new("RGB", (size, size)) 145 for g in range(0, size): 146 for b in range(0, size): 147 value = get_value((r, (g * 256) / size, (b * 256 / size))) 148 im.putpixel((g, b), value) 149 im.save("rgb%d.png" % r) 150 151 def test_flat(rgb): 152 size = 64 153 im = PIL.Image.new("RGB", (size, size)) 154 for y in range(0, size): 155 for x in range(0, size): 156 im.putpixel((x, y), get_value(rgb)) 157 im.save("rgb%02d%02d%02d.png" % rgb) 158 159 def rotate_and_scale(exif, im, width, height, rotate): 160 if rotate or exif and exif["Image Orientation"].values == [6L]: 161 im = im.rotate(270) 162 163 w, h = im.size 164 if w > h: 165 height = (width * h) / w 166 else: 167 width = (height * w) / h 168 169 return im.resize((width, height)) 170 171 def count_colours(im, colours): 172 width, height = im.size 173 for y in range(0, height): 174 l = set() 175 for x in range(0, width): 176 l.add(im.getpixel((x, y))) 177 if len(l) > colours: 178 return (y, l) 179 return None 180 181 if __name__ == "__main__": 182 if "--test" in sys.argv: 183 test() 184 sys.exit(0) 185 elif "--test-flat" in sys.argv: 186 test_flat((120, 40, 60)) 187 sys.exit(0) 188 189 width = 320 190 height = 256 191 192 input_filename, output_filename = sys.argv[1:3] 193 basename, ext = splitext(output_filename) 194 preview_filename = "".join([basename + "_preview", ext]) 195 196 options = sys.argv[3:] 197 198 rotate = "-r" in options 199 saturate = options.count("-s") 200 desaturate = options.count("-d") 201 darken = options.count("-D") 202 brighten = options.count("-B") 203 square = "-2" in options and square or (lambda x: x) 204 preview = "-p" in options 205 half_resolution_preview = "-h" in options 206 verify = "-v" in options 207 no_normal_output = "-n" in options 208 make_image = not no_normal_output 209 210 if make_image or preview: 211 exif = EXIF.process_file(open(input_filename)) 212 im = PIL.Image.open(input_filename).convert("RGB") 213 im = rotate_and_scale(exif, im, width, height, rotate) 214 215 width, height = im.size 216 217 if saturate or desaturate or darken or brighten: 218 for y in range(0, height): 219 for x in range(0, width): 220 rgb = im.getpixel((x, y)) 221 if saturate or desaturate: 222 rgb = saturate_rgb(rgb, saturate and math.pow(0.5, saturate) or math.pow(2, desaturate)) 223 if darken or brighten: 224 rgb = amplify_rgb(rgb, brighten and math.pow(0.5, brighten) or math.pow(2, darken)) 225 im.putpixel((x, y), rgb) 226 227 if preview: 228 imp = im.copy() 229 step = half_resolution_preview and 2 or 1 230 for y in range(0, height): 231 for x in range(0, width, step): 232 rgb = imp.getpixel((x, y)) 233 value = get_value(rgb) 234 imp.putpixel((x, y), value) 235 if half_resolution_preview: 236 imp.putpixel((x+1, y), value) 237 238 imp.save(preview_filename) 239 240 if make_image: 241 for y in range(0, height): 242 c = get_colours(im, y) 243 most = [value for n, value in c[:4]] 244 least = [value for n, value in c[4:]] 245 246 for x in range(0, width): 247 rgb = im.getpixel((x, y)) 248 value = get_value(rgb, most) 249 im.putpixel((x, y), value) 250 251 if y < height - 1: 252 253 # Spread the error across adjacent pixels. 254 255 error = tuple(map(lambda i: (i[0] - i[1]) / 2.0, zip(rgb, value))) 256 for xn in range(x, min(x + 2, width)): 257 rgbn = im.getpixel((xn, y+1)) 258 rgbn = tuple(map(lambda i: clip(i[0] + i[1]), zip(rgbn, error))) 259 im.putpixel((xn, y+1), rgbn) 260 261 im.save(output_filename) 262 263 if verify: 264 if no_normal_output: 265 im = PIL.Image.open(output_filename).convert("RGB") 266 267 result = count_colours(im, 4) 268 if result is not None: 269 y, colours = result 270 print "Row %d has the following colours: %s" % (y, "; ".join([repr(c) for c in colours])) 271 272 # vim: tabstop=4 expandtab shiftwidth=4