PaletteOptimiser

optimiser.py

31:244292b4816c
2015-10-02 Paul Boddie Use a proper main program!
     1 #!/usr/bin/env python     2      3 from dither import dither     4 from random import random     5 from os.path import splitext     6 import EXIF     7 import PIL.Image     8 import itertools     9 import math    10 import sys    11     12 corners = [    13     (0, 0, 0), (255, 0, 0), (0, 255, 0), (255, 255, 0),    14     (0, 0, 255), (255, 0, 255), (0, 255, 255), (255, 255, 255)    15     ]    16     17 def distance(rgb1, rgb2):    18     r1, g1, b1 = rgb1    19     r2, g2, b2 = rgb2    20     return math.sqrt(pow(r1 - r2, 2) + pow(g1 - g2, 2) + pow(b1 - b2, 2))    21     22 def distribution(rgb, values=None):    23     l = []    24     total = 0    25     for c in values or corners:    26         d = distance(rgb, c)    27         if d == 0:    28             return [(1, c)]    29         l.append((pow(d, -3), c))    30         total += pow(d, -3)    31     return [(d / total, c) for d, c in l]    32     33 def pattern(rgb, values=None):    34     l = distribution(rgb, values or corners)    35     l.sort(reverse=True)    36     return l    37     38 def get_value(rgb, values=None):    39     choose = random()    40     threshold = 0    41     for f, c in pattern(rgb, values):    42         threshold += f    43         if choose < threshold:    44             return c    45     return c    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 get_best(im, width, y, colours):    57     best = 0    58     best_values = None    59     60     for values in itertools.combinations([value for n, value in colours], 4):    61         current = 0    62         for x in range(0, width):    63             rgb = im.getpixel((x, y))    64             for f, value in distribution(rgb, values):    65                 current += f    66         if current >= best:    67             best_values = values    68             best = current    69     70     return best_values    71     72 def test():    73     size = 512    74     for r in (0, 63, 127, 191, 255):    75         im = PIL.Image.new("RGB", (size, size))    76         for g in range(0, size):    77             for b in range(0, size):    78                 value = get_value((r, (g * 256) / size, (b * 256 / size)))    79                 im.putpixel((g, b), value)    80         im.save("rgb%d.png" % r)    81     82 def test_flat(rgb):    83     size = 64    84     im = PIL.Image.new("RGB", (size, size))    85     for y in range(0, size):    86         for x in range(0, size):    87             im.putpixel((x, y), get_value(rgb))    88     im.save("rgb%02d%02d%02d.png" % rgb)    89     90 def rotate_and_scale(im, width, height, rotate):    91     if rotate or x and x["Image Orientation"].values == [6L]:    92         im = im.rotate(270)    93     94     w, h = im.size    95     if w > h:    96         height = (width * h) / w    97     else:    98         width = (height * w) / h    99    100     return im.resize((width, height))   101    102 if __name__ == "__main__":   103     if "--test" in sys.argv:   104         test()   105         sys.exit(0)   106     elif "--test-flat" in sys.argv:   107         test_flat((120, 40, 60))   108         sys.exit(0)   109    110     width = 320   111     height = 256   112    113     input_filename, output_filename = sys.argv[1:3]   114     basename, ext = splitext(output_filename)   115    116     rotate = "-r" in sys.argv[3:]   117     saturate = sys.argv[3:].count("-s")   118     desaturate = sys.argv[3:].count("-d")   119     best = "-b" in sys.argv[3:]   120     dithering = "-d" in sys.argv[3:]   121    122     x = EXIF.process_file(open(input_filename))   123     im = PIL.Image.open(input_filename).convert("RGB")   124     im = rotate_and_scale(im, width, height, rotate)   125    126     width, height = im.size   127    128     im_original = im   129    130     if dithering:   131         im = dither(im).convert("RGB")   132    133     colours = []   134    135     for y in range(0, height):   136         c = {}   137         for x in range(0, width):   138             rgb = im.getpixel((x, y))   139    140             # Saturate if requested.   141    142             if saturate or desaturate:   143                 rgb = saturate_rgb(rgb, saturate and math.pow(0.5, saturate) or math.pow(2, desaturate))   144                 im.putpixel((x, y), rgb)   145    146             # Sum the colour probabilities.   147    148             for f, value in distribution(rgb):   149                 if not c.has_key(value):   150                     c[value] = f   151                 else:   152                     c[value] += f   153    154         c = [(n, value) for value, n in c.items()]   155         c.sort(reverse=True)   156         colours.append(c)   157    158     for y, c in enumerate(colours):   159         most = [value for n, value in c[:4]]   160         least = [value for n, value in c[4:]]   161    162         if least and best:   163             most = get_best(im_original, width, y, c)   164    165         for x in range(0, width):   166             rgb = im.getpixel((x, y))   167    168             # Get the requested colours and choose the closest alternative for   169             # less common colours.   170    171             value = get_value(rgb)   172             if value in least:   173                 rgb = im_original.getpixel((x, y))   174                 value = get_value(rgb, most)   175    176             im.putpixel((x, y), value)   177    178     im.save(output_filename)   179    180 # vim: tabstop=4 expandtab shiftwidth=4