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