PaletteOptimiser

optimiser.py

30:fa52036b3d75
2015-10-02 Paul Boddie Fixed the preview filename.
     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 math     8 import sys     9     10 corners = [    11     (0, 0, 0), (255, 0, 0), (0, 255, 0), (255, 255, 0),    12     (0, 0, 255), (255, 0, 255), (0, 255, 255), (255, 255, 255)    13     ]    14     15 def distance(rgb1, rgb2):    16     r1, g1, b1 = rgb1    17     r2, g2, b2 = rgb2    18     return math.sqrt(pow(r1 - r2, 2) + pow(g1 - g2, 2) + pow(b1 - b2, 2))    19     20 def factor(start, end, rgb):    21     r1, g1, b1 = start    22     r2, g2, b2 = end    23     gr, gg, gb = r2 - r1, g2 - g1, b2 - b1    24     r, g, b = rgb    25     pr, pg, pb = r - r1, g - g1, b - b1    26     dp = pr * gr + pg * gg + pb * gb    27     return float(dp) / pow(distance(start, end), 2)    28     29 def nearest(rgb, values=None):    30     l = map(lambda c: (distance(rgb, c), c), values or corners)    31     l.sort()    32     return l    33     34 def pattern(rgb, values=None):    35     l = nearest(rgb, values)    36     start, end = l[0][1], l[1][1]    37     f = factor(start, end, rgb)    38     return start, end, f    39     40 def get_value(rgb, values=None):    41     rgb1, rgb2, f = pattern(rgb, values)    42     if random() < pow(f, 2):    43         return rgb2    44     else:    45         return rgb1    46     47 def sign(x):    48     return x >= 0 and 1 or -1    49     50 def saturate_rgb(rgb, exp):    51     return tuple([saturate_value(x, exp) for x in rgb])    52     53 def saturate_value(x, exp):    54     return int(127.5 + sign(x - 127.5) * 127.5 * pow(abs(x - 127.5) / 127.5, exp))    55     56 def test():    57     size = 512    58     for r in (0, 63, 127, 191, 255):    59         im = PIL.Image.new("RGB", (size, size))    60         for g in range(0, size):    61             for b in range(0, size):    62                 value = get_value((r, (g * 256) / size, (b * 256 / size)))    63                 im.putpixel((g, b), value)    64         im.save("rgb%d.png" % r)    65     66 def test_flat(rgb):    67     size = 64    68     im = PIL.Image.new("RGB", (size, size))    69     for y in range(0, size):    70         for x in range(0, size):    71             im.putpixel((x, y), get_value(rgb))    72     im.save("rgb%02d%02d%02d.png" % rgb)    73     74 def rotate_and_scale(im, width, height, rotate):    75     if rotate or x and x["Image Orientation"].values == [6L]:    76         im = im.rotate(270)    77     78     w, h = im.size    79     if w > h:    80         height = (width * h) / w    81     else:    82         width = (height * w) / h    83     84     return im.resize((width, height))    85     86 if __name__ == "__main__":    87     if "--test" in sys.argv:    88         test()    89         sys.exit(0)    90     elif "--test-flat" in sys.argv:    91         test_flat((120, 40, 60))    92         sys.exit(0)    93     94     width = 320    95     height = 256    96     97     input_filename, output_filename = sys.argv[1:3]    98     basename, ext = splitext(output_filename)    99     preview_filename = "".join([basename + "_preview", ext])   100    101     preview = "-p" in sys.argv[3:]   102     rotate = "-r" in sys.argv[3:]   103     saturate = sys.argv[3:].count("-s")   104     desaturate = sys.argv[3:].count("-d")   105    106     x = EXIF.process_file(open(input_filename))   107     im = PIL.Image.open(input_filename)   108     im = rotate_and_scale(im, width, height, rotate)   109    110     if preview:   111         im_preview = im.copy()   112    113     width, height = im.size   114    115     colours = []   116    117     for y in range(0, height):   118         c = {}   119         for x in range(0, width):   120             rgb = im.getpixel((x, y))   121    122             # Saturate if requested.   123    124             if saturate or desaturate:   125                 rgb = saturate_rgb(rgb, saturate and math.pow(0.5, saturate) or math.pow(2, desaturate))   126                 im.putpixel((x, y), rgb)   127    128             # Count the number of requested colours.   129    130             value = get_value(rgb)   131             if not c.has_key(value):   132                 c[value] = 1   133             else:   134                 c[value] += 1   135    136             if preview:   137                 im_preview.putpixel((x, y), value)   138    139         c = [(n, value) for value, n in c.items()]   140         c.sort(reverse=True)   141         colours.append(c)   142    143     if preview:   144         im_preview.save(preview_filename)   145    146     for y, c in enumerate(colours):   147         most = [value for n, value in c[:4]]   148         least = [value for n, value in c[4:]]   149    150         for x in range(0, width):   151             rgb = im.getpixel((x, y))   152    153             # Get the requested colours and choose the closest alternative for   154             # less common colours.   155    156             value = get_value(rgb)   157             if value in least:   158                 value = get_value(rgb, most)   159             im.putpixel((x, y), value)   160    161     im.save(output_filename)   162    163 # vim: tabstop=4 expandtab shiftwidth=4