84 lines
2.9 KiB
Python
84 lines
2.9 KiB
Python
from typing import Generator
|
|
from PIL import Image
|
|
from pathlib import Path
|
|
import numpy as np
|
|
import math
|
|
|
|
def load_palette(path: Path) -> np.ndarray:
|
|
palette = []
|
|
with open(path, 'r') as f:
|
|
for line in f:
|
|
hex_color = line.strip()
|
|
# Convert hex to RGB tuple
|
|
rgb_color = [int(hex_color[i:i+2], 16) for i in (0, 2, 4)]
|
|
palette.append(rgb_color)
|
|
return np.array(palette)
|
|
|
|
pico8_palette = load_palette('../pico-8-secret-palette.hex')
|
|
|
|
def load_image_128(path: Path) -> Image.Image:
|
|
img = Image.open(path).convert('RGB')
|
|
w, h = img.size
|
|
if w > 128 or h > 128:
|
|
if w % 128 == 0 and h % 128 == 0:
|
|
resample = Image.Resampling.NEAREST
|
|
else:
|
|
resample = Image.Resampling.BICUBIC
|
|
if w > h:
|
|
w, h = 128, int(round(h * (128 / w)))
|
|
else:
|
|
w, h = int(round(w * (128 / h))), 128
|
|
img = img.resize((w, h), resample=resample)
|
|
return img
|
|
|
|
|
|
def quantize_image(img: Image.Image) -> Image.Image:
|
|
pal_idxs = choose_subpalette_indexes(img, 16, pico8_palette)
|
|
pal = np.array([pico8_palette[i] for i in pal_idxs]).astype(np.float32)
|
|
img_np = np.array(img)
|
|
pixels = img_np.reshape(-1, 3).astype(np.float32)
|
|
distances = np.sum((pixels[:, None, :] - pal[None, :, :]) ** 2, axis=2)
|
|
closest_idxs = np.argmin(distances, axis=1).astype(np.uint8).reshape(img_np.shape[1], -1)
|
|
result = Image.fromarray(closest_idxs, mode="P")
|
|
result.putpalette(pal.astype(np.uint8).reshape(-1).tolist())
|
|
return result
|
|
|
|
|
|
def choose_subpalette_indexes(img: Image.Image, n: int, palette: np.ndarray):
|
|
assert len(palette.shape) == 2;
|
|
assert palette.shape[1] == 3
|
|
img_np = np.array(img)
|
|
result = list(range(n))
|
|
num_steps = 0
|
|
pal_idx = 0
|
|
while num_steps < n:
|
|
num_steps += 1
|
|
old_val = result[pal_idx]
|
|
result[pal_idx] = None
|
|
best_err = math.inf
|
|
best_val = None
|
|
for val in [i for i in range(len(palette)) if i not in result]:
|
|
result[pal_idx] = val
|
|
err = _squared_quantization_error(img_np, [palette[i] for i in result])
|
|
if err < best_err:
|
|
best_err = err
|
|
best_val = val
|
|
result[pal_idx] = best_val
|
|
if best_val != old_val:
|
|
num_steps = 0
|
|
pal_idx = (pal_idx + 1) % len(result)
|
|
return result
|
|
|
|
def _squared_quantization_error(img_np: np.ndarray, palette: np.ndarray):
|
|
pixels = img_np.reshape(-1, 3).astype(np.float32)
|
|
pal = np.array(palette, dtype=np.float32)
|
|
sq_distances = np.sum((pixels[:, None, :] - pal[None, :, :]) ** 2, axis=2)
|
|
min_errs = np.min(sq_distances, axis=1)
|
|
return np.sum(min_errs)
|
|
|
|
def walk_corpus(quantized=True) -> Generator[Path, None, None]:
|
|
if quantized:
|
|
return Path('../corpus/quantized/').glob('*.png')
|
|
else:
|
|
return Path('../corpus/').glob('*.*')
|