PaletteOptimiser

optimiser.py

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