PaletteOptimiser

optimiser.py

47:4b9192e71fb8
2015-10-04 Paul Boddie Introduced a distance-based approach to obtain suitable alternative colours.
     1 #!/usr/bin/env python     2      3 from random import choice, random, randrange     4 from os.path import splitext     5 import EXIF     6 import PIL.Image     7 import itertools     8 import math     9 import sys    10     11 corners = [    12     (0, 0, 0), (255, 0, 0), (0, 255, 0), (255, 255, 0),    13     (0, 0, 255), (255, 0, 255), (0, 255, 255), (255, 255, 255)    14     ]    15     16 def distance(rgb1, rgb2):    17     r1, g1, b1 = rgb1    18     r2, g2, b2 = rgb2    19     return math.sqrt(pow(r1 - r2, 2) + pow(g1 - g2, 2) + pow(b1 - b2, 2))    20     21 def nearest(rgb, values):    22     l = [(distance(rgb, value), value) for value in values]    23     total = sum([d for d, value in l])    24     l = [(d / total, value) for d, value in l]    25     l.sort(reverse=True)    26     lc = [(d, value) for d, value in l if common(rgb, value)]    27     return get_from_distribution(rgb, lc or l)    28     29 def common(rgb1, rgb2):    30     return sum(map(lambda x: x[0] * x[1], zip(rgb1, rgb2))) != 0    31     32 def restore(srgb):    33     return tuple(map(lambda x: int(x * 255.0), srgb))    34     35 def scale(rgb):    36     return tuple(map(lambda x: x / 255.0, rgb))    37     38 def square(srgb):    39     return tuple(map(lambda x: pow(x, 2), srgb))    40     41 def invert(srgb):    42     return tuple(map(lambda x: 1.0 - x, srgb))    43     44 def combination(rgb):    45     rgb = square(scale(rgb))    46     rgbi = invert(rgb)    47     pairs = zip(rgbi, rgb)    48     d = []    49     for corner in corners:    50         rs, gs, bs = scale(corner)    51         d.append((pairs[0][int(rs)] * pairs[1][int(gs)] * pairs[2][int(bs)], corner))    52     return balance(d)    53     54 def complements(rgb):    55     r, g, b = rgb    56     return rgb, restore(invert(scale(rgb)))    57     58 def balance(d):    59     d = dict([(value, f) for f, value in d])    60     for primary, secondary in map(complements, [(255, 0, 0), (0, 255, 0), (0, 0, 255)]):    61         common = min(d[primary], d[secondary])    62         d[primary] -= common    63         d[secondary] -= common    64         d[(0, 0, 0)] += common    65         d[(255, 255, 255)] += common    66     return [(f, value) for value, f in d.items()]    67     68 def combine(d):    69     out = [0, 0, 0]    70     for v, rgb in d:    71         out[0] += v * rgb[0]    72         out[1] += v * rgb[1]    73         out[2] += v * rgb[2]    74     return out    75     76 def pattern(rgb):    77     l = combination(rgb)    78     l.sort(reverse=True)    79     return l    80     81 def get_from_distribution(rgb, dist):    82     choose = random()    83     threshold = 0    84     for f, c in dist:    85         threshold += f    86         if choose < threshold:    87             return c    88     return c    89     90 def get_value(rgb):    91     return get_from_distribution(rgb, pattern(rgb))    92     93 def sign(x):    94     return x >= 0 and 1 or -1    95     96 def saturate_rgb(rgb, exp):    97     return tuple([saturate_value(x, exp) for x in rgb])    98     99 def saturate_value(x, exp):   100     return int(127.5 + sign(x - 127.5) * 127.5 * pow(abs(x - 127.5) / 127.5, exp))   101    102 def replace(value, values):   103     if value not in values:   104         for i, v in list(enumerate(values))[::-1]:   105             if v != value:   106                 values[i] = value   107                 return   108    109 def test():   110     size = 512   111     for r in (0, 63, 127, 191, 255):   112         im = PIL.Image.new("RGB", (size, size))   113         for g in range(0, size):   114             for b in range(0, size):   115                 value = get_value((r, (g * 256) / size, (b * 256 / size)))   116                 im.putpixel((g, b), value)   117         im.save("rgb%d.png" % r)   118    119 def test_flat(rgb):   120     size = 64   121     im = PIL.Image.new("RGB", (size, size))   122     for y in range(0, size):   123         for x in range(0, size):   124             im.putpixel((x, y), get_value(rgb))   125     im.save("rgb%02d%02d%02d.png" % rgb)   126    127 def rotate_and_scale(im, width, height, rotate):   128     if rotate or x and x["Image Orientation"].values == [6L]:   129         im = im.rotate(270)   130    131     w, h = im.size   132     if w > h:   133         height = (width * h) / w   134     else:   135         width = (height * w) / h   136    137     return im.resize((width, height))   138    139 if __name__ == "__main__":   140     if "--test" in sys.argv:   141         test()   142         sys.exit(0)   143     elif "--test-flat" in sys.argv:   144         test_flat((120, 40, 60))   145         sys.exit(0)   146    147     width = 320   148     height = 256   149    150     input_filename, output_filename = sys.argv[1:3]   151     basename, ext = splitext(output_filename)   152     preview_filename = "".join([basename + "_preview", ext])   153    154     rotate = "-r" in sys.argv[3:]   155     saturate = sys.argv[3:].count("-s")   156     desaturate = sys.argv[3:].count("-d")   157     preview = "-p" in sys.argv[3:]   158     square = "-2" in sys.argv[3:] and square or (lambda x: x)   159    160     x = EXIF.process_file(open(input_filename))   161     im = PIL.Image.open(input_filename).convert("RGB")   162     im = rotate_and_scale(im, width, height, rotate)   163    164     width, height = im.size   165    166     colours = []   167    168     for y in range(0, height):   169         c = {}   170         for x in range(0, width):   171             rgb = im.getpixel((x, y))   172    173             # Saturate if requested.   174    175             if saturate or desaturate:   176                 rgb = saturate_rgb(rgb, saturate and math.pow(0.5, saturate) or math.pow(2, desaturate))   177                 im.putpixel((x, y), rgb)   178    179             # Sum the colour probabilities.   180    181             for f, value in combination(rgb):   182                 if not c.has_key(value):   183                     c[value] = f   184                 else:   185                     c[value] += f   186    187         c = [(n/width, value) for value, n in c.items()]   188         c.sort(reverse=True)   189         colours.append(c)   190    191     if preview:   192         imp = im.copy()   193         for y in range(0, height):   194             for x in range(0, width):   195                 rgb = imp.getpixel((x, y))   196                 value = get_value(rgb)   197                 imp.putpixel((x, y), value)   198         imp.save(preview_filename)   199    200     for y, c in enumerate(colours):   201         most = [value for n, value in c[:4]]   202         least = [value for n, value in c[4:]]   203    204         if least:   205             switched = []   206             for j in 1, 2:   207                 i = randrange(0, 4)   208                 n, value = c[i]   209                 if n < 0.1:   210                     switched.append(c[i])   211                     del c[i]   212             c += switched   213             most = [value for n, value in c[:4]]   214             least = [value for n, value in c[4:]]   215    216         for x in range(0, width):   217             rgb = im.getpixel((x, y))   218    219             # Get the requested colours and choose the closest alternative for   220             # less common colours.   221    222             value = get_value(rgb)   223             if value in least:   224                 value = nearest(rgb, most)   225    226             im.putpixel((x, y), value)   227    228     im.save(output_filename)   229    230 # vim: tabstop=4 expandtab shiftwidth=4