PaletteOptimiser

optimiser.py

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