PaletteOptimiser

optimiser.py

33:da4132364cfd
2015-10-03 Paul Boddie Experiment with replacing colours with black and white.
     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 distribution(rgb, values=None):    22     l = []    23     total = 0    24     for c in values or corners:    25         d = distance(rgb, c)    26         if d == 0:    27             return [(1, c)]    28         l.append((pow(d, -3), c))    29         total += pow(d, -3)    30     return [(d / total, c) for d, c in l]    31     32 def pattern(rgb, values=None):    33     l = distribution(rgb, values or corners)    34     l.sort(reverse=True)    35     return l    36     37 def get_value(rgb, values=None):    38     choose = random()    39     threshold = 0    40     for f, c in pattern(rgb, values):    41         threshold += f    42         if choose < threshold:    43             return c    44     return c    45     46 def sign(x):    47     return x >= 0 and 1 or -1    48     49 def saturate_rgb(rgb, exp):    50     return tuple([saturate_value(x, exp) for x in rgb])    51     52 def saturate_value(x, exp):    53     return int(127.5 + sign(x - 127.5) * 127.5 * pow(abs(x - 127.5) / 127.5, exp))    54     55 def replace(value, values):    56     if value not in values:    57         for i, v in list(enumerate(values))[::-1]:    58             if v != value:    59                 values[i] = value    60                 return    61     62 def test():    63     size = 512    64     for r in (0, 63, 127, 191, 255):    65         im = PIL.Image.new("RGB", (size, size))    66         for g in range(0, size):    67             for b in range(0, size):    68                 value = get_value((r, (g * 256) / size, (b * 256 / size)))    69                 im.putpixel((g, b), value)    70         im.save("rgb%d.png" % r)    71     72 def test_flat(rgb):    73     size = 64    74     im = PIL.Image.new("RGB", (size, size))    75     for y in range(0, size):    76         for x in range(0, size):    77             im.putpixel((x, y), get_value(rgb))    78     im.save("rgb%02d%02d%02d.png" % rgb)    79     80 def rotate_and_scale(im, width, height, rotate):    81     if rotate or x and x["Image Orientation"].values == [6L]:    82         im = im.rotate(270)    83     84     w, h = im.size    85     if w > h:    86         height = (width * h) / w    87     else:    88         width = (height * w) / h    89     90     return im.resize((width, height))    91     92 if __name__ == "__main__":    93     if "--test" in sys.argv:    94         test()    95         sys.exit(0)    96     elif "--test-flat" in sys.argv:    97         test_flat((120, 40, 60))    98         sys.exit(0)    99    100     width = 320   101     height = 256   102    103     input_filename, output_filename = sys.argv[1:3]   104     basename, ext = splitext(output_filename)   105    106     rotate = "-r" in sys.argv[3:]   107     saturate = sys.argv[3:].count("-s")   108     desaturate = sys.argv[3:].count("-d")   109    110     x = EXIF.process_file(open(input_filename))   111     im = PIL.Image.open(input_filename).convert("RGB")   112     im = rotate_and_scale(im, width, height, rotate)   113    114     width, height = im.size   115    116     colours = []   117    118     for y in range(0, height):   119         c = {}   120         for x in range(0, width):   121             rgb = im.getpixel((x, y))   122    123             # Saturate if requested.   124    125             if saturate or desaturate:   126                 rgb = saturate_rgb(rgb, saturate and math.pow(0.5, saturate) or math.pow(2, desaturate))   127                 im.putpixel((x, y), rgb)   128    129             # Sum the colour probabilities.   130    131             for f, value in distribution(rgb):   132                 if not c.has_key(value):   133                     c[value] = f   134                 else:   135                     c[value] += f   136    137         c = [(n, value) for value, n in c.items()]   138         c.sort(reverse=True)   139         colours.append(c)   140    141     for y, c in enumerate(colours):   142         most = [value for n, value in c[:4]]   143         least = [value for n, value in c[4:]]   144    145         if least:   146             if (0, 0, 0) in least[:2]:   147                 replace((0, 0, 0), most)   148             if (255, 255, 255) in least[:2]:   149                 replace((255, 255, 255), most)   150    151         for x in range(0, width):   152             rgb = im.getpixel((x, y))   153    154             # Get the requested colours and choose the closest alternative for   155             # less common colours.   156    157             value = get_value(rgb)   158             if value in least:   159                 rgb = im.getpixel((x, y))   160                 value = get_value(value, most)   161    162             im.putpixel((x, y), value)   163    164     im.save(output_filename)   165    166 # vim: tabstop=4 expandtab shiftwidth=4