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