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 distribution(rgb, values=None): 22 l = [] 23 total = 0 24 for c in values or corners: 25 d = distance(rgb, c) 26 if d == 0: 27 return [(1, c)] 28 l.append((pow(d, -3), c)) 29 total += pow(d, -3) 30 return [(d / total, c) for d, c in l] 31 32 def pattern(rgb, values=None): 33 l = distribution(rgb, values or corners) 34 l.sort(reverse=True) 35 return l 36 37 def get_value(rgb, values=None): 38 choose = random() 39 threshold = 0 40 for f, c in pattern(rgb, values): 41 threshold += f 42 if choose < threshold: 43 return c 44 return c 45 46 def sign(x): 47 return x >= 0 and 1 or -1 48 49 def saturate_rgb(rgb, exp): 50 return tuple([saturate_value(x, exp) for x in rgb]) 51 52 def saturate_value(x, exp): 53 return int(127.5 + sign(x - 127.5) * 127.5 * pow(abs(x - 127.5) / 127.5, exp)) 54 55 def replace(value, values): 56 if value not in values: 57 for i, v in list(enumerate(values))[::-1]: 58 if v != value: 59 values[i] = value 60 return 61 62 def test(): 63 size = 512 64 for r in (0, 63, 127, 191, 255): 65 im = PIL.Image.new("RGB", (size, size)) 66 for g in range(0, size): 67 for b in range(0, size): 68 value = get_value((r, (g * 256) / size, (b * 256 / size))) 69 im.putpixel((g, b), value) 70 im.save("rgb%d.png" % r) 71 72 def test_flat(rgb): 73 size = 64 74 im = PIL.Image.new("RGB", (size, size)) 75 for y in range(0, size): 76 for x in range(0, size): 77 im.putpixel((x, y), get_value(rgb)) 78 im.save("rgb%02d%02d%02d.png" % rgb) 79 80 def rotate_and_scale(im, width, height, rotate): 81 if rotate or x and x["Image Orientation"].values == [6L]: 82 im = im.rotate(270) 83 84 w, h = im.size 85 if w > h: 86 height = (width * h) / w 87 else: 88 width = (height * w) / h 89 90 return im.resize((width, height)) 91 92 if __name__ == "__main__": 93 if "--test" in sys.argv: 94 test() 95 sys.exit(0) 96 elif "--test-flat" in sys.argv: 97 test_flat((120, 40, 60)) 98 sys.exit(0) 99 100 width = 320 101 height = 256 102 103 input_filename, output_filename = sys.argv[1:3] 104 basename, ext = splitext(output_filename) 105 106 rotate = "-r" in sys.argv[3:] 107 saturate = sys.argv[3:].count("-s") 108 desaturate = sys.argv[3:].count("-d") 109 110 x = EXIF.process_file(open(input_filename)) 111 im = PIL.Image.open(input_filename).convert("RGB") 112 im = rotate_and_scale(im, width, height, rotate) 113 114 width, height = im.size 115 116 colours = [] 117 118 for y in range(0, height): 119 c = {} 120 for x in range(0, width): 121 rgb = im.getpixel((x, y)) 122 123 # Saturate if requested. 124 125 if saturate or desaturate: 126 rgb = saturate_rgb(rgb, saturate and math.pow(0.5, saturate) or math.pow(2, desaturate)) 127 im.putpixel((x, y), rgb) 128 129 # Sum the colour probabilities. 130 131 for f, value in distribution(rgb): 132 if not c.has_key(value): 133 c[value] = f 134 else: 135 c[value] += f 136 137 c = [(n, value) for value, n in c.items()] 138 c.sort(reverse=True) 139 colours.append(c) 140 141 for y, c in enumerate(colours): 142 most = [value for n, value in c[:4]] 143 least = [value for n, value in c[4:]] 144 145 if least: 146 if (0, 0, 0) in least[:2]: 147 replace((0, 0, 0), most) 148 if (255, 255, 255) in least[:2]: 149 replace((255, 255, 255), most) 150 151 for x in range(0, width): 152 rgb = im.getpixel((x, y)) 153 154 # Get the requested colours and choose the closest alternative for 155 # less common colours. 156 157 value = get_value(rgb) 158 if value in least: 159 rgb = im.getpixel((x, y)) 160 value = get_value(value, most) 161 162 im.putpixel((x, y), value) 163 164 im.save(output_filename) 165 166 # vim: tabstop=4 expandtab shiftwidth=4