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