PaletteOptimiser

Change of optimiser.py

32:3a950efbc68f
optimiser.py
     1.1 --- a/optimiser.py	Fri Oct 02 13:52:42 2015 +0200
     1.2 +++ b/optimiser.py	Sat Oct 03 00:32:58 2015 +0200
     1.3 @@ -1,9 +1,10 @@
     1.4  #!/usr/bin/env python
     1.5  
     1.6  from random import random
     1.7 -from os.path import extsep, splitext
     1.8 +from os.path import splitext
     1.9  import EXIF
    1.10  import PIL.Image
    1.11 +import itertools
    1.12  import math
    1.13  import sys
    1.14  
    1.15 @@ -17,32 +18,30 @@
    1.16      r2, g2, b2 = rgb2
    1.17      return math.sqrt(pow(r1 - r2, 2) + pow(g1 - g2, 2) + pow(b1 - b2, 2))
    1.18  
    1.19 -def factor(start, end, rgb):
    1.20 -    r1, g1, b1 = start
    1.21 -    r2, g2, b2 = end
    1.22 -    gr, gg, gb = r2 - r1, g2 - g1, b2 - b1
    1.23 -    r, g, b = rgb
    1.24 -    pr, pg, pb = r - r1, g - g1, b - b1
    1.25 -    dp = pr * gr + pg * gg + pb * gb
    1.26 -    return float(dp) / pow(distance(start, end), 2)
    1.27 +def distribution(rgb, values=None):
    1.28 +    l = []
    1.29 +    total = 0
    1.30 +    for c in values or corners:
    1.31 +        d = distance(rgb, c)
    1.32 +        if d == 0:
    1.33 +            return [(1, c)]
    1.34 +        l.append((pow(d, -3), c))
    1.35 +        total += pow(d, -3)
    1.36 +    return [(d / total, c) for d, c in l]
    1.37  
    1.38 -def nearest(rgb, values=None):
    1.39 -    l = map(lambda c: (distance(rgb, c), c), values or corners)
    1.40 -    l.sort()
    1.41 +def pattern(rgb, values=None):
    1.42 +    l = distribution(rgb, values or corners)
    1.43 +    l.sort(reverse=True)
    1.44      return l
    1.45  
    1.46 -def pattern(rgb, values=None):
    1.47 -    l = nearest(rgb, values)
    1.48 -    start, end = l[0][1], l[1][1]
    1.49 -    f = factor(start, end, rgb)
    1.50 -    return start, end, f
    1.51 -
    1.52  def get_value(rgb, values=None):
    1.53 -    rgb1, rgb2, f = pattern(rgb, values)
    1.54 -    if random() < pow(f, 2):
    1.55 -        return rgb2
    1.56 -    else:
    1.57 -        return rgb1
    1.58 +    choose = random()
    1.59 +    threshold = 0
    1.60 +    for f, c in pattern(rgb, values):
    1.61 +        threshold += f
    1.62 +        if choose < threshold:
    1.63 +            return c
    1.64 +    return c
    1.65  
    1.66  def sign(x):
    1.67      return x >= 0 and 1 or -1
    1.68 @@ -96,20 +95,15 @@
    1.69  
    1.70      input_filename, output_filename = sys.argv[1:3]
    1.71      basename, ext = splitext(output_filename)
    1.72 -    preview_filename = extsep.join([basename + "_preview", ext])
    1.73  
    1.74 -    preview = "-p" in sys.argv[3:]
    1.75      rotate = "-r" in sys.argv[3:]
    1.76      saturate = sys.argv[3:].count("-s")
    1.77      desaturate = sys.argv[3:].count("-d")
    1.78  
    1.79      x = EXIF.process_file(open(input_filename))
    1.80 -    im = PIL.Image.open(input_filename)
    1.81 +    im = PIL.Image.open(input_filename).convert("RGB")
    1.82      im = rotate_and_scale(im, width, height, rotate)
    1.83  
    1.84 -    if preview:
    1.85 -        im_preview = im.copy()
    1.86 -
    1.87      width, height = im.size
    1.88  
    1.89      colours = []
    1.90 @@ -125,24 +119,18 @@
    1.91                  rgb = saturate_rgb(rgb, saturate and math.pow(0.5, saturate) or math.pow(2, desaturate))
    1.92                  im.putpixel((x, y), rgb)
    1.93  
    1.94 -            # Count the number of requested colours.
    1.95 +            # Sum the colour probabilities.
    1.96  
    1.97 -            value = get_value(rgb)
    1.98 -            if not c.has_key(value):
    1.99 -                c[value] = 1
   1.100 -            else:
   1.101 -                c[value] += 1
   1.102 -
   1.103 -            if preview:
   1.104 -                im_preview.putpixel((x, y), value)
   1.105 +            for f, value in distribution(rgb):
   1.106 +                if not c.has_key(value):
   1.107 +                    c[value] = f
   1.108 +                else:
   1.109 +                    c[value] += f
   1.110  
   1.111          c = [(n, value) for value, n in c.items()]
   1.112          c.sort(reverse=True)
   1.113          colours.append(c)
   1.114  
   1.115 -    if preview:
   1.116 -        im_preview.save(preview_filename)
   1.117 -
   1.118      for y, c in enumerate(colours):
   1.119          most = [value for n, value in c[:4]]
   1.120          least = [value for n, value in c[4:]]
   1.121 @@ -155,7 +143,9 @@
   1.122  
   1.123              value = get_value(rgb)
   1.124              if value in least:
   1.125 -                value = get_value(rgb, most)
   1.126 +                rgb = im.getpixel((x, y))
   1.127 +                value = pattern(value, most)[0][1]
   1.128 +
   1.129              im.putpixel((x, y), value)
   1.130  
   1.131      im.save(output_filename)