PaletteOptimiser

optimiser.py

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