PaletteOptimiser

optimiser.py

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