PaletteOptimiser

Change of optimiser.py

20:99d5a002b1ce
optimiser.py
     1.1 --- a/optimiser.py	Fri Sep 18 13:55:41 2015 +0200
     1.2 +++ b/optimiser.py	Wed Sep 30 18:19:47 2015 +0200
     1.3 @@ -1,71 +1,89 @@
     1.4  #!/usr/bin/env python
     1.5  
     1.6 -from array import array
     1.7  from itertools import combinations
     1.8 +from random import randint
     1.9  import EXIF
    1.10  import PIL.Image
    1.11 +import math
    1.12  import sys
    1.13  
    1.14 -def scale(v):
    1.15 -    return (v + 43) / 85
    1.16 -
    1.17 -def point(rgb):
    1.18 -    return tuple(map(scale, rgb))
    1.19 -
    1.20 -def index(p):
    1.21 -    return p[0] * 16 + p[1] * 4 + p[2]
    1.22 -
    1.23 -def colour(i):
    1.24 -    return (255 * (i % 2), 255 * ((i / 2) % 2), 255 * ((i / 4) % 2))
    1.25 -
    1.26 -def add(d, v):
    1.27 -    d[v] = (d.has_key(v) and d[v] or 0) + 1
    1.28 -
    1.29 -def by_frequency(d):
    1.30 -    l = [(f, t) for (t, f) in d.items()]
    1.31 -    l.sort(reverse=True)
    1.32 -    return [i[1] for i in l]
    1.33 -
    1.34 -def match(b, bases):
    1.35 -    return b in bases and b
    1.36 -
    1.37 -def fallback(bases):
    1.38 -    for b in by_frequency(bases):
    1.39 -        if b not in ["_", "W"]:
    1.40 -            return b
    1.41 -    return by_frequency(bases)[0]
    1.42 -
    1.43 -tones = [
    1.44 -    "___", "_BB", "_BB", "BBB", # 00x
    1.45 -    "_GG", "__C", "_BC", "BCC", # 01x
    1.46 -    "_GG", "GGC", "BCW", "CCC", # 02x
    1.47 -    "GGG", "GGC", "GCW", "CCW", # 03x
    1.48 -    "_RR", "_MM", "BMM", "MBB", # 10x
    1.49 -    "_YY", "_**", "**B", "BBC", # 11x
    1.50 -    "_GY", "GGC", "*CC", "BCW", # 12x
    1.51 -    "BGY", "GGC", "CCW", "CCW", # 13x
    1.52 -    "_RR", "_RM", "*MM", "RMM", # 20x
    1.53 -    "RYY", "*RY", "RMM", "MMM", # 21x
    1.54 -    "YYY", "YYW", "*WW", "*WW", # 22x
    1.55 -    "YYY", "GYY", "GWW", "CWW", # 23x
    1.56 -    "RRR", "RRM", "BMR", "BMR", # 30x
    1.57 -    "RRY", "RRY", "RMW", "RMW", # 31x
    1.58 -    "YYY", "YYW", "YYW", "WWW", # 32x
    1.59 -    "YYY", "YYW", "YYW", "WWW", # 33x
    1.60 +corners = [
    1.61 +    (0, 0, 0), (255, 0, 0), (0, 255, 0), (255, 255, 0),
    1.62 +    (0, 0, 255), (255, 0, 255), (0, 255, 255), (255, 255, 255)
    1.63      ]
    1.64  
    1.65 -colours = ["_", "R", "G", "Y", "B", "M", "C", "W"]
    1.66 +def distance(rgb1, rgb2):
    1.67 +    r1, g1, b1 = rgb1
    1.68 +    r2, g2, b2 = rgb2
    1.69 +    return math.sqrt(pow(r1 - r2, 2) + pow(g1 - g2, 2) + pow(b1 - b2, 2))
    1.70 +
    1.71 +def brightness(rgb):
    1.72 +    return distance(rgb, (0, 0, 0))
    1.73  
    1.74 -if __name__ == "__main__":
    1.75 -    width = 320
    1.76 -    height = 256
    1.77 +def factor(start, end, rgb):
    1.78 +    r1, g1, b1 = start
    1.79 +    r2, g2, b2 = end
    1.80 +    gr, gg, gb = r2 - r1, g2 - g1, b2 - b1
    1.81 +    r, g, b = rgb
    1.82 +    pr, pg, pb = r - r1, g - g1, b - b1
    1.83 +    dp = pr * gr + pg * gg + pb * gb
    1.84 +    return dp / pow(distance(start, end), 2)
    1.85 +
    1.86 +def darklight(rgb1, rgb2):
    1.87 +    if brightness(rgb1) <= brightness(rgb2):
    1.88 +        return rgb1, rgb2
    1.89 +    else:
    1.90 +        return rgb2, rgb1
    1.91 +
    1.92 +def nearest(rgb, values=None):
    1.93 +    l = map(lambda c: (distance(rgb, c), c), values or corners)
    1.94 +    l.sort()
    1.95 +    return l
    1.96  
    1.97 -    input_filename, output_filename = sys.argv[1:3]
    1.98 -    rotate = "-r" in sys.argv[3:]
    1.99 +def pattern(rgb):
   1.100 +    l = nearest(rgb)
   1.101 +    start, end = l[0][1], l[1][1]
   1.102 +    f = factor(start, end, rgb)
   1.103 +    #if f > 0.5:
   1.104 +    #    start, end = end, start
   1.105 +    #    f = 1 - f
   1.106 +    return start, end, f
   1.107 +
   1.108 +def choose(seq, f):
   1.109 +    last = int(seq * math.sqrt(f))
   1.110 +    current = int((seq + 1) * math.sqrt(f))
   1.111 +    return last != current
   1.112 +
   1.113 +def get_value(xy, rgb, width, height):
   1.114 +    x, y = xy
   1.115 +    rgb1, rgb2, f = pattern(rgb)
   1.116 +    if choose(x + randint(0, width), f) and choose(y + randint(0, height), f):
   1.117 +        return rgb2
   1.118 +    else:
   1.119 +        return rgb1
   1.120  
   1.121 -    x = EXIF.process_file(open(input_filename))
   1.122 -    im = PIL.Image.open(input_filename)
   1.123 +def get_best(rgb, values):
   1.124 +    return nearest(rgb, values)[0][1]
   1.125  
   1.126 +def test():
   1.127 +    size = 512
   1.128 +    for r in (0, 63, 127, 191, 255):
   1.129 +        im = PIL.Image.new("RGB", (size, size))
   1.130 +        for g in range(0, size):
   1.131 +            for b in range(0, size):
   1.132 +                value = get_value((g, b), (r, (g * 256) / size, (b * 256 / size)), size, size)
   1.133 +                im.putpixel((g, b), value)
   1.134 +        im.save("rgb%d.png" % r)
   1.135 +
   1.136 +def test_flat(rgb):
   1.137 +    size = 64
   1.138 +    im = PIL.Image.new("RGB", (size, size))
   1.139 +    for y in range(0, size):
   1.140 +        for x in range(0, size):
   1.141 +            im.putpixel((x, y), get_value((x, y), rgb, size, size))
   1.142 +    im.save("rgb%02d%02d%02d.png" % rgb)
   1.143 +
   1.144 +def rotate_and_scale(im, width, height, rotate):
   1.145      if rotate or x and x["Image Orientation"].values == [6L]:
   1.146          im = im.rotate(270)
   1.147  
   1.148 @@ -75,79 +93,55 @@
   1.149      else:
   1.150          width = (height * w) / h
   1.151  
   1.152 -    im = im.resize((width, height))
   1.153 +    return im.resize((width, height))
   1.154  
   1.155 -    usage = []
   1.156 -    base_usage = []
   1.157 -    toned = []
   1.158 +if __name__ == "__main__":
   1.159 +    if "--test" in sys.argv:
   1.160 +        test()
   1.161 +        sys.exit(0)
   1.162 +    elif "--test-flat" in sys.argv:
   1.163 +        test_flat((120, 40, 60))
   1.164 +        sys.exit(0)
   1.165  
   1.166 -    for row in range(0, height):
   1.167 -        u = {}
   1.168 -        usage.append(u)
   1.169 -        bu = {}
   1.170 -        base_usage.append(bu)
   1.171 -        tr = []
   1.172 -        toned.append(tr)
   1.173 -        for column in range(0, width):
   1.174 -            rgb = im.getpixel((column, row))
   1.175 -            p = point(rgb)
   1.176 -            i = index(p)
   1.177 -            t = tones[i]
   1.178 -            add(u, t)
   1.179 -            if t[0] != "*":
   1.180 -                add(bu, t[0])
   1.181 -            if t[1] != "*":
   1.182 -                add(bu, t[1])
   1.183 -            if t[2] != "*":
   1.184 -                add(bu, t[2])
   1.185 -            tr.append(t)
   1.186 +    width = 320
   1.187 +    height = 256
   1.188 +
   1.189 +    input_filename, output_filename = sys.argv[1:3]
   1.190 +    rotate = "-r" in sys.argv[3:]
   1.191  
   1.192 -    chosen = []
   1.193 +    x = EXIF.process_file(open(input_filename))
   1.194 +    im = PIL.Image.open(input_filename)
   1.195 +    im = rotate_and_scale(im, width, height, rotate)
   1.196  
   1.197 -    for row, (u, bu) in enumerate(zip(usage, base_usage)):
   1.198 -        light = row % 2
   1.199 -        best = 0
   1.200 -        best_bases = None
   1.201 -        best_missing = None
   1.202 -        best_map = None
   1.203 +    width, height = im.size
   1.204 +
   1.205 +    colours = []
   1.206  
   1.207 -        for bases in combinations(bu, min(len(bu), 4)):
   1.208 -            bases = dict([(base, bu[base]) for base in bases])
   1.209 -            count = 0
   1.210 -            missing = []
   1.211 -            tone_map = {}
   1.212 -            for tone, freq in u.items():
   1.213 -                base = match(tone[1], bases)
   1.214 -                if base:
   1.215 -                    tone_map[tone] = base
   1.216 -                    count += freq
   1.217 -                else:
   1.218 -                    base = light and match(tone[2], bases) or match(tone[0], bases) or match(tone[2], bases)
   1.219 -                    if base:
   1.220 -                        tone_map[tone] = base
   1.221 -                        count += freq / 2
   1.222 -                    else:
   1.223 -                        missing.append(tone)
   1.224 -            if count > best:
   1.225 -                best_bases = bases
   1.226 -                best_missing = missing
   1.227 -                best_map = tone_map
   1.228 -                best = count
   1.229 +    for y in range(0, height):
   1.230 +        c = {}
   1.231 +        for x in range(0, width):
   1.232 +            rgb = im.getpixel((x, y))
   1.233 +            value = get_value((x, y), rgb, width, height)
   1.234 +            im.putpixel((x, y), value)
   1.235 +            if not c.has_key(value):
   1.236 +                c[value] = 1
   1.237 +            else:
   1.238 +                c[value] += 1
   1.239 +        c = [(n, value) for value, n in c.items()]
   1.240 +        c.sort(reverse=True)
   1.241 +        colours.append(c)
   1.242  
   1.243 -        chosen.append((best, best_bases or bases, best_map or tone_map, best_missing or missing))
   1.244 -
   1.245 -    output = []
   1.246 +    for y, c in enumerate(colours):
   1.247 +        if len(c) <= 4:
   1.248 +            continue
   1.249 +        most = [value for n, value in c[:4]]
   1.250 +        least = [value for n, value in c[4:]]
   1.251  
   1.252 -    for row, (tr, ch) in enumerate(zip(toned, chosen)):
   1.253 -        o = []
   1.254 -        for column, t in enumerate(tr):
   1.255 -            best, bases, tone_map, missing = ch
   1.256 -            base = tone_map.get(t) or fallback(bases)
   1.257 -            o.append(base)
   1.258 -            i = colours.index(base)
   1.259 -            im.putpixel((column, row), colour(i))
   1.260 -
   1.261 -        output.append("".join(o))
   1.262 +        for x in range(0, width):
   1.263 +            rgb = im.getpixel((x, y))
   1.264 +            if rgb in least:
   1.265 +                value = get_best(rgb, most)
   1.266 +                im.putpixel((x, y), value)
   1.267  
   1.268      im.save(output_filename)
   1.269