PaletteOptimiser

optimiser.py

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