PaletteOptimiser

optimiser.py

66:e82caf5f5a8a
2015-10-09 Paul Boddie Fixed option handling at the end of the argument list.
     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 def get_float(options, flag):   278     try:   279         i = options.index(flag)   280         if i+1 < len(options) and options[i+1].isdigit():   281             return float(options[i+1])   282         else:   283             return 1.0   284     except ValueError:   285         return 0.0   286    287 # Main program.   288    289 if __name__ == "__main__":   290    291     # Test options.   292    293     if "--test" in sys.argv:   294         test()   295         sys.exit(0)   296     elif "--test-flat" in sys.argv:   297         test_flat((120, 40, 60))   298         sys.exit(0)   299     elif "--help" in sys.argv:   300         print >>sys.stderr, """\   301 Usage: %s <input filename> <output filename> [ <options> ]   302    303 Options are...   304    305 -s - Saturate the input image (can be repeated)   306 -d - Desaturate the input image (can be repeated)   307 -D - Darken the input image (can be followed by a float, default 1.0)   308 -B - Brighten the input image (can be followed by a float, default 1.0)   309 -2 - Square/diminish the bright corner colour contributions (experimental)   310    311 -r - Rotate the input image clockwise   312 -p - Generate a separate preview image   313 -h - Make the preview image with half horizontal resolution (MODE 2)   314 -v - Verify the output image (loaded if -n is given)   315 -n - Generate no output image   316 """ % split(sys.argv[0])[1]   317         sys.exit(1)   318    319     width = 320   320     height = 256   321    322     input_filename, output_filename = sys.argv[1:3]   323     basename, ext = splitext(output_filename)   324     preview_filename = "".join([basename + "_preview", ext])   325    326     options = sys.argv[3:]   327    328     # Preprocessing options that can be repeated for extra effect.   329    330     saturate = options.count("-s")   331     desaturate = options.count("-d")   332     darken = get_float(options, "-D")   333     brighten = get_float(options, "-B")   334    335     # Experimental colour distribution modification.   336    337     use_square = "-2" in options   338     if use_square:   339         extra = square   340    341     # General output options.   342    343     rotate = "-r" in options   344     preview = "-p" in options   345     half_resolution_preview = "-h" in options   346     verify = "-v" in options   347     no_normal_output = "-n" in options   348     make_image = not no_normal_output   349    350     # Load the input image if requested.   351    352     if make_image or preview:   353         exif = EXIF.process_file(open(input_filename))   354         im = PIL.Image.open(input_filename).convert("RGB")   355         im = rotate_and_scale(exif, im, width, height, rotate)   356    357         width, height = im.size   358    359         if saturate or desaturate or darken or brighten:   360             for y in range(0, height):   361                 for x in range(0, width):   362                     rgb = im.getpixel((x, y))   363                     if saturate or desaturate:   364                         rgb = saturate_rgb(rgb, saturate and math.pow(0.5, saturate) or math.pow(2, desaturate))   365                     if darken or brighten:   366                         rgb = amplify_rgb(rgb, brighten and 0.5 / brighten or 2 * darken)   367                     im.putpixel((x, y), rgb)   368    369     # Generate a preview if requested.   370    371     if preview:   372         imp = im.copy()   373         step = half_resolution_preview and 2 or 1   374         for y in range(0, height):   375             for x in range(0, width, step):   376                 rgb = imp.getpixel((x, y))   377                 value = get_value(rgb)   378                 imp.putpixel((x, y), value)   379                 if half_resolution_preview:   380                     imp.putpixel((x+1, y), value)   381    382         imp.save(preview_filename)   383    384     # Generate an output image if requested.   385    386     if make_image:   387         for y in range(0, height):   388             c = get_colours(im, y)   389    390             for l in get_combinations(c, 4):   391                 most = [value for f, value in l]   392                 for x in range(0, width):   393                     rgb = im.getpixel((x, y))   394                     value = get_value(rgb, most, True)   395                     if value is None:   396                         break # try next combination   397                 else:   398                     break # use this combination   399             else:   400                 most = [value for f, value in c[:4]] # use the first four   401    402             for x in range(0, width):   403                 rgb = im.getpixel((x, y))   404                 value = get_value(rgb, most)   405                 im.putpixel((x, y), value)   406    407                 if y < height - 1:   408                     rgbn = im.getpixel((x, y+1))   409                     rgbn = tuple(map(lambda i: clip(i[0] + i[1] - i[2]), zip(rgbn, rgb, value)))   410                     im.putpixel((x, y+1), rgbn)   411    412         im.save(output_filename)   413    414     # Verify the output image (which may be loaded) if requested.   415    416     if verify:   417         if no_normal_output:   418             im = PIL.Image.open(output_filename).convert("RGB")   419    420         result = count_colours(im, 4)   421         if result is not None:   422             y, colours = result   423             print "Image %s: row %d has the following colours: %s" % (output_filename, y, "; ".join([repr(c) for c in colours]))   424    425 # vim: tabstop=4 expandtab shiftwidth=4