PaletteOptimiser

optimiser.py

32:3a950efbc68f
2015-10-03 Paul Boddie Sum colour probabilities instead of counting requested colours. Find the best alternative for each unallocated colour from the allocated ones. Removed the preview image saving.
     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 test():    56     size = 512    57     for r in (0, 63, 127, 191, 255):    58         im = PIL.Image.new("RGB", (size, size))    59         for g in range(0, size):    60             for b in range(0, size):    61                 value = get_value((r, (g * 256) / size, (b * 256 / size)))    62                 im.putpixel((g, b), value)    63         im.save("rgb%d.png" % r)    64     65 def test_flat(rgb):    66     size = 64    67     im = PIL.Image.new("RGB", (size, size))    68     for y in range(0, size):    69         for x in range(0, size):    70             im.putpixel((x, y), get_value(rgb))    71     im.save("rgb%02d%02d%02d.png" % rgb)    72     73 def rotate_and_scale(im, width, height, rotate):    74     if rotate or x and x["Image Orientation"].values == [6L]:    75         im = im.rotate(270)    76     77     w, h = im.size    78     if w > h:    79         height = (width * h) / w    80     else:    81         width = (height * w) / h    82     83     return im.resize((width, height))    84     85 if __name__ == "__main__":    86     if "--test" in sys.argv:    87         test()    88         sys.exit(0)    89     elif "--test-flat" in sys.argv:    90         test_flat((120, 40, 60))    91         sys.exit(0)    92     93     width = 320    94     height = 256    95     96     input_filename, output_filename = sys.argv[1:3]    97     basename, ext = splitext(output_filename)    98     99     rotate = "-r" in sys.argv[3:]   100     saturate = sys.argv[3:].count("-s")   101     desaturate = sys.argv[3:].count("-d")   102    103     x = EXIF.process_file(open(input_filename))   104     im = PIL.Image.open(input_filename).convert("RGB")   105     im = rotate_and_scale(im, width, height, rotate)   106    107     width, height = im.size   108    109     colours = []   110    111     for y in range(0, height):   112         c = {}   113         for x in range(0, width):   114             rgb = im.getpixel((x, y))   115    116             # Saturate if requested.   117    118             if saturate or desaturate:   119                 rgb = saturate_rgb(rgb, saturate and math.pow(0.5, saturate) or math.pow(2, desaturate))   120                 im.putpixel((x, y), rgb)   121    122             # Sum the colour probabilities.   123    124             for f, value in distribution(rgb):   125                 if not c.has_key(value):   126                     c[value] = f   127                 else:   128                     c[value] += f   129    130         c = [(n, value) for value, n in c.items()]   131         c.sort(reverse=True)   132         colours.append(c)   133    134     for y, c in enumerate(colours):   135         most = [value for n, value in c[:4]]   136         least = [value for n, value in c[4:]]   137    138         for x in range(0, width):   139             rgb = im.getpixel((x, y))   140    141             # Get the requested colours and choose the closest alternative for   142             # less common colours.   143    144             value = get_value(rgb)   145             if value in least:   146                 rgb = im.getpixel((x, y))   147                 value = pattern(value, most)[0][1]   148    149             im.putpixel((x, y), value)   150    151     im.save(output_filename)   152    153 # vim: tabstop=4 expandtab shiftwidth=4