1.1 --- a/optimiser.py	Fri Oct 02 13:37:29 2015 +0200
     1.2 +++ b/optimiser.py	Fri Oct 02 19:22:17 2015 +0200
     1.3 @@ -1,9 +1,11 @@
     1.4  #!/usr/bin/env python
     1.5  
     1.6 +from dither import dither
     1.7  from random import random
     1.8 -from os.path import extsep, splitext
     1.9 +from os.path import splitext
    1.10  import EXIF
    1.11  import PIL.Image
    1.12 +import itertools
    1.13  import math
    1.14  import sys
    1.15  
    1.16 @@ -17,10 +19,10 @@
    1.17      r2, g2, b2 = rgb2
    1.18      return math.sqrt(pow(r1 - r2, 2) + pow(g1 - g2, 2) + pow(b1 - b2, 2))
    1.19  
    1.20 -def distribution(rgb, points):
    1.21 +def distribution(rgb, values=None):
    1.22      l = []
    1.23      total = 0
    1.24 -    for c in points:
    1.25 +    for c in values or corners:
    1.26          d = distance(rgb, c)
    1.27          if d == 0:
    1.28              return [(1, c)]
    1.29 @@ -51,6 +53,22 @@
    1.30  def saturate_value(x, exp):
    1.31      return int(127.5 + sign(x - 127.5) * 127.5 * pow(abs(x - 127.5) / 127.5, exp))
    1.32  
    1.33 +def get_best(im, width, y, colours):
    1.34 +    best = 0
    1.35 +    best_values = None
    1.36 +
    1.37 +    for values in itertools.combinations([value for n, value in colours], 4):
    1.38 +        current = 0
    1.39 +        for x in range(0, width):
    1.40 +            rgb = im.getpixel((x, y))
    1.41 +            for f, value in distribution(rgb, values):
    1.42 +                current += f
    1.43 +        if current >= best:
    1.44 +            best_values = values
    1.45 +            best = current
    1.46 +
    1.47 +    return best_values
    1.48 +
    1.49  def test():
    1.50      size = 512
    1.51      for r in (0, 63, 127, 191, 255):
    1.52 @@ -94,21 +112,23 @@
    1.53  
    1.54      input_filename, output_filename = sys.argv[1:3]
    1.55      basename, ext = splitext(output_filename)
    1.56 -    preview_filename = extsep.join([basename + "_preview", ext])
    1.57  
    1.58 -    preview = "-p" in sys.argv[3:]
    1.59      rotate = "-r" in sys.argv[3:]
    1.60      saturate = sys.argv[3:].count("-s")
    1.61      desaturate = sys.argv[3:].count("-d")
    1.62 +    best = "-b" in sys.argv[3:]
    1.63 +    dithering = "-d" in sys.argv[3:]
    1.64  
    1.65      x = EXIF.process_file(open(input_filename))
    1.66      im = PIL.Image.open(input_filename).convert("RGB")
    1.67      im = rotate_and_scale(im, width, height, rotate)
    1.68  
    1.69 -    if preview:
    1.70 -        im_preview = im.copy()
    1.71 +    width, height = im.size
    1.72  
    1.73 -    width, height = im.size
    1.74 +    im_original = im
    1.75 +
    1.76 +    if dithering:
    1.77 +        im = dither(im).convert("RGB")
    1.78  
    1.79      colours = []
    1.80  
    1.81 @@ -123,28 +143,25 @@
    1.82                  rgb = saturate_rgb(rgb, saturate and math.pow(0.5, saturate) or math.pow(2, desaturate))
    1.83                  im.putpixel((x, y), rgb)
    1.84  
    1.85 -            # Count the number of requested colours.
    1.86 +            # Sum the colour probabilities.
    1.87  
    1.88 -            value = get_value(rgb)
    1.89 -            if not c.has_key(value):
    1.90 -                c[value] = 1
    1.91 -            else:
    1.92 -                c[value] += 1
    1.93 -
    1.94 -            if preview:
    1.95 -                im_preview.putpixel((x, y), value)
    1.96 +            for f, value in distribution(rgb):
    1.97 +                if not c.has_key(value):
    1.98 +                    c[value] = f
    1.99 +                else:
   1.100 +                    c[value] += f
   1.101  
   1.102          c = [(n, value) for value, n in c.items()]
   1.103          c.sort(reverse=True)
   1.104          colours.append(c)
   1.105  
   1.106 -    if preview:
   1.107 -        im_preview.save(preview_filename)
   1.108 -
   1.109      for y, c in enumerate(colours):
   1.110          most = [value for n, value in c[:4]]
   1.111          least = [value for n, value in c[4:]]
   1.112  
   1.113 +        if least and best:
   1.114 +            most = get_best(im_original, width, y, c)
   1.115 +
   1.116          for x in range(0, width):
   1.117              rgb = im.getpixel((x, y))
   1.118  
   1.119 @@ -153,7 +170,9 @@
   1.120  
   1.121              value = get_value(rgb)
   1.122              if value in least:
   1.123 +                rgb = im_original.getpixel((x, y))
   1.124                  value = get_value(rgb, most)
   1.125 +
   1.126              im.putpixel((x, y), value)
   1.127  
   1.128      im.save(output_filename)