PaletteOptimiser

optimiser.py

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