130 lines
5.2 KiB
Python
130 lines
5.2 KiB
Python
"""Apple IIGS Super Hi-Res encoder: 320x200, 16 palettes x 16 colours (of 4096),
|
|
one palette per scanline (the SCB).
|
|
|
|
Lines are clustered into 16 palette groups; each group's 16 colours come from a
|
|
k-means of its pixels (quantised to 12-bit); each line is then assigned its best
|
|
palette and dithered to that palette's 16 colours. Mono uses a single 16-grey
|
|
palette for all lines.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
from ... import dither, palette as c64pal
|
|
from ...convert.base import perceptual_error
|
|
from .. import palette as gs
|
|
|
|
W, H = 320, 200
|
|
NPAL = 16
|
|
PALSIZE = 16
|
|
# SHR RAM block ($2000-$9FFF) offsets, relative to the 32K block base ($2000):
|
|
PIX_OFF, SCB_OFF, PAL_OFF = 0x0000, 0x7D00, 0x7E00
|
|
BLOCK_LEN = 0x8000
|
|
|
|
|
|
def _kmeans(pts, k, iters=12, seed=0):
|
|
rng = np.random.default_rng(seed)
|
|
if len(pts) <= k:
|
|
c = np.zeros((k, pts.shape[1])); c[:len(pts)] = pts; return c
|
|
cen = pts[rng.choice(len(pts), k, replace=False)].copy()
|
|
for _ in range(iters):
|
|
lab = np.argmin(((pts[:, None, :] - cen[None]) ** 2).sum(2), 1)
|
|
for j in range(k):
|
|
m = pts[lab == j]
|
|
if len(m):
|
|
cen[j] = m.mean(0)
|
|
return cen
|
|
|
|
|
|
def _build_palettes(img_lab, img_rgb, mono, base_color):
|
|
"""Return (pal_rgb (NPAL,16,3) float 0-255, pal_q4 (NPAL,16,3) 4-bit,
|
|
line_pal (H,) palette index per scanline)."""
|
|
if mono:
|
|
rgb, q4 = gs.greys12(PALSIZE)
|
|
if base_color is not None:
|
|
# tint: black -> base hue -> white ramp
|
|
base = np.array(base_color, float)
|
|
t = np.linspace(0, 1, PALSIZE)[:, None]
|
|
ramp = np.where(t < 0.5, base * (t / 0.5),
|
|
base + (255 - base) * ((t - 0.5) / 0.5))
|
|
rgb, q4 = gs.quantize12(ramp)
|
|
pal_rgb = np.repeat(rgb[None], NPAL, 0)
|
|
pal_q4 = np.repeat(q4[None], NPAL, 0)
|
|
return pal_rgb, pal_q4, np.zeros(H, np.int64)
|
|
|
|
# cluster the 200 lines into NPAL groups by mean colour
|
|
line_mean = img_lab.mean(1) # (H,3)
|
|
gc = _kmeans(line_mean, NPAL)
|
|
grp = np.argmin(((line_mean[:, None, :] - gc[None]) ** 2).sum(2), 1)
|
|
|
|
pal_rgb = np.zeros((NPAL, PALSIZE, 3))
|
|
pal_q4 = np.zeros((NPAL, PALSIZE, 3), np.int64)
|
|
rng = np.random.default_rng(0)
|
|
for g in range(NPAL):
|
|
lines = np.where(grp == g)[0]
|
|
if len(lines) == 0:
|
|
lines = np.array([g % H])
|
|
pool = img_rgb[lines].reshape(-1, 3)
|
|
if len(pool) > 4000:
|
|
pool = pool[rng.choice(len(pool), 4000, replace=False)]
|
|
cen = _kmeans(c64pal.srgb_to_lab(pool.reshape(1, -1, 3))[0], PALSIZE)
|
|
# k-means in Lab -> map centroids back to RGB via nearest pool pixel
|
|
plab_pool = c64pal.srgb_to_lab(pool.reshape(1, -1, 3))[0]
|
|
rgb16 = np.zeros((PALSIZE, 3))
|
|
for i in range(PALSIZE):
|
|
rgb16[i] = pool[np.argmin(((plab_pool - cen[i]) ** 2).sum(1))]
|
|
rgb_q, q4 = gs.quantize12(rgb16)
|
|
pal_rgb[g], pal_q4[g] = rgb_q, q4
|
|
|
|
# assign each line its best palette (min nearest-colour error)
|
|
pal_lab = c64pal.srgb_to_lab(pal_rgb.reshape(1, -1, 3))[0].reshape(NPAL, PALSIZE, 3)
|
|
line_pal = np.zeros(H, np.int64)
|
|
for y in range(H):
|
|
row = img_lab[y] # (W,3)
|
|
best, berr = 0, np.inf
|
|
for g in range(NPAL):
|
|
e = np.min(((row[:, None, :] - pal_lab[g][None]) ** 2).sum(2), 1).sum()
|
|
if e < berr:
|
|
berr, best = e, g
|
|
line_pal[y] = best
|
|
return pal_rgb, pal_q4, line_pal
|
|
|
|
|
|
def encode(img_rgb, dither_mode, mono=False, base_color=None):
|
|
img_lab = c64pal.srgb_to_lab(img_rgb)
|
|
pal_rgb, pal_q4, line_pal = _build_palettes(img_lab, img_rgb, mono, base_color)
|
|
|
|
# combined colour table (NPAL*16) for the dither; per-pixel allowed = its
|
|
# line's 16 palette colours
|
|
rgb_all = pal_rgb.reshape(-1, 3)
|
|
plab_all = c64pal.srgb_to_lab(rgb_all.reshape(1, -1, 3))[0]
|
|
allowed = np.zeros((H, W, PALSIZE), np.int64)
|
|
for y in range(H):
|
|
base = line_pal[y] * PALSIZE
|
|
allowed[y, :, :] = np.arange(base, base + PALSIZE)
|
|
idx = dither.quantize(img_lab, allowed, plab_all, dither_mode).astype(np.int64)
|
|
nib = (idx - line_pal[:, None] * PALSIZE).astype(np.uint8) # 0-15 within palette
|
|
|
|
# ---- emit the 32K SHR block ($2000-$9FFF) ----
|
|
block = bytearray(BLOCK_LEN)
|
|
for y in range(H):
|
|
row = nib[y]
|
|
off = PIX_OFF + y * 160
|
|
for bx in range(160):
|
|
block[off + bx] = (int(row[bx * 2]) << 4) | (int(row[bx * 2 + 1]) & 0x0F)
|
|
block[SCB_OFF + y] = line_pal[y] & 0x0F # 320 mode, palette n
|
|
for p in range(NPAL):
|
|
for c in range(PALSIZE):
|
|
r4, g4, b4 = pal_q4[p, c]
|
|
block[PAL_OFF + p * 32 + c * 2:PAL_OFF + p * 32 + c * 2 + 2] = \
|
|
gs.color_word(int(r4), int(g4), int(b4))
|
|
|
|
preview = rgb_all[idx].astype(np.uint8)
|
|
if mono:
|
|
# measure against luminance (the greys have no hue), like other mono modes
|
|
lum = img_lab.copy(); lum[..., 1:] = 0.0
|
|
plum = plab_all.copy(); plum[:, 1:] = 0.0
|
|
err = perceptual_error(idx, lum, plum)
|
|
else:
|
|
err = perceptual_error(idx, img_lab, plab_all)
|
|
return bytes(block), preview, err
|