PaletteOptimiser

optimiser.py

62:aa75aca1ebd7
2015-10-09 Paul Boddie Removed the nearest function.
     1 #!/usr/bin/env python     2      3 """     4 Convert and optimise images for display in an Acorn Electron MODE 1 variant     5 with four colours per line but eight colours available for selection for each     6 line.     7      8 Copyright (C) 2015 Paul Boddie <paul@boddie.org.uk>     9     10 This program is free software; you can redistribute it and/or modify it under    11 the terms of the GNU General Public License as published by the Free Software    12 Foundation; either version 3 of the License, or (at your option) any later    13 version.    14     15 This program is distributed in the hope that it will be useful, but WITHOUT ANY    16 WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A    17 PARTICULAR PURPOSE.  See the GNU General Public License for more details.    18     19 You should have received a copy of the GNU General Public License along    20 with this program.  If not, see <http://www.gnu.org/licenses/>.    21 """    22     23 from random import random, randrange    24 from os.path import split, splitext    25 import EXIF    26 import PIL.Image    27 import itertools    28 import math    29 import sys    30     31 corners = [    32     (0, 0, 0), (255, 0, 0), (0, 255, 0), (255, 255, 0),    33     (0, 0, 255), (255, 0, 255), (0, 255, 255), (255, 255, 255)    34     ]    35     36 # Basic colour operations.    37     38 def extra(x): return x    39     40 def within(v, lower, upper):    41     return min(max(v, lower), upper)    42     43 def clip(v):    44     return within(v, 0, 255)    45     46 def distance(rgb1, rgb2):    47     r1, g1, b1 = rgb1    48     r2, g2, b2 = rgb2    49     return math.sqrt(pow(r1 - r2, 2) + pow(g1 - g2, 2) + pow(b1 - b2, 2))    50     51 def restore(srgb):    52     return tuple(map(lambda x: int(x * 255.0), srgb))    53     54 def scale(rgb):    55     return tuple(map(lambda x: x / 255.0, rgb))    56     57 def square(srgb):    58     return tuple(map(lambda x: pow(x, 2), srgb))    59     60 def invert(srgb):    61     return tuple(map(lambda x: 1.0 - x, srgb))    62     63 # Colour distribution functions.    64     65 cache = {}    66     67 def combination(rgb):    68     69     "Return the colour distribution for 'rgb'."    70     71     if not cache.has_key(rgb):    72     73         # Get the colour with components scaled from 0 to 1, plus the inverted    74         # component values.    75     76         rgb = extra(scale(rgb))    77         rgbi = invert(rgb)    78         pairs = zip(rgbi, rgb)    79     80         # For each corner of the colour cube (primary and secondary colours plus    81         # black and white), calculate the corner value's contribution to the    82         # input colour.    83     84         d = []    85         for corner in corners:    86             rs, gs, bs = scale(corner)    87     88             # Obtain inverted channel values where corner channels are low;    89             # obtain original channel values where corner channels are high.    90     91             d.append((pairs[0][int(rs)] * pairs[1][int(gs)] * pairs[2][int(bs)], corner))    92     93         # Balance the corner contributions.    94     95         cache[rgb] = balance(d)    96     97     return cache[rgb]    98     99 def complements(rgb):   100    101     "Return 'rgb' and its complement."   102    103     r, g, b = rgb   104     return rgb, restore(invert(scale(rgb)))   105    106 def balance(d):   107    108     """   109     Balance distribution 'd', cancelling opposing values and their complements   110     and replacing their common contributions with black and white contributions.   111     """   112    113     d = dict([(value, f) for f, value in d])   114     for primary, secondary in map(complements, [(0, 0, 0), (255, 0, 0), (0, 255, 0), (0, 0, 255)]):   115         common = min(d[primary], d[secondary])   116         d[primary] -= common   117         d[secondary] -= common   118     return [(f, value) for value, f in d.items()]   119    120 def combine(d):   121    122     "Combine distribution 'd' to get a colour value."   123    124     out = [0, 0, 0]   125     for v, rgb in d:   126         out[0] += v * rgb[0]   127         out[1] += v * rgb[1]   128         out[2] += v * rgb[2]   129     return out   130    131 def pattern(rgb, chosen=None):   132    133     """   134     Obtain a sorted colour distribution for 'rgb', optionally limited to any   135     specified 'chosen' colours.   136     """   137    138     l = [(f, value) for f, value in combination(rgb) if not chosen or value in chosen]   139     l.sort(reverse=True)   140     return l   141    142 def get_value(rgb, chosen=None, fail=False):   143    144     """   145     Get an output colour for 'rgb', optionally limited to any specified 'chosen'   146     colours. If 'fail' is set to a true value, return None if the colour cannot   147     be expressed using any of the chosen colours.   148     """   149    150     l = pattern(rgb, chosen)   151     limit = sum([f for f, c in l])   152     if not limit:   153         if fail:   154             return None   155         else:   156             return l[randrange(0, len(l))][1]   157    158     choose = random() * limit   159     threshold = 0   160     for f, c in l:   161         threshold += f   162         if choose < threshold:   163             return c   164     return c   165    166 # Colour processing operations.   167    168 def sign(x):   169     return x >= 0 and 1 or -1   170    171 def saturate_rgb(rgb, exp):   172     return tuple([saturate_value(x, exp) for x in rgb])   173    174 def saturate_value(x, exp):   175     return int(127.5 + sign(x - 127.5) * 127.5 * pow(abs(x - 127.5) / 127.5, exp))   176    177 def amplify_rgb(rgb, exp):   178     return tuple([amplify_value(x, exp) for x in rgb])   179    180 def amplify_value(x, exp):   181     return int(pow(x / 255.0, exp) * 255.0)   182    183 # Image operations.   184    185 def get_colours(im, y):   186    187     "Get a colour distribution from image 'im' for the row 'y'."   188    189     width, height = im.size   190     c = {}   191     for x in range(0, width):   192         rgb = im.getpixel((x, y))   193    194         # Sum the colour probabilities.   195    196         for f, value in combination(rgb):   197             if not c.has_key(value):   198                 c[value] = f   199             else:   200                 c[value] += f   201    202     c = [(n/width, value) for value, n in c.items()]   203     c.sort(reverse=True)   204     return c   205    206 def get_combinations(c, n):   207    208     """   209     Get combinations of colours from 'c' of size 'n' in decreasing order of   210     probability.   211     """   212    213     all = []   214     for l in itertools.combinations(c, n):   215         total = 0   216         for f, value in l:   217             total += f   218         all.append((total, l))   219     all.sort(reverse=True)   220     return [l for total, l in all]   221    222 def test():   223    224     "Generate slices of the colour cube."   225    226     size = 512   227     for r in (0, 63, 127, 191, 255):   228         im = PIL.Image.new("RGB", (size, size))   229         for g in range(0, size):   230             for b in range(0, size):   231                 value = get_value((r, (g * 256) / size, (b * 256 / size)))   232                 im.putpixel((g, b), value)   233         im.save("rgb%d.png" % r)   234    235 def test_flat(rgb):   236    237     "Generate a flat image for the colour 'rgb'."   238    239     size = 64   240     im = PIL.Image.new("RGB", (size, size))   241     for y in range(0, size):   242         for x in range(0, size):   243             im.putpixel((x, y), get_value(rgb))   244     im.save("rgb%02d%02d%02d.png" % rgb)   245    246 def rotate_and_scale(exif, im, width, height, rotate):   247    248     """   249     Using the given 'exif' information, rotate and scale image 'im' given the   250     indicated 'width' and 'height' constraints and any explicit 'rotate'   251     indication. The returned image will be within the given 'width' and   252     'height', filling either or both, and preserve its original aspect ratio.   253     """   254    255     if rotate or exif and exif["Image Orientation"].values == [6L]:   256         im = im.rotate(270)   257    258     w, h = im.size   259     if w > h:   260         height = (width * h) / w   261     else:   262         width = (height * w) / h   263    264     return im.resize((width, height))   265    266 def count_colours(im, colours):   267    268     """   269     Count colours on each row of image 'im', returning a tuple indicating the   270     first row with more than the given number of 'colours' together with the   271     found colours; otherwise returning None.   272     """   273    274     width, height = im.size   275     for y in range(0, height):   276         l = set()   277         for x in range(0, width):   278             l.add(im.getpixel((x, y)))   279         if len(l) > colours:   280             return (y, l)   281     return None   282    283 # Main program.   284    285 if __name__ == "__main__":   286    287     # Test options.   288    289     if "--test" in sys.argv:   290         test()   291         sys.exit(0)   292     elif "--test-flat" in sys.argv:   293         test_flat((120, 40, 60))   294         sys.exit(0)   295     elif "--help" in sys.argv:   296         print >>sys.stderr, """\   297 Usage: %s <input filename> <output filename> [ <options> ]   298    299 Options are...   300    301 -s - Saturate the input image (can be repeated)   302 -d - Desaturate the input image (can be repeated)   303 -D - Darken the input image (can be repeated)   304 -B - Brighten the input image (can be repeated)   305 -2 - Square/diminish the bright corner colour contributions (experimental)   306    307 -r - Rotate the input image clockwise   308 -p - Generate a separate preview image   309 -h - Make the preview image with half horizontal resolution (MODE 2)   310 -v - Verify the output image (loaded if -n is given)   311 -n - Generate no output image   312 """ % split(sys.argv[0])[1]   313         sys.exit(1)   314    315     width = 320   316     height = 256   317    318     input_filename, output_filename = sys.argv[1:3]   319     basename, ext = splitext(output_filename)   320     preview_filename = "".join([basename + "_preview", ext])   321    322     options = sys.argv[3:]   323    324     # Preprocessing options that can be repeated for extra effect.   325    326     saturate = options.count("-s")   327     desaturate = options.count("-d")   328     darken = options.count("-D")   329     brighten = options.count("-B")   330    331     # Experimental colour distribution modification.   332    333     use_square = "-2" in options   334     if use_square:   335         extra = square   336    337     # General output options.   338    339     rotate = "-r" in options   340     preview = "-p" in options   341     half_resolution_preview = "-h" in options   342     verify = "-v" in options   343     no_normal_output = "-n" in options   344     make_image = not no_normal_output   345    346     # Load the input image if requested.   347    348     if make_image or preview:   349         exif = EXIF.process_file(open(input_filename))   350         im = PIL.Image.open(input_filename).convert("RGB")   351         im = rotate_and_scale(exif, im, width, height, rotate)   352    353         width, height = im.size   354    355         if saturate or desaturate or darken or brighten:   356             for y in range(0, height):   357                 for x in range(0, width):   358                     rgb = im.getpixel((x, y))   359                     if saturate or desaturate:   360                         rgb = saturate_rgb(rgb, saturate and math.pow(0.5, saturate) or math.pow(2, desaturate))   361                     if darken or brighten:   362                         rgb = amplify_rgb(rgb, brighten and math.pow(0.5, brighten) or math.pow(2, darken))   363                     im.putpixel((x, y), rgb)   364    365     # Generate a preview if requested.   366    367     if preview:   368         imp = im.copy()   369         step = half_resolution_preview and 2 or 1   370         for y in range(0, height):   371             for x in range(0, width, step):   372                 rgb = imp.getpixel((x, y))   373                 value = get_value(rgb)   374                 imp.putpixel((x, y), value)   375                 if half_resolution_preview:   376                     imp.putpixel((x+1, y), value)   377    378         imp.save(preview_filename)   379    380     # Generate an output image if requested.   381    382     if make_image:   383         for y in range(0, height):   384             c = get_colours(im, y)   385    386             for l in get_combinations(c, 4):   387                 most = [value for f, value in l]   388                 for x in range(0, width):   389                     rgb = im.getpixel((x, y))   390                     value = get_value(rgb, most, True)   391                     if value is None:   392                         break # try next combination   393                 else:   394                     break # use this combination   395             else:   396                 most = [value for f, value in c[:4]] # use the first four   397    398             for x in range(0, width):   399                 rgb = im.getpixel((x, y))   400                 value = get_value(rgb, most)   401                 im.putpixel((x, y), value)   402    403                 if y < height - 1:   404                     rgbn = im.getpixel((x, y+1))   405                     rgbn = tuple(map(lambda i: clip(i[0] + i[1] - i[2]), zip(rgbn, rgb, value)))   406                     im.putpixel((x, y+1), rgbn)   407    408         im.save(output_filename)   409    410     # Verify the output image (which may be loaded) if requested.   411    412     if verify:   413         if no_normal_output:   414             im = PIL.Image.open(output_filename).convert("RGB")   415    416         result = count_colours(im, 4)   417         if result is not None:   418             y, colours = result   419             print "Image %s: row %d has the following colours: %s" % (output_filename, y, "; ".join([repr(c) for c in colours]))   420    421 # vim: tabstop=4 expandtab shiftwidth=4