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 test(): 56 size = 512 57 for r in (0, 63, 127, 191, 255): 58 im = PIL.Image.new("RGB", (size, size)) 59 for g in range(0, size): 60 for b in range(0, size): 61 value = get_value((r, (g * 256) / size, (b * 256 / size))) 62 im.putpixel((g, b), value) 63 im.save("rgb%d.png" % r) 64 65 def test_flat(rgb): 66 size = 64 67 im = PIL.Image.new("RGB", (size, size)) 68 for y in range(0, size): 69 for x in range(0, size): 70 im.putpixel((x, y), get_value(rgb)) 71 im.save("rgb%02d%02d%02d.png" % rgb) 72 73 def rotate_and_scale(im, width, height, rotate): 74 if rotate or x and x["Image Orientation"].values == [6L]: 75 im = im.rotate(270) 76 77 w, h = im.size 78 if w > h: 79 height = (width * h) / w 80 else: 81 width = (height * w) / h 82 83 return im.resize((width, height)) 84 85 if __name__ == "__main__": 86 if "--test" in sys.argv: 87 test() 88 sys.exit(0) 89 elif "--test-flat" in sys.argv: 90 test_flat((120, 40, 60)) 91 sys.exit(0) 92 93 width = 320 94 height = 256 95 96 input_filename, output_filename = sys.argv[1:3] 97 basename, ext = splitext(output_filename) 98 99 rotate = "-r" in sys.argv[3:] 100 saturate = sys.argv[3:].count("-s") 101 desaturate = sys.argv[3:].count("-d") 102 103 x = EXIF.process_file(open(input_filename)) 104 im = PIL.Image.open(input_filename).convert("RGB") 105 im = rotate_and_scale(im, width, height, rotate) 106 107 width, height = im.size 108 109 colours = [] 110 111 for y in range(0, height): 112 c = {} 113 for x in range(0, width): 114 rgb = im.getpixel((x, y)) 115 116 # Saturate if requested. 117 118 if saturate or desaturate: 119 rgb = saturate_rgb(rgb, saturate and math.pow(0.5, saturate) or math.pow(2, desaturate)) 120 im.putpixel((x, y), rgb) 121 122 # Sum the colour probabilities. 123 124 for f, value in distribution(rgb): 125 if not c.has_key(value): 126 c[value] = f 127 else: 128 c[value] += f 129 130 c = [(n, value) for value, n in c.items()] 131 c.sort(reverse=True) 132 colours.append(c) 133 134 for y, c in enumerate(colours): 135 most = [value for n, value in c[:4]] 136 least = [value for n, value in c[4:]] 137 138 for x in range(0, width): 139 rgb = im.getpixel((x, y)) 140 141 # Get the requested colours and choose the closest alternative for 142 # less common colours. 143 144 value = get_value(rgb) 145 if value in least: 146 rgb = im.getpixel((x, y)) 147 value = pattern(value, most)[0][1] 148 149 im.putpixel((x, y), value) 150 151 im.save(output_filename) 152 153 # vim: tabstop=4 expandtab shiftwidth=4