PaletteOptimiser

optimiser.py

55:e017e76791a0
2015-10-07 Paul Boddie Reduced the propagated error by half.
     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 within(v, lower, upper):    17     return min(max(v, lower), upper)    18     19 def clip(v):    20     return int(within(v, 0, 255))    21     22 def distance(rgb1, rgb2):    23     r1, g1, b1 = rgb1    24     r2, g2, b2 = rgb2    25     return math.sqrt(pow(r1 - r2, 2) + pow(g1 - g2, 2) + pow(b1 - b2, 2))    26     27 def nearest(rgb, values):    28     l = [(distance(rgb, value), value) for value in values]    29     l.sort()    30     return l[0][1]    31     32 def restore(srgb):    33     return tuple(map(lambda x: int(x * 255.0), srgb))    34     35 def scale(rgb):    36     return tuple(map(lambda x: x / 255.0, rgb))    37     38 def square(srgb):    39     return tuple(map(lambda x: pow(x, 2), srgb))    40     41 def invert(srgb):    42     return tuple(map(lambda x: 1.0 - x, srgb))    43     44 cache = {}    45     46 def combination(rgb):    47     if not cache.has_key(rgb):    48         rgb = square(scale(rgb))    49         rgbi = invert(rgb)    50         pairs = zip(rgbi, rgb)    51         d = []    52         for corner in corners:    53             rs, gs, bs = scale(corner)    54             d.append((pairs[0][int(rs)] * pairs[1][int(gs)] * pairs[2][int(bs)], corner))    55         cache[rgb] = balance(d)    56     return cache[rgb]    57     58 def complements(rgb):    59     r, g, b = rgb    60     return rgb, restore(invert(scale(rgb)))    61     62 def balance(d):    63     d = dict([(value, f) for f, value in d])    64     for primary, secondary in map(complements, [(255, 0, 0), (0, 255, 0), (0, 0, 255)]):    65         common = min(d[primary], d[secondary])    66         d[primary] -= common    67         d[secondary] -= common    68         d[(0, 0, 0)] += common    69         d[(255, 255, 255)] += common    70     return [(f, value) for value, f in d.items()]    71     72 def compensate(d, chosen):    73     dd = dict([(value, f) for f, value in d])    74     for f, value in d:    75         if value in chosen:    76             _value, complement = complements(value)    77             if complement not in chosen:    78                 f = max(0, f - dd[complement])    79                 dd[value] = f    80     return [(f, value) for value, f in dd.items() if value in chosen]    81     82 def combine(d):    83     out = [0, 0, 0]    84     for v, rgb in d:    85         out[0] += v * rgb[0]    86         out[1] += v * rgb[1]    87         out[2] += v * rgb[2]    88     return out    89     90 def pattern(rgb, chosen=None):    91     l = combination(rgb)    92     if chosen:    93         l = compensate(l, chosen)    94     l.sort(reverse=True)    95     return l    96     97 def get_value(rgb, chosen=None):    98     l = pattern(rgb, chosen)    99     limit = sum([f for f, c in l])   100     choose = random() * limit   101     threshold = 0   102     for f, c in l:   103         threshold += f   104         if choose < threshold:   105             return c   106     return c   107    108 def sign(x):   109     return x >= 0 and 1 or -1   110    111 def saturate_rgb(rgb, exp):   112     return tuple([saturate_value(x, exp) for x in rgb])   113    114 def saturate_value(x, exp):   115     return int(127.5 + sign(x - 127.5) * 127.5 * pow(abs(x - 127.5) / 127.5, exp))   116    117 def amplify_rgb(rgb, exp):   118     return tuple([amplify_value(x, exp) for x in rgb])   119    120 def amplify_value(x, exp):   121     return int(pow(x / 255.0, exp) * 255.0)   122    123 def get_colours(im, y):   124     width, height = im.size   125     c = {}   126     for x in range(0, width):   127         rgb = im.getpixel((x, y))   128    129         # Sum the colour probabilities.   130    131         for f, value in combination(rgb):   132             if not c.has_key(value):   133                 c[value] = f   134             else:   135                 c[value] += f   136    137     c = [(n/width, value) for value, n in c.items()]   138     c.sort(reverse=True)   139     return c   140    141 def test():   142     size = 512   143     for r in (0, 63, 127, 191, 255):   144         im = PIL.Image.new("RGB", (size, size))   145         for g in range(0, size):   146             for b in range(0, size):   147                 value = get_value((r, (g * 256) / size, (b * 256 / size)))   148                 im.putpixel((g, b), value)   149         im.save("rgb%d.png" % r)   150    151 def test_flat(rgb):   152     size = 64   153     im = PIL.Image.new("RGB", (size, size))   154     for y in range(0, size):   155         for x in range(0, size):   156             im.putpixel((x, y), get_value(rgb))   157     im.save("rgb%02d%02d%02d.png" % rgb)   158    159 def rotate_and_scale(exif, im, width, height, rotate):   160     if rotate or exif and exif["Image Orientation"].values == [6L]:   161         im = im.rotate(270)   162    163     w, h = im.size   164     if w > h:   165         height = (width * h) / w   166     else:   167         width = (height * w) / h   168    169     return im.resize((width, height))   170    171 def count_colours(im, colours):   172     width, height = im.size   173     for y in range(0, height):   174         l = set()   175         for x in range(0, width):   176             l.add(im.getpixel((x, y)))   177         if len(l) > colours:   178             return (y, l)   179     return None   180    181 if __name__ == "__main__":   182     if "--test" in sys.argv:   183         test()   184         sys.exit(0)   185     elif "--test-flat" in sys.argv:   186         test_flat((120, 40, 60))   187         sys.exit(0)   188    189     width = 320   190     height = 256   191    192     input_filename, output_filename = sys.argv[1:3]   193     basename, ext = splitext(output_filename)   194     preview_filename = "".join([basename + "_preview", ext])   195    196     options = sys.argv[3:]   197    198     rotate = "-r" in options   199     saturate = options.count("-s")   200     desaturate = options.count("-d")   201     darken = options.count("-D")   202     brighten = options.count("-B")   203     square = "-2" in options and square or (lambda x: x)   204     preview = "-p" in options   205     half_resolution_preview = "-h" in options   206     verify = "-v" in options   207     no_normal_output = "-n" in options   208     make_image = not no_normal_output   209    210     if make_image or preview:   211         exif = EXIF.process_file(open(input_filename))   212         im = PIL.Image.open(input_filename).convert("RGB")   213         im = rotate_and_scale(exif, im, width, height, rotate)   214    215         width, height = im.size   216    217         if saturate or desaturate or darken or brighten:   218             for y in range(0, height):   219                 for x in range(0, width):   220                     rgb = im.getpixel((x, y))   221                     if saturate or desaturate:   222                         rgb = saturate_rgb(rgb, saturate and math.pow(0.5, saturate) or math.pow(2, desaturate))   223                     if darken or brighten:   224                         rgb = amplify_rgb(rgb, brighten and math.pow(0.5, brighten) or math.pow(2, darken))   225                     im.putpixel((x, y), rgb)   226    227     if preview:   228         imp = im.copy()   229         step = half_resolution_preview and 2 or 1   230         for y in range(0, height):   231             for x in range(0, width, step):   232                 rgb = imp.getpixel((x, y))   233                 value = get_value(rgb)   234                 imp.putpixel((x, y), value)   235                 if half_resolution_preview:   236                     imp.putpixel((x+1, y), value)   237    238         imp.save(preview_filename)   239    240     if make_image:   241         for y in range(0, height):   242             c = get_colours(im, y)   243             most = [value for n, value in c[:4]]   244             least = [value for n, value in c[4:]]   245    246             for x in range(0, width):   247                 rgb = im.getpixel((x, y))   248                 value = get_value(rgb, most)   249                 im.putpixel((x, y), value)   250    251                 if y < height - 1:   252                     rgbn = im.getpixel((x, y+1))   253                     rgbn = tuple(map(lambda i: clip(i[0] + (i[1] - i[2]) / 2.0), zip(rgbn, rgb, value)))   254                     im.putpixel((x, y+1), rgbn)   255    256         im.save(output_filename)   257    258     if verify:   259         if no_normal_output:   260             im = PIL.Image.open(output_filename).convert("RGB")   261    262         result = count_colours(im, 4)   263         if result is not None:   264             y, colours = result   265             print "Row %d has the following colours: %s" % (y, "; ".join([repr(c) for c in colours]))   266    267 # vim: tabstop=4 expandtab shiftwidth=4