PaletteOptimiser

optimiser.py

24:c768f24e5233
2015-10-01 Paul Boddie Choose the nearest alternative colours to the original colour.
     1 #!/usr/bin/env python     2      3 from itertools import combinations     4 from random import randint     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 * math.sqrt(f))    54     current = int((seq + 1) * math.sqrt(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     rotate = "-r" in sys.argv[3:]   120     saturate = sys.argv[3:].count("-s")   121     desaturate = sys.argv[3:].count("-d")   122    123     x = EXIF.process_file(open(input_filename))   124     im = PIL.Image.open(input_filename)   125     im = rotate_and_scale(im, width, height, rotate)   126    127     width, height = im.size   128    129     colours = []   130    131     for y in range(0, height):   132         c = {}   133         for x in range(0, width):   134             rgb = im.getpixel((x, y))   135    136             # Saturate if requested.   137    138             if saturate or desaturate:   139                 rgb = saturate_rgb(rgb, saturate and math.pow(0.5, saturate) or math.pow(2, desaturate))   140                 im.putpixel((x, y), rgb)   141    142             # Count the number of requested colours.   143    144             value = get_value((x, y), rgb, width, height)   145             if not c.has_key(value):   146                 c[value] = 1   147             else:   148                 c[value] += 1   149    150         c = [(n, value) for value, n in c.items()]   151         c.sort(reverse=True)   152         colours.append(c)   153    154     for y, c in enumerate(colours):   155         most = [value for n, value in c[:4]]   156         least = [value for n, value in c[4:]]   157    158         for x in range(0, width):   159             rgb = im.getpixel((x, y))   160    161             # Get the requested colours and choose the closest alternative for   162             # less common colours.   163    164             value = get_value((x, y), rgb, width, height)   165             if value in least:   166                 value = get_value((x, y), rgb, width, height, most)   167             im.putpixel((x, y), value)   168    169     im.save(output_filename)   170    171 # vim: tabstop=4 expandtab shiftwidth=4