PaletteOptimiser

optimiser.py

10:c1d09a02aa21
2015-09-09 Paul Boddie Handle lines with fewer than four base colours properly.
     1 #!/usr/bin/env python     2      3 from array import array     4 from itertools import combinations     5 import PIL.Image     6 import sys     7      8 def scale(v):     9     return (v + 43) / 85    10     11 def point(rgb):    12     return tuple(map(scale, rgb))    13     14 def index(p):    15     return p[0] * 16 + p[1] * 4 + p[2]    16     17 def colour(i):    18     return (255 * (i % 2), 255 * ((i / 2) % 2), 255 * ((i / 4) % 2))    19     20 def add(d, v):    21     d[v] = (d.has_key(v) and d[v] or 0) + 1    22     23 def by_frequency(d):    24     l = [(f, t) for (t, f) in d.items()]    25     l.sort(reverse=True)    26     return [i[1] for i in l]    27     28 def match(b, bases):    29     return b in bases and b    30     31 def fallback(bases):    32     for b in by_frequency(bases):    33         if b not in ["_", "W"]:    34             return b    35     return by_frequency(bases)[0]    36     37 tones = [    38     "___", "_BB", "BBB", "BBB", # 00x    39     "_GG", "_GC", "GGC", "GCC", # 01x    40     "GGG", "GGC", "GCC", "GCC", # 02x    41     "GGG", "GCC", "CCC", "CCC", # 03x    42     "_RR", "_MM", "MMB", "MBB", # 10x    43     "_YY", "_**", "_*B", "BBW", # 11x    44     "GYY", "GGW", "GCC", "CCW", # 12x    45     "GGY", "GGG", "GCC", "CCW", # 13x    46     "RRR", "RRM", "RMM", "MMM", # 20x    47     "RYY", "RRW", "RMW", "MMW", # 21x    48     "YYY", "YYW", "**W", "WWW", # 22x    49     "YYY", "YYW", "YWW", "WWW", # 23x    50     "RRR", "RMM", "RMM", "MMM", # 30x    51     "RRY", "RRY", "RMM", "MMW", # 31x    52     "YYY", "YYW", "YYW", "WWW", # 32x    53     "YYY", "YYW", "YYW", "WWW", # 33x    54     ]    55     56 colours = ["_", "R", "G", "Y", "B", "M", "C", "W"]    57     58 if __name__ == "__main__":    59     width = 320    60     input_filename, output_filename = sys.argv[1:3]    61     62     im = PIL.Image.open(input_filename)    63     w, h = im.size    64     height = (width * h) / w    65     im = im.resize((width, height))    66     67     usage = []    68     base_usage = []    69     toned = []    70     71     for row in range(0, height):    72         u = {}    73         usage.append(u)    74         bu = {}    75         base_usage.append(bu)    76         tr = []    77         toned.append(tr)    78         for column in range(0, width):    79             rgb = im.getpixel((column, row))    80             p = point(rgb)    81             i = index(p)    82             t = tones[i]    83             add(u, t)    84             if t[0] != "*":    85                 add(bu, t[0])    86             if t[1] != "*":    87                 add(bu, t[1])    88             if t[2] != "*":    89                 add(bu, t[2])    90             tr.append(t)    91     92     chosen = []    93     94     for row, (u, bu) in enumerate(zip(usage, base_usage)):    95         light = row % 2    96         best = 0    97         best_bases = None    98         best_missing = None    99         best_map = None   100    101         for bases in combinations(bu, min(len(bu), 4)):   102             bases = dict([(base, bu[base]) for base in bases])   103             count = 0   104             missing = []   105             tone_map = {}   106             for tone, freq in u.items():   107                 base = match(tone[1], bases) or match(light and tone[2] or tone[0], bases)   108                 if base:   109                     tone_map[tone] = base   110                     count += freq   111                 else:   112                     missing.append(tone)   113             if count > best:   114                 best_bases = bases   115                 best_missing = missing   116                 best_map = tone_map   117                 best = count   118    119         chosen.append((best, best_bases or bases, best_map or tone_map, best_missing or missing))   120    121     output = []   122    123     for row, (tr, ch) in enumerate(zip(toned, chosen)):   124         o = []   125         for column, t in enumerate(tr):   126             best, bases, tone_map, missing = ch   127             base = tone_map.get(t) or fallback(bases)   128             o.append(base)   129             i = colours.index(base)   130             im.putpixel((column, row), colour(i))   131    132         output.append("".join(o))   133    134     im.save(output_filename)   135    136 # vim: tabstop=4 expandtab shiftwidth=4