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 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 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 compensate(d, chosen): 69 dd = dict([(value, f) for f, value in d]) 70 for f, value in d: 71 if value not in chosen: 72 _value, complement = complements(value) 73 if complement in chosen: 74 f = max(0, f - dd[complement]) 75 dd[value] = f 76 return [(f, value) for value, f in dd.items() if value in chosen] 77 78 def combine(d): 79 out = [0, 0, 0] 80 for v, rgb in d: 81 out[0] += v * rgb[0] 82 out[1] += v * rgb[1] 83 out[2] += v * rgb[2] 84 return out 85 86 def pattern(rgb, chosen=None): 87 l = combination(rgb) 88 if chosen: 89 l = compensate(l, chosen) 90 l.sort(reverse=True) 91 return l 92 93 def get_value(rgb, chosen=None): 94 l = pattern(rgb, chosen) 95 limit = sum([f for f, c in l]) 96 choose = random() * limit 97 threshold = 0 98 for f, c in l: 99 threshold += f 100 if choose < threshold: 101 return c 102 return c 103 104 def sign(x): 105 return x >= 0 and 1 or -1 106 107 def saturate_rgb(rgb, exp): 108 return tuple([saturate_value(x, exp) for x in rgb]) 109 110 def saturate_value(x, exp): 111 return int(127.5 + sign(x - 127.5) * 127.5 * pow(abs(x - 127.5) / 127.5, exp)) 112 113 def amplify_rgb(rgb, exp): 114 return tuple([amplify_value(x, exp) for x in rgb]) 115 116 def amplify_value(x, exp): 117 return int(pow(x / 255.0, exp) * 255.0) 118 119 def get_colours(im, y): 120 width, height = im.size 121 c = {} 122 for x in range(0, width): 123 rgb = im.getpixel((x, y)) 124 125 # Sum the colour probabilities. 126 127 for f, value in combination(rgb): 128 if not c.has_key(value): 129 c[value] = f 130 else: 131 c[value] += f 132 133 c = [(n/width, value) for value, n in c.items()] 134 c.sort(reverse=True) 135 return c 136 137 def test(): 138 size = 512 139 for r in (0, 63, 127, 191, 255): 140 im = PIL.Image.new("RGB", (size, size)) 141 for g in range(0, size): 142 for b in range(0, size): 143 value = get_value((r, (g * 256) / size, (b * 256 / size))) 144 im.putpixel((g, b), value) 145 im.save("rgb%d.png" % r) 146 147 def test_flat(rgb): 148 size = 64 149 im = PIL.Image.new("RGB", (size, size)) 150 for y in range(0, size): 151 for x in range(0, size): 152 im.putpixel((x, y), get_value(rgb)) 153 im.save("rgb%02d%02d%02d.png" % rgb) 154 155 def rotate_and_scale(exif, im, width, height, rotate): 156 if rotate or exif and exif["Image Orientation"].values == [6L]: 157 im = im.rotate(270) 158 159 w, h = im.size 160 if w > h: 161 height = (width * h) / w 162 else: 163 width = (height * w) / h 164 165 return im.resize((width, height)) 166 167 def count_colours(im, colours): 168 width, height = im.size 169 for y in range(0, height): 170 l = set() 171 for x in range(0, width): 172 l.add(im.getpixel((x, y))) 173 if len(l) > colours: 174 return (y, l) 175 return None 176 177 if __name__ == "__main__": 178 if "--test" in sys.argv: 179 test() 180 sys.exit(0) 181 elif "--test-flat" in sys.argv: 182 test_flat((120, 40, 60)) 183 sys.exit(0) 184 185 width = 320 186 height = 256 187 188 input_filename, output_filename = sys.argv[1:3] 189 basename, ext = splitext(output_filename) 190 preview_filename = "".join([basename + "_preview", ext]) 191 192 options = sys.argv[3:] 193 194 rotate = "-r" in options 195 saturate = options.count("-s") 196 desaturate = options.count("-d") 197 darken = options.count("-D") 198 brighten = options.count("-B") 199 square = "-2" in options and square or (lambda x: x) 200 preview = "-p" in options 201 half_resolution_preview = "-h" in options 202 verify = "-v" in options 203 no_normal_output = "-n" in options 204 make_image = not no_normal_output 205 206 if make_image or preview: 207 exif = EXIF.process_file(open(input_filename)) 208 im = PIL.Image.open(input_filename).convert("RGB") 209 im = rotate_and_scale(exif, im, width, height, rotate) 210 211 width, height = im.size 212 213 if saturate or desaturate or darken or brighten: 214 for y in range(0, height): 215 for x in range(0, width): 216 if saturate or desaturate: 217 rgb = im.getpixel((x, y)) 218 rgb = saturate_rgb(rgb, saturate and math.pow(0.5, saturate) or math.pow(2, desaturate)) 219 im.putpixel((x, y), rgb) 220 221 if darken or brighten: 222 rgb = im.getpixel((x, y)) 223 rgb = amplify_rgb(rgb, brighten and math.pow(0.5, brighten) or math.pow(2, darken)) 224 im.putpixel((x, y), rgb) 225 226 if preview: 227 imp = im.copy() 228 step = half_resolution_preview and 2 or 1 229 for y in range(0, height): 230 for x in range(0, width, step): 231 rgb = imp.getpixel((x, y)) 232 value = get_value(rgb) 233 imp.putpixel((x, y), value) 234 if half_resolution_preview: 235 imp.putpixel((x+1, y), value) 236 237 imp.save(preview_filename) 238 239 if make_image: 240 for y in range(0, height): 241 c = get_colours(im, y) 242 most = [value for n, value in c[:4]] 243 least = [value for n, value in c[4:]] 244 245 for x in range(0, width): 246 rgb = im.getpixel((x, y)) 247 value = get_value(rgb, most) 248 im.putpixel((x, y), value) 249 250 if y < height - 1: 251 rgbn = im.getpixel((x, y+1)) 252 rgbn = tuple(map(lambda i: clip(i[0] + i[1] - i[2]), zip(rgbn, rgb, value))) 253 im.putpixel((x, y+1), rgbn) 254 255 im.save(output_filename) 256 257 if verify: 258 if no_normal_output: 259 im = PIL.Image.open(output_filename).convert("RGB") 260 261 result = count_colours(im, 4) 262 if result is not None: 263 y, colours = result 264 print "Row %d has the following colours: %s" % (y, "; ".join([repr(c) for c in colours])) 265 266 # vim: tabstop=4 expandtab shiftwidth=4