156 lines
4.7 KiB
Python
156 lines
4.7 KiB
Python
import math
|
|
import PIL
|
|
from PIL.Image import Image
|
|
from functools import lru_cache
|
|
|
|
RGB = tuple[int, int, int]
|
|
|
|
# Load image, palette
|
|
img: Image = PIL.Image.open('peacock-bayer.png')
|
|
|
|
# Get palette info
|
|
assert img.mode == 'P'
|
|
raw_pal = img.getpalette()
|
|
palette: list[RGB] = []
|
|
for i in range(0, len(raw_pal), 3):
|
|
palette.append(tuple(raw_pal[i:i+3]))
|
|
palette_indexes = {c: i for i, c in enumerate(palette)}
|
|
|
|
# Get darkest/lightest colors
|
|
@lru_cache
|
|
def brightness(c: RGB):
|
|
return c[0]*0.3 + c[1]*0.6 + c[2]*0.1
|
|
|
|
darkest_color = min(palette, key=brightness)
|
|
lightest_color = max(palette, key=brightness)
|
|
|
|
@lru_cache
|
|
def distance_sq(a: RGB, b: RGB) -> int:
|
|
result = 0
|
|
for i in range(3):
|
|
result += (a[i] - b[i])**2
|
|
return result
|
|
|
|
def nearest_color(color: RGB, colors: list[RGB]) -> RGB:
|
|
return min(colors, key=lambda c: distance_sq(c, color))
|
|
|
|
def choose_color_pair_kmeans(counts: dict[RGB,int]) -> tuple[RGB, RGB]:
|
|
colors = [k for k, v in counts.items() if v > 0]
|
|
total_count = sum(counts.values())
|
|
if len(colors) == 0:
|
|
raise ValueError('no colors present')
|
|
elif len(colors) == 1:
|
|
return colors[0], None
|
|
|
|
def weighted_avg(colors: list[RGB]) -> RGB:
|
|
result = [0, 0, 0]
|
|
total_weight = 0
|
|
for color in colors:
|
|
weight = counts[color]
|
|
total_weight += weight
|
|
for i in range(3):
|
|
result[i] += weight*color[i]
|
|
if total_weight == 0:
|
|
total_weight += 1 # Hack to handle division-by-zero
|
|
return tuple(round(x/total_weight) for x in result)
|
|
|
|
# k-means to choose the two colors
|
|
old_result = (None, None)
|
|
result = (min(colors, key=brightness), max(colors, key=brightness))
|
|
print(f"k-means using {counts=}")
|
|
print(f"starting value: {result}")
|
|
while True:
|
|
lists = ([], [])
|
|
for color in colors:
|
|
d_dark = distance_sq(color, result[0])
|
|
d_light = distance_sq(color, result[1])
|
|
if d_dark < d_light:
|
|
lists[0].append(color)
|
|
else:
|
|
lists[1].append(color)
|
|
old_result = result
|
|
print(f"{lists=}")
|
|
result = tuple(weighted_avg(cs) for cs in lists)
|
|
print(f"updated: {result}")
|
|
if old_result == result:
|
|
break
|
|
return tuple(nearest_color(x, palette) for x in result)
|
|
|
|
|
|
def choose_color_pair_luminance(counts: dict[RGB, int]) -> tuple[RGB,RGB]:
|
|
# Calculate brightness threshold
|
|
mean_brightness = 0
|
|
total_count = 0
|
|
for rgb, count in counts.items():
|
|
mean_brightness += brightness(rgb)*count
|
|
total_count += count
|
|
mean_brightness = mean_brightness/total_count
|
|
|
|
# Separate colors into two lists
|
|
lists = ([], [])
|
|
for color in counts.keys():
|
|
if brightness(color) <= mean_brightness:
|
|
lists[0].append(color)
|
|
else:
|
|
lists[1].append(color)
|
|
print(lists)
|
|
|
|
# Find least-bad approximations for each color
|
|
def choose_color(cs: list[RGB]) -> RGB:
|
|
best = None
|
|
best_cost = math.inf
|
|
for rgb in palette:
|
|
cost = sum(distance_sq(rgb, c) * counts[c] for c in cs)
|
|
if cost < best_cost:
|
|
best = rgb
|
|
best_cost = cost
|
|
return best
|
|
dark, light = tuple(choose_color(l) for l in lists)
|
|
if dark == light:
|
|
light = None
|
|
print(dark, light)
|
|
return dark, light
|
|
|
|
|
|
def two_tone(img: Image, xy1: tuple[int,int] = None, xy2: tuple[int,int] = None) -> None:
|
|
# choose the two colors
|
|
x1, y1 = xy1 or (0, 0)
|
|
x2, y2 = xy2 or (img.width, img.height)
|
|
x1 = max(x1, 0)
|
|
y1 = max(y1, 0)
|
|
x2 = min(x2, img.width)
|
|
y2 = min(y2, img.height)
|
|
counts = {}
|
|
for x in range(x1, x2):
|
|
for y in range(y1, y2):
|
|
color = palette[img.getpixel((x,y))]
|
|
if color not in counts:
|
|
counts[color] = 0
|
|
counts[color] += 1
|
|
#dark, light = choose_color_pair_luminance(counts)
|
|
print(f"choosing color pair for {xy1=}")
|
|
dark, light = choose_color_pair_kmeans(counts)
|
|
print(f"final: {dark=}, {light=}\n")
|
|
if light == None:
|
|
# Only one color, *ought* to be a no-op
|
|
light = dark
|
|
colors = [dark, light]
|
|
|
|
# change each pixel's color to nearest
|
|
for x in range(x1, x2):
|
|
for y in range(y1, y2):
|
|
color = palette[img.getpixel((x, y))]
|
|
replacement = nearest_color(color, colors)
|
|
img.putpixel((x, y), palette_indexes[replacement])
|
|
|
|
|
|
BLOCK_SIZE = 10
|
|
blocked = img.copy()
|
|
for x in range(0, img.width, BLOCK_SIZE):
|
|
for y in range(0, img.height, BLOCK_SIZE):
|
|
two_tone(blocked, (x, y), (x + BLOCK_SIZE, y + BLOCK_SIZE))
|
|
blocked.save("blocked.png")
|
|
for x in range(5):
|
|
for y in range(5):
|
|
blocked.putpixel((x, y), 0)
|