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