PaletteOptimiser

optimiser.py

83:3f96e51bfacb
2015-10-10 Paul Boddie Made a new module for Shedskin to compile as an extension module. simpleimage-shedskin
     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 optimiserlib import *    24 from os.path import split, splitext    25 import EXIF    26 import PIL.Image    27 import itertools    28 import sys    29     30 # Image operations.    31     32 def get_colours(im, y):    33     34     "Get a colour distribution from image 'im' for the row 'y'."    35     36     width, height = im.size    37     c = {}    38     for x in range(0, width):    39         rgb = im.getpixel((x, y))    40     41         # Sum the colour probabilities.    42     43         for f, value in combination(rgb):    44             if not c.has_key(value):    45                 c[value] = f    46             else:    47                 c[value] += f    48     49     d = [(n/width, value) for value, n in c.items()]    50     d.sort(reverse=True)    51     return d    52     53 def get_combinations(c, n):    54     55     """    56     Get combinations of colours from 'c' of size 'n' in decreasing order of    57     probability.    58     """    59     60     all = []    61     for l in itertools.combinations(c, n):    62         total = 0    63         for f, value in l:    64             total += f    65         all.append((total, l))    66     all.sort(reverse=True)    67     return [l for total, l in all]    68     69 def test():    70     71     "Generate slices of the colour cube."    72     73     size = 512    74     for r in (0, 63, 127, 191, 255):    75         im = PIL.Image.new("RGB", (size, size))    76         for g in range(0, size):    77             for b in range(0, size):    78                 value = get_value((r, (g * 256) / size, (b * 256 / size)))    79                 im.putpixel((g, b), value)    80         im.save("rgb%d.png" % r)    81     82 def test_flat(rgb):    83     84     "Generate a flat image for the colour 'rgb'."    85     86     size = 64    87     im = PIL.Image.new("RGB", (size, size))    88     for y in range(0, size):    89         for x in range(0, size):    90             im.putpixel((x, y), get_value(rgb))    91     im.save("rgb%02d%02d%02d.png" % rgb)    92     93 def rotate_and_scale(exif, im, width, height, rotate):    94     95     """    96     Using the given 'exif' information, rotate and scale image 'im' given the    97     indicated 'width' and 'height' constraints and any explicit 'rotate'    98     indication. The returned image will be within the given 'width' and    99     'height', filling either or both, and preserve its original aspect ratio.   100     """   101    102     if rotate or exif and exif["Image Orientation"].values == [6L]:   103         im = im.rotate(270)   104    105     w, h = im.size   106     if w > h:   107         height = (width * h) / w   108     else:   109         width = (height * w) / h   110    111     return im.resize((width, height))   112    113 def count_colours(im, colours):   114    115     """   116     Count colours on each row of image 'im', returning a tuple indicating the   117     first row with more than the given number of 'colours' together with the   118     found colours; otherwise returning None.   119     """   120    121     width, height = im.size   122    123     for y in range(0, height):   124         l = set()   125         for x in range(0, width):   126             l.add(im.getpixel((x, y)))   127         if len(l) > colours:   128             return (y, l)   129     return None   130    131 def process_image(pim, saturate, desaturate, darken, brighten):   132    133     """   134     Process image 'pim' using the given options: 'saturate', 'desaturate',   135     'darken', 'brighten'.   136     """   137    138     width, height = pim.size   139     im = SimpleImage(list(pim.getdata()), pim.size)   140    141     if saturate or desaturate or darken or brighten:   142         for y in range(0, height):   143             for x in range(0, width):   144                 rgb = im.getpixel((x, y))   145                 if saturate or desaturate:   146                     rgb = saturate_rgb(rgb, saturate and 0.5 / saturate or 2 * desaturate)   147                 if darken or brighten:   148                     rgb = amplify_rgb(rgb, brighten and 0.5 / brighten or 2 * darken)   149                 im.putpixel((x, y), rgb)   150    151     pim.putdata(im.getdata())   152    153 def preview_image(pim, half_resolution_preview=False):   154    155     "Return a preview copy of image 'pim'."   156    157     width, height = pim.size   158     imp = pim.copy()   159     im = SimpleImage(list(pim.getdata()), pim.size)   160     step = half_resolution_preview and 2 or 1   161    162     for y in range(0, height):   163         for x in range(0, width):   164             rgb = im.getpixel((x, y))   165             value = get_value(rgb)   166             im.putpixel((x, y), value)   167             if half_resolution_preview:   168                 im.putpixel((x+1, y), value)   169    170     imp.putdata(im.getdata())   171     return imp   172    173 def convert_image(pim):   174    175     "Convert image 'pim' to an appropriate output representation."   176    177     width, height = pim.size   178     im = SimpleImage(list(pim.getdata()), pim.size)   179    180     for y in range(0, height):   181         c = get_colours(im, y)   182    183         for l in get_combinations(c, 4):   184             most = [value for f, value in l]   185             for x in range(0, width):   186                 rgb = im.getpixel((x, y))   187                 value = get_value(rgb, most, True)   188                 if value is None:   189                     break # try next combination   190             else:   191                 break # use this combination   192         else:   193             most = [value for f, value in c[:4]] # use the first four   194    195         for x in range(0, width):   196             rgb = im.getpixel((x, y))   197             value = get_value(rgb, most)   198             im.putpixel((x, y), value)   199    200             if x < width - 1:   201                 rgbn = im.getpixel((x+1, y))   202                 rgbn = (   203                     clip(rgbn[0] + (rgb[0] - value[0]) / 4.0),   204                     clip(rgbn[1] + (rgb[1] - value[1]) / 4.0),   205                     clip(rgbn[2] + (rgb[2] - value[2]) / 4.0)   206                     )   207                 im.putpixel((x+1, y), rgbn)   208    209             if y < height - 1:   210                 rgbn = im.getpixel((x, y+1))   211                 rgbn = (   212                     clip(rgbn[0] + (rgb[0] - value[0]) / 2.0),   213                     clip(rgbn[1] + (rgb[1] - value[1]) / 2.0),   214                     clip(rgbn[2] + (rgb[2] - value[2]) / 2.0)   215                     )   216                 im.putpixel((x, y+1), rgbn)   217    218     pim.putdata(im.getdata())   219    220 def get_float(options, flag):   221     try:   222         i = options.index(flag)   223         if i+1 < len(options) and options[i+1].isdigit():   224             return float(options[i+1])   225         else:   226             return 1.0   227     except ValueError:   228         return 0.0   229    230 class SimpleImage:   231    232     "An image behaving like PIL.Image."   233    234     def __init__(self, data, size):   235         self.data = data   236         self.width, self.height = self.size = size   237    238     def copy(self):   239         return SimpleImage(self.data[:], self.size)   240    241     def getpixel(self, xy):   242         x, y = xy   243         return self.data[y * self.width + x]   244    245     def putpixel(self, xy, value):   246         x, y = xy   247         self.data[y * self.width + x] = value   248    249     def getdata(self):   250         return self.data   251    252 # Main program.   253    254 if __name__ == "__main__":   255    256     # Test options.   257    258     if "--test" in sys.argv:   259         test()   260         sys.exit(0)   261     elif "--test-flat" in sys.argv:   262         test_flat((120, 40, 60))   263         sys.exit(0)   264     elif "--help" in sys.argv:   265         print >>sys.stderr, """\   266 Usage: %s <input filename> <output filename> [ <options> ]   267    268 Options are...   269    270 -s - Saturate the input image (can be followed by a float, default 1.0)   271 -d - Desaturate the input image (can be followed by a float, default 1.0)   272 -D - Darken the input image (can be followed by a float, default 1.0)   273 -B - Brighten the input image (can be followed by a float, default 1.0)   274    275 -r - Rotate the input image clockwise   276 -p - Generate a separate preview image   277 -h - Make the preview image with half horizontal resolution (MODE 2)   278 -v - Verify the output image (loaded if -n is given)   279 -n - Generate no output image   280 """ % split(sys.argv[0])[1]   281         sys.exit(1)   282    283     width = 320   284     height = 256   285    286     input_filename, output_filename = sys.argv[1:3]   287     basename, ext = splitext(output_filename)   288     preview_filename = "".join([basename + "_preview", ext])   289    290     options = sys.argv[3:]   291    292     # Preprocessing options that can be repeated for extra effect.   293    294     saturate = get_float(options, "-s")   295     desaturate = get_float(options, "-d")   296     darken = get_float(options, "-D")   297     brighten = get_float(options, "-B")   298    299     # General output options.   300    301     rotate = "-r" in options   302     preview = "-p" in options   303     half_resolution_preview = "-h" in options   304     verify = "-v" in options   305     no_normal_output = "-n" in options   306     make_image = not no_normal_output   307    308     # Load the input image if requested.   309    310     if make_image or preview:   311         exif = EXIF.process_file(open(input_filename))   312         im = PIL.Image.open(input_filename).convert("RGB")   313         im = rotate_and_scale(exif, im, width, height, rotate)   314    315         process_image(im, saturate, desaturate, darken, brighten)   316    317     # Generate a preview if requested.   318    319     if preview:   320         preview_image(im, half_resolution_preview).save(preview_filename)   321    322     # Generate an output image if requested.   323    324     if make_image:   325         convert_image(im)   326         im.save(output_filename)   327    328     # Verify the output image (which may be loaded) if requested.   329    330     if verify:   331         if no_normal_output:   332             im = PIL.Image.open(output_filename).convert("RGB")   333    334         result = count_colours(im, 4)   335         if result is not None:   336             y, colours = result   337             print "Image %s: row %d has the following colours: %s" % (output_filename, y, "; ".join([repr(c) for c in colours]))   338    339 # vim: tabstop=4 expandtab shiftwidth=4