paul@0 | 1 | #!/usr/bin/env python |
paul@0 | 2 | |
paul@38 | 3 | from random import random, randrange |
paul@32 | 4 | from os.path import splitext |
paul@12 | 5 | import EXIF |
paul@0 | 6 | import PIL.Image |
paul@32 | 7 | import itertools |
paul@20 | 8 | import math |
paul@0 | 9 | import sys |
paul@0 | 10 | |
paul@20 | 11 | corners = [ |
paul@20 | 12 | (0, 0, 0), (255, 0, 0), (0, 255, 0), (255, 255, 0), |
paul@20 | 13 | (0, 0, 255), (255, 0, 255), (0, 255, 255), (255, 255, 255) |
paul@0 | 14 | ] |
paul@0 | 15 | |
paul@49 | 16 | def within(v, lower, upper): |
paul@49 | 17 | return min(max(v, lower), upper) |
paul@49 | 18 | |
paul@49 | 19 | def clip(v): |
paul@49 | 20 | return within(v, 0, 255) |
paul@49 | 21 | |
paul@20 | 22 | def distance(rgb1, rgb2): |
paul@20 | 23 | r1, g1, b1 = rgb1 |
paul@20 | 24 | r2, g2, b2 = rgb2 |
paul@20 | 25 | return math.sqrt(pow(r1 - r2, 2) + pow(g1 - g2, 2) + pow(b1 - b2, 2)) |
paul@20 | 26 | |
paul@34 | 27 | def nearest(rgb, values): |
paul@34 | 28 | l = [(distance(rgb, value), value) for value in values] |
paul@34 | 29 | l.sort() |
paul@34 | 30 | return l[0][1] |
paul@34 | 31 | |
paul@35 | 32 | def restore(srgb): |
paul@35 | 33 | return tuple(map(lambda x: int(x * 255.0), srgb)) |
paul@35 | 34 | |
paul@34 | 35 | def scale(rgb): |
paul@34 | 36 | return tuple(map(lambda x: x / 255.0, rgb)) |
paul@34 | 37 | |
paul@39 | 38 | def square(srgb): |
paul@39 | 39 | return tuple(map(lambda x: pow(x, 2), srgb)) |
paul@39 | 40 | |
paul@34 | 41 | def invert(srgb): |
paul@34 | 42 | return tuple(map(lambda x: 1.0 - x, srgb)) |
paul@34 | 43 | |
paul@51 | 44 | cache = {} |
paul@51 | 45 | |
paul@34 | 46 | def combination(rgb): |
paul@51 | 47 | if not cache.has_key(rgb): |
paul@51 | 48 | rgb = square(scale(rgb)) |
paul@51 | 49 | rgbi = invert(rgb) |
paul@51 | 50 | pairs = zip(rgbi, rgb) |
paul@51 | 51 | d = [] |
paul@51 | 52 | for corner in corners: |
paul@51 | 53 | rs, gs, bs = scale(corner) |
paul@51 | 54 | d.append((pairs[0][int(rs)] * pairs[1][int(gs)] * pairs[2][int(bs)], corner)) |
paul@51 | 55 | cache[rgb] = balance(d) |
paul@51 | 56 | return cache[rgb] |
paul@35 | 57 | |
paul@35 | 58 | def complements(rgb): |
paul@35 | 59 | r, g, b = rgb |
paul@35 | 60 | return rgb, restore(invert(scale(rgb))) |
paul@35 | 61 | |
paul@35 | 62 | def balance(d): |
paul@35 | 63 | d = dict([(value, f) for f, value in d]) |
paul@35 | 64 | for primary, secondary in map(complements, [(255, 0, 0), (0, 255, 0), (0, 0, 255)]): |
paul@35 | 65 | common = min(d[primary], d[secondary]) |
paul@35 | 66 | d[primary] -= common |
paul@35 | 67 | d[secondary] -= common |
paul@35 | 68 | d[(0, 0, 0)] += common |
paul@35 | 69 | d[(255, 255, 255)] += common |
paul@35 | 70 | return [(f, value) for value, f in d.items()] |
paul@20 | 71 | |
paul@48 | 72 | def compensate(d, chosen): |
paul@48 | 73 | dd = dict([(value, f) for f, value in d]) |
paul@48 | 74 | for f, value in d: |
paul@52 | 75 | if value in chosen: |
paul@48 | 76 | _value, complement = complements(value) |
paul@52 | 77 | if complement not in chosen: |
paul@48 | 78 | f = max(0, f - dd[complement]) |
paul@48 | 79 | dd[value] = f |
paul@48 | 80 | return [(f, value) for value, f in dd.items() if value in chosen] |
paul@48 | 81 | |
paul@34 | 82 | def combine(d): |
paul@34 | 83 | out = [0, 0, 0] |
paul@34 | 84 | for v, rgb in d: |
paul@34 | 85 | out[0] += v * rgb[0] |
paul@34 | 86 | out[1] += v * rgb[1] |
paul@34 | 87 | out[2] += v * rgb[2] |
paul@34 | 88 | return out |
paul@34 | 89 | |
paul@48 | 90 | def pattern(rgb, chosen=None): |
paul@34 | 91 | l = combination(rgb) |
paul@48 | 92 | if chosen: |
paul@48 | 93 | l = compensate(l, chosen) |
paul@32 | 94 | l.sort(reverse=True) |
paul@20 | 95 | return l |
paul@19 | 96 | |
paul@48 | 97 | def get_value(rgb, chosen=None): |
paul@48 | 98 | l = pattern(rgb, chosen) |
paul@48 | 99 | limit = sum([f for f, c in l]) |
paul@48 | 100 | choose = random() * limit |
paul@32 | 101 | threshold = 0 |
paul@48 | 102 | for f, c in l: |
paul@32 | 103 | threshold += f |
paul@32 | 104 | if choose < threshold: |
paul@32 | 105 | return c |
paul@32 | 106 | return c |
paul@0 | 107 | |
paul@21 | 108 | def sign(x): |
paul@21 | 109 | return x >= 0 and 1 or -1 |
paul@21 | 110 | |
paul@22 | 111 | def saturate_rgb(rgb, exp): |
paul@22 | 112 | return tuple([saturate_value(x, exp) for x in rgb]) |
paul@21 | 113 | |
paul@22 | 114 | def saturate_value(x, exp): |
paul@24 | 115 | return int(127.5 + sign(x - 127.5) * 127.5 * pow(abs(x - 127.5) / 127.5, exp)) |
paul@21 | 116 | |
paul@49 | 117 | def amplify_rgb(rgb, exp): |
paul@49 | 118 | return tuple([amplify_value(x, exp) for x in rgb]) |
paul@49 | 119 | |
paul@49 | 120 | def amplify_value(x, exp): |
paul@49 | 121 | return int(pow(x / 255.0, exp) * 255.0) |
paul@49 | 122 | |
paul@50 | 123 | def get_colours(im, y): |
paul@50 | 124 | width, height = im.size |
paul@49 | 125 | c = {} |
paul@49 | 126 | for x in range(0, width): |
paul@49 | 127 | rgb = im.getpixel((x, y)) |
paul@49 | 128 | |
paul@49 | 129 | # Sum the colour probabilities. |
paul@49 | 130 | |
paul@49 | 131 | for f, value in combination(rgb): |
paul@49 | 132 | if not c.has_key(value): |
paul@49 | 133 | c[value] = f |
paul@49 | 134 | else: |
paul@49 | 135 | c[value] += f |
paul@49 | 136 | |
paul@49 | 137 | c = [(n/width, value) for value, n in c.items()] |
paul@49 | 138 | c.sort(reverse=True) |
paul@49 | 139 | return c |
paul@33 | 140 | |
paul@20 | 141 | def test(): |
paul@20 | 142 | size = 512 |
paul@20 | 143 | for r in (0, 63, 127, 191, 255): |
paul@20 | 144 | im = PIL.Image.new("RGB", (size, size)) |
paul@20 | 145 | for g in range(0, size): |
paul@20 | 146 | for b in range(0, size): |
paul@28 | 147 | value = get_value((r, (g * 256) / size, (b * 256 / size))) |
paul@20 | 148 | im.putpixel((g, b), value) |
paul@20 | 149 | im.save("rgb%d.png" % r) |
paul@20 | 150 | |
paul@20 | 151 | def test_flat(rgb): |
paul@20 | 152 | size = 64 |
paul@20 | 153 | im = PIL.Image.new("RGB", (size, size)) |
paul@20 | 154 | for y in range(0, size): |
paul@20 | 155 | for x in range(0, size): |
paul@28 | 156 | im.putpixel((x, y), get_value(rgb)) |
paul@20 | 157 | im.save("rgb%02d%02d%02d.png" % rgb) |
paul@20 | 158 | |
paul@50 | 159 | def rotate_and_scale(exif, im, width, height, rotate): |
paul@50 | 160 | if rotate or exif and exif["Image Orientation"].values == [6L]: |
paul@11 | 161 | im = im.rotate(270) |
paul@19 | 162 | |
paul@4 | 163 | w, h = im.size |
paul@19 | 164 | if w > h: |
paul@19 | 165 | height = (width * h) / w |
paul@19 | 166 | else: |
paul@19 | 167 | width = (height * w) / h |
paul@19 | 168 | |
paul@20 | 169 | return im.resize((width, height)) |
paul@1 | 170 | |
paul@50 | 171 | def count_colours(im, colours): |
paul@50 | 172 | width, height = im.size |
paul@50 | 173 | for y in range(0, height): |
paul@50 | 174 | l = set() |
paul@50 | 175 | for x in range(0, width): |
paul@50 | 176 | l.add(im.getpixel((x, y))) |
paul@50 | 177 | if len(l) > colours: |
paul@50 | 178 | return (y, l) |
paul@50 | 179 | return None |
paul@50 | 180 | |
paul@20 | 181 | if __name__ == "__main__": |
paul@20 | 182 | if "--test" in sys.argv: |
paul@20 | 183 | test() |
paul@20 | 184 | sys.exit(0) |
paul@20 | 185 | elif "--test-flat" in sys.argv: |
paul@20 | 186 | test_flat((120, 40, 60)) |
paul@20 | 187 | sys.exit(0) |
paul@0 | 188 | |
paul@20 | 189 | width = 320 |
paul@20 | 190 | height = 256 |
paul@20 | 191 | |
paul@20 | 192 | input_filename, output_filename = sys.argv[1:3] |
paul@26 | 193 | basename, ext = splitext(output_filename) |
paul@35 | 194 | preview_filename = "".join([basename + "_preview", ext]) |
paul@26 | 195 | |
paul@50 | 196 | options = sys.argv[3:] |
paul@1 | 197 | |
paul@50 | 198 | rotate = "-r" in options |
paul@50 | 199 | saturate = options.count("-s") |
paul@50 | 200 | desaturate = options.count("-d") |
paul@50 | 201 | darken = options.count("-D") |
paul@50 | 202 | brighten = options.count("-B") |
paul@50 | 203 | square = "-2" in options and square or (lambda x: x) |
paul@50 | 204 | preview = "-p" in options |
paul@50 | 205 | half_resolution_preview = "-h" in options |
paul@50 | 206 | verify = "-v" in options |
paul@50 | 207 | no_normal_output = "-n" in options |
paul@50 | 208 | make_image = not no_normal_output |
paul@1 | 209 | |
paul@50 | 210 | if make_image or preview: |
paul@50 | 211 | exif = EXIF.process_file(open(input_filename)) |
paul@50 | 212 | im = PIL.Image.open(input_filename).convert("RGB") |
paul@50 | 213 | im = rotate_and_scale(exif, im, width, height, rotate) |
paul@50 | 214 | |
paul@50 | 215 | width, height = im.size |
paul@20 | 216 | |
paul@50 | 217 | if saturate or desaturate or darken or brighten: |
paul@50 | 218 | for y in range(0, height): |
paul@50 | 219 | for x in range(0, width): |
paul@51 | 220 | rgb = im.getpixel((x, y)) |
paul@50 | 221 | if saturate or desaturate: |
paul@50 | 222 | rgb = saturate_rgb(rgb, saturate and math.pow(0.5, saturate) or math.pow(2, desaturate)) |
paul@50 | 223 | if darken or brighten: |
paul@50 | 224 | rgb = amplify_rgb(rgb, brighten and math.pow(0.5, brighten) or math.pow(2, darken)) |
paul@51 | 225 | im.putpixel((x, y), rgb) |
paul@10 | 226 | |
paul@35 | 227 | if preview: |
paul@35 | 228 | imp = im.copy() |
paul@50 | 229 | step = half_resolution_preview and 2 or 1 |
paul@35 | 230 | for y in range(0, height): |
paul@50 | 231 | for x in range(0, width, step): |
paul@35 | 232 | rgb = imp.getpixel((x, y)) |
paul@35 | 233 | value = get_value(rgb) |
paul@35 | 234 | imp.putpixel((x, y), value) |
paul@50 | 235 | if half_resolution_preview: |
paul@50 | 236 | imp.putpixel((x+1, y), value) |
paul@50 | 237 | |
paul@35 | 238 | imp.save(preview_filename) |
paul@35 | 239 | |
paul@50 | 240 | if make_image: |
paul@50 | 241 | for y in range(0, height): |
paul@50 | 242 | c = get_colours(im, y) |
paul@50 | 243 | most = [value for n, value in c[:4]] |
paul@50 | 244 | least = [value for n, value in c[4:]] |
paul@50 | 245 | |
paul@50 | 246 | for x in range(0, width): |
paul@50 | 247 | rgb = im.getpixel((x, y)) |
paul@50 | 248 | value = get_value(rgb, most) |
paul@50 | 249 | im.putpixel((x, y), value) |
paul@1 | 250 | |
paul@50 | 251 | if y < height - 1: |
paul@50 | 252 | rgbn = im.getpixel((x, y+1)) |
paul@50 | 253 | rgbn = tuple(map(lambda i: clip(i[0] + i[1] - i[2]), zip(rgbn, rgb, value))) |
paul@50 | 254 | im.putpixel((x, y+1), rgbn) |
paul@50 | 255 | |
paul@50 | 256 | im.save(output_filename) |
paul@1 | 257 | |
paul@50 | 258 | if verify: |
paul@50 | 259 | if no_normal_output: |
paul@50 | 260 | im = PIL.Image.open(output_filename).convert("RGB") |
paul@49 | 261 | |
paul@50 | 262 | result = count_colours(im, 4) |
paul@50 | 263 | if result is not None: |
paul@50 | 264 | y, colours = result |
paul@50 | 265 | print "Row %d has the following colours: %s" % (y, "; ".join([repr(c) for c in colours])) |
paul@0 | 266 | |
paul@0 | 267 | # vim: tabstop=4 expandtab shiftwidth=4 |