PaletteOptimiser

optimiser.py

82:7d63ab284ab8
2015-10-10 Paul Boddie Merged tidying changes. simpleimage
     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 sys    29     30 corners = [    31     (0, 0, 0), (255, 0, 0), (0, 255, 0), (255, 255, 0),    32     (0, 0, 255), (255, 0, 255), (0, 255, 255), (255, 255, 255)    33     ]    34     35 # Basic colour operations.    36     37 def within(v, lower, upper):    38     return min(max(v, lower), upper)    39     40 def clip(v):    41     return int(within(v, 0, 255))    42     43 def restore(srgb):    44     r, g, b = srgb    45     return int(r * 255.0), int(g * 255.0), int(b * 255.0)    46     47 def scale(rgb):    48     r, g, b = rgb    49     return r / 255.0, g / 255.0, b / 255.0    50     51 def invert(srgb):    52     r, g, b = srgb    53     return 1.0 - r, 1.0 - g, 1.0 - b    54     55 scaled_corners = map(scale, corners)    56 zipped_corners = zip(corners, scaled_corners)    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     srgb = scale(rgb)    68     rgbi = invert(srgb)    69     pairs = zip(rgbi, srgb)    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, scaled in zipped_corners:    77         rs, gs, bs = scaled    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 bases = [(0, 0, 0), (255, 0, 0), (0, 255, 0), (0, 0, 255)]    96 base_complements = map(complements, bases)    97     98 def balance(d):    99    100     """   101     Balance distribution 'd', cancelling opposing values and their complements   102     and replacing their common contributions with black and white contributions.   103     """   104    105     d = dict([(value, f) for f, value in d])   106     for primary, secondary in base_complements:   107         common = min(d[primary], d[secondary])   108         d[primary] -= common   109         d[secondary] -= common   110     return [(f, value) for value, f in d.items()]   111    112 def combine(d):   113    114     "Combine distribution 'd' to get a colour value."   115    116     out = [0, 0, 0]   117     for v, rgb in d:   118         out[0] += v * rgb[0]   119         out[1] += v * rgb[1]   120         out[2] += v * rgb[2]   121     return tuple(map(int, out))   122    123 def pattern(rgb, chosen=None):   124    125     """   126     Obtain a sorted colour distribution for 'rgb', optionally limited to any   127     specified 'chosen' colours.   128     """   129    130     l = [(f, value) for f, value in combination(rgb) if not chosen or value in chosen]   131     l.sort(reverse=True)   132     return l   133    134 def get_value(rgb, chosen=None, fail=False):   135    136     """   137     Get an output colour for 'rgb', optionally limited to any specified 'chosen'   138     colours. If 'fail' is set to a true value, return None if the colour cannot   139     be expressed using any of the chosen colours.   140     """   141    142     l = pattern(rgb, chosen)   143     limit = sum([f for f, c in l])   144     if not limit:   145         if fail:   146             return None   147         else:   148             return l[randrange(0, len(l))][1]   149    150     choose = random() * limit   151     threshold = 0   152     for f, c in l:   153         threshold += f   154         if choose < threshold:   155             return c   156     return c   157    158 # Colour processing operations.   159    160 def sign(x):   161     return x >= 0 and 1 or -1   162    163 def saturate_rgb(rgb, exp):   164     r, g, b = rgb   165     return saturate_value(r, exp), saturate_value(g, exp), saturate_value(b, exp)   166    167 def saturate_value(x, exp):   168     return int(127.5 + sign(x - 127.5) * 127.5 * pow(abs(x - 127.5) / 127.5, exp))   169    170 def amplify_rgb(rgb, exp):   171     r, g, b = rgb   172     return amplify_value(r, exp), amplify_value(g, exp), amplify_value(b, exp)   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     d = [(n/width, value) for value, n in c.items()]   197     d.sort(reverse=True)   198     return d   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    270     for y in range(0, height):   271         l = set()   272         for x in range(0, width):   273             l.add(im.getpixel((x, y)))   274         if len(l) > colours:   275             return (y, l)   276     return None   277    278 def process_image(pim, saturate, desaturate, darken, brighten):   279    280     """   281     Process image 'pim' using the given options: 'saturate', 'desaturate',   282     'darken', 'brighten'.   283     """   284    285     width, height = pim.size   286     im = SimpleImage(list(pim.getdata()), pim.size)   287    288     if saturate or desaturate or darken or brighten:   289         for y in range(0, height):   290             for x in range(0, width):   291                 rgb = im.getpixel((x, y))   292                 if saturate or desaturate:   293                     rgb = saturate_rgb(rgb, saturate and 0.5 / saturate or 2 * desaturate)   294                 if darken or brighten:   295                     rgb = amplify_rgb(rgb, brighten and 0.5 / brighten or 2 * darken)   296                 im.putpixel((x, y), rgb)   297    298     pim.putdata(im.getdata())   299    300 def preview_image(pim, half_resolution_preview=False):   301    302     "Return a preview copy of image 'pim'."   303    304     width, height = pim.size   305     imp = pim.copy()   306     im = SimpleImage(list(pim.getdata()), pim.size)   307     step = half_resolution_preview and 2 or 1   308    309     for y in range(0, height):   310         for x in range(0, width):   311             rgb = im.getpixel((x, y))   312             value = get_value(rgb)   313             im.putpixel((x, y), value)   314             if half_resolution_preview:   315                 im.putpixel((x+1, y), value)   316    317     imp.putdata(im.getdata())   318     return imp   319    320 def convert_image(pim):   321    322     "Convert image 'pim' to an appropriate output representation."   323    324     width, height = pim.size   325     im = SimpleImage(list(pim.getdata()), pim.size)   326    327     for y in range(0, height):   328         c = get_colours(im, y)   329    330         for l in get_combinations(c, 4):   331             most = [value for f, value in l]   332             for x in range(0, width):   333                 rgb = im.getpixel((x, y))   334                 value = get_value(rgb, most, True)   335                 if value is None:   336                     break # try next combination   337             else:   338                 break # use this combination   339         else:   340             most = [value for f, value in c[:4]] # use the first four   341    342         for x in range(0, width):   343             rgb = im.getpixel((x, y))   344             value = get_value(rgb, most)   345             im.putpixel((x, y), value)   346    347             if x < width - 1:   348                 rgbn = im.getpixel((x+1, y))   349                 rgbn = (   350                     clip(rgbn[0] + (rgb[0] - value[0]) / 4.0),   351                     clip(rgbn[1] + (rgb[1] - value[1]) / 4.0),   352                     clip(rgbn[2] + (rgb[2] - value[2]) / 4.0)   353                     )   354                 im.putpixel((x+1, y), rgbn)   355    356             if y < height - 1:   357                 rgbn = im.getpixel((x, y+1))   358                 rgbn = (   359                     clip(rgbn[0] + (rgb[0] - value[0]) / 2.0),   360                     clip(rgbn[1] + (rgb[1] - value[1]) / 2.0),   361                     clip(rgbn[2] + (rgb[2] - value[2]) / 2.0)   362                     )   363                 im.putpixel((x, y+1), rgbn)   364    365     pim.putdata(im.getdata())   366    367 def get_float(options, flag):   368     try:   369         i = options.index(flag)   370         if i+1 < len(options) and options[i+1].isdigit():   371             return float(options[i+1])   372         else:   373             return 1.0   374     except ValueError:   375         return 0.0   376    377 class SimpleImage:   378    379     "An image behaving like PIL.Image."   380    381     def __init__(self, data, size):   382         self.data = data   383         self.width, self.height = self.size = size   384    385     def copy(self):   386         return SimpleImage(self.data[:], self.size)   387    388     def getpixel(self, xy):   389         x, y = xy   390         return self.data[y * self.width + x]   391    392     def putpixel(self, xy, value):   393         x, y = xy   394         self.data[y * self.width + x] = value   395    396     def getdata(self):   397         return self.data   398    399 # Main program.   400    401 if __name__ == "__main__":   402    403     # Test options.   404    405     if "--test" in sys.argv:   406         test()   407         sys.exit(0)   408     elif "--test-flat" in sys.argv:   409         test_flat((120, 40, 60))   410         sys.exit(0)   411     elif "--help" in sys.argv:   412         print >>sys.stderr, """\   413 Usage: %s <input filename> <output filename> [ <options> ]   414    415 Options are...   416    417 -s - Saturate the input image (can be followed by a float, default 1.0)   418 -d - Desaturate the input image (can be followed by a float, default 1.0)   419 -D - Darken the input image (can be followed by a float, default 1.0)   420 -B - Brighten the input image (can be followed by a float, default 1.0)   421    422 -r - Rotate the input image clockwise   423 -p - Generate a separate preview image   424 -h - Make the preview image with half horizontal resolution (MODE 2)   425 -v - Verify the output image (loaded if -n is given)   426 -n - Generate no output image   427 """ % split(sys.argv[0])[1]   428         sys.exit(1)   429    430     width = 320   431     height = 256   432    433     input_filename, output_filename = sys.argv[1:3]   434     basename, ext = splitext(output_filename)   435     preview_filename = "".join([basename + "_preview", ext])   436    437     options = sys.argv[3:]   438    439     # Preprocessing options that can be repeated for extra effect.   440    441     saturate = get_float(options, "-s")   442     desaturate = get_float(options, "-d")   443     darken = get_float(options, "-D")   444     brighten = get_float(options, "-B")   445    446     # General output options.   447    448     rotate = "-r" in options   449     preview = "-p" in options   450     half_resolution_preview = "-h" in options   451     verify = "-v" in options   452     no_normal_output = "-n" in options   453     make_image = not no_normal_output   454    455     # Load the input image if requested.   456    457     if make_image or preview:   458         exif = EXIF.process_file(open(input_filename))   459         im = PIL.Image.open(input_filename).convert("RGB")   460         im = rotate_and_scale(exif, im, width, height, rotate)   461    462         process_image(im, saturate, desaturate, darken, brighten)   463    464     # Generate a preview if requested.   465    466     if preview:   467         preview_image(im, half_resolution_preview).save(preview_filename)   468    469     # Generate an output image if requested.   470    471     if make_image:   472         convert_image(im)   473         im.save(output_filename)   474    475     # Verify the output image (which may be loaded) if requested.   476    477     if verify:   478         if no_normal_output:   479             im = PIL.Image.open(output_filename).convert("RGB")   480    481         result = count_colours(im, 4)   482         if result is not None:   483             y, colours = result   484             print "Image %s: row %d has the following colours: %s" % (output_filename, y, "; ".join([repr(c) for c in colours]))   485    486 # vim: tabstop=4 expandtab shiftwidth=4