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