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