paul@0 | 1 | #!/usr/bin/env python |
paul@0 | 2 | |
paul@1 | 3 | from itertools import combinations |
paul@20 | 4 | from random import randint |
paul@12 | 5 | import EXIF |
paul@0 | 6 | import PIL.Image |
paul@20 | 7 | import math |
paul@0 | 8 | import sys |
paul@0 | 9 | |
paul@20 | 10 | corners = [ |
paul@20 | 11 | (0, 0, 0), (255, 0, 0), (0, 255, 0), (255, 255, 0), |
paul@20 | 12 | (0, 0, 255), (255, 0, 255), (0, 255, 255), (255, 255, 255) |
paul@0 | 13 | ] |
paul@0 | 14 | |
paul@20 | 15 | def distance(rgb1, rgb2): |
paul@20 | 16 | r1, g1, b1 = rgb1 |
paul@20 | 17 | r2, g2, b2 = rgb2 |
paul@20 | 18 | return math.sqrt(pow(r1 - r2, 2) + pow(g1 - g2, 2) + pow(b1 - b2, 2)) |
paul@20 | 19 | |
paul@20 | 20 | def brightness(rgb): |
paul@20 | 21 | return distance(rgb, (0, 0, 0)) |
paul@0 | 22 | |
paul@20 | 23 | def factor(start, end, rgb): |
paul@20 | 24 | r1, g1, b1 = start |
paul@20 | 25 | r2, g2, b2 = end |
paul@20 | 26 | gr, gg, gb = r2 - r1, g2 - g1, b2 - b1 |
paul@20 | 27 | r, g, b = rgb |
paul@20 | 28 | pr, pg, pb = r - r1, g - g1, b - b1 |
paul@20 | 29 | dp = pr * gr + pg * gg + pb * gb |
paul@20 | 30 | return dp / pow(distance(start, end), 2) |
paul@20 | 31 | |
paul@20 | 32 | def darklight(rgb1, rgb2): |
paul@20 | 33 | if brightness(rgb1) <= brightness(rgb2): |
paul@20 | 34 | return rgb1, rgb2 |
paul@20 | 35 | else: |
paul@20 | 36 | return rgb2, rgb1 |
paul@20 | 37 | |
paul@20 | 38 | def nearest(rgb, values=None): |
paul@20 | 39 | l = map(lambda c: (distance(rgb, c), c), values or corners) |
paul@20 | 40 | l.sort() |
paul@20 | 41 | return l |
paul@19 | 42 | |
paul@20 | 43 | def pattern(rgb): |
paul@20 | 44 | l = nearest(rgb) |
paul@20 | 45 | start, end = l[0][1], l[1][1] |
paul@20 | 46 | f = factor(start, end, rgb) |
paul@20 | 47 | #if f > 0.5: |
paul@20 | 48 | # start, end = end, start |
paul@20 | 49 | # f = 1 - f |
paul@20 | 50 | return start, end, f |
paul@20 | 51 | |
paul@20 | 52 | def choose(seq, f): |
paul@20 | 53 | last = int(seq * math.sqrt(f)) |
paul@20 | 54 | current = int((seq + 1) * math.sqrt(f)) |
paul@20 | 55 | return last != current |
paul@20 | 56 | |
paul@20 | 57 | def get_value(xy, rgb, width, height): |
paul@20 | 58 | x, y = xy |
paul@20 | 59 | rgb1, rgb2, f = pattern(rgb) |
paul@20 | 60 | if choose(x + randint(0, width), f) and choose(y + randint(0, height), f): |
paul@20 | 61 | return rgb2 |
paul@20 | 62 | else: |
paul@20 | 63 | return rgb1 |
paul@0 | 64 | |
paul@20 | 65 | def get_best(rgb, values): |
paul@20 | 66 | return nearest(rgb, values)[0][1] |
paul@19 | 67 | |
paul@20 | 68 | def test(): |
paul@20 | 69 | size = 512 |
paul@20 | 70 | for r in (0, 63, 127, 191, 255): |
paul@20 | 71 | im = PIL.Image.new("RGB", (size, size)) |
paul@20 | 72 | for g in range(0, size): |
paul@20 | 73 | for b in range(0, size): |
paul@20 | 74 | value = get_value((g, b), (r, (g * 256) / size, (b * 256 / size)), size, size) |
paul@20 | 75 | im.putpixel((g, b), value) |
paul@20 | 76 | im.save("rgb%d.png" % r) |
paul@20 | 77 | |
paul@20 | 78 | def test_flat(rgb): |
paul@20 | 79 | size = 64 |
paul@20 | 80 | im = PIL.Image.new("RGB", (size, size)) |
paul@20 | 81 | for y in range(0, size): |
paul@20 | 82 | for x in range(0, size): |
paul@20 | 83 | im.putpixel((x, y), get_value((x, y), rgb, size, size)) |
paul@20 | 84 | im.save("rgb%02d%02d%02d.png" % rgb) |
paul@20 | 85 | |
paul@20 | 86 | def rotate_and_scale(im, width, height, rotate): |
paul@19 | 87 | if rotate or x and x["Image Orientation"].values == [6L]: |
paul@11 | 88 | im = im.rotate(270) |
paul@19 | 89 | |
paul@4 | 90 | w, h = im.size |
paul@19 | 91 | if w > h: |
paul@19 | 92 | height = (width * h) / w |
paul@19 | 93 | else: |
paul@19 | 94 | width = (height * w) / h |
paul@19 | 95 | |
paul@20 | 96 | return im.resize((width, height)) |
paul@1 | 97 | |
paul@20 | 98 | if __name__ == "__main__": |
paul@20 | 99 | if "--test" in sys.argv: |
paul@20 | 100 | test() |
paul@20 | 101 | sys.exit(0) |
paul@20 | 102 | elif "--test-flat" in sys.argv: |
paul@20 | 103 | test_flat((120, 40, 60)) |
paul@20 | 104 | sys.exit(0) |
paul@0 | 105 | |
paul@20 | 106 | width = 320 |
paul@20 | 107 | height = 256 |
paul@20 | 108 | |
paul@20 | 109 | input_filename, output_filename = sys.argv[1:3] |
paul@20 | 110 | rotate = "-r" in sys.argv[3:] |
paul@1 | 111 | |
paul@20 | 112 | x = EXIF.process_file(open(input_filename)) |
paul@20 | 113 | im = PIL.Image.open(input_filename) |
paul@20 | 114 | im = rotate_and_scale(im, width, height, rotate) |
paul@1 | 115 | |
paul@20 | 116 | width, height = im.size |
paul@20 | 117 | |
paul@20 | 118 | colours = [] |
paul@10 | 119 | |
paul@20 | 120 | for y in range(0, height): |
paul@20 | 121 | c = {} |
paul@20 | 122 | for x in range(0, width): |
paul@20 | 123 | rgb = im.getpixel((x, y)) |
paul@20 | 124 | value = get_value((x, y), rgb, width, height) |
paul@20 | 125 | im.putpixel((x, y), value) |
paul@20 | 126 | if not c.has_key(value): |
paul@20 | 127 | c[value] = 1 |
paul@20 | 128 | else: |
paul@20 | 129 | c[value] += 1 |
paul@20 | 130 | c = [(n, value) for value, n in c.items()] |
paul@20 | 131 | c.sort(reverse=True) |
paul@20 | 132 | colours.append(c) |
paul@10 | 133 | |
paul@20 | 134 | for y, c in enumerate(colours): |
paul@20 | 135 | if len(c) <= 4: |
paul@20 | 136 | continue |
paul@20 | 137 | most = [value for n, value in c[:4]] |
paul@20 | 138 | least = [value for n, value in c[4:]] |
paul@1 | 139 | |
paul@20 | 140 | for x in range(0, width): |
paul@20 | 141 | rgb = im.getpixel((x, y)) |
paul@20 | 142 | if rgb in least: |
paul@20 | 143 | value = get_best(rgb, most) |
paul@20 | 144 | im.putpixel((x, y), value) |
paul@1 | 145 | |
paul@0 | 146 | im.save(output_filename) |
paul@0 | 147 | |
paul@0 | 148 | # vim: tabstop=4 expandtab shiftwidth=4 |