130 lines
5.5 KiB
Python
130 lines
5.5 KiB
Python
"""Amiga encoders: HAM6 (4096 colours), flat low-res (<=32 of 4096), greyscale.
|
|
|
|
HAM (Hold-And-Modify) is the showpiece: 6 bitplanes where the top 2 bits choose
|
|
whether a pixel is one of 16 base palette colours or modifies one R/G/B channel
|
|
of the pixel to its left -- giving up to 4096 colours on screen. We pick 16 base
|
|
colours, then walk each scanline left-to-right choosing per pixel the option
|
|
(set / modify R / modify G / modify B) closest to the target.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
from ... import dither, palette as c64pal
|
|
from ...convert.base import _box_blur
|
|
from .. import palette as apal
|
|
|
|
W, H = 320, 200
|
|
|
|
|
|
def _perr(final_rgb, img_rgb):
|
|
a = _box_blur(c64pal.srgb_to_lab(final_rgb.astype(np.float64)))
|
|
b = _box_blur(c64pal.srgb_to_lab(img_rgb.astype(np.float64)))
|
|
return float(np.sqrt(((a - b) ** 2).sum(-1)).mean())
|
|
|
|
|
|
def planar_split(codes, nplanes):
|
|
"""codes (H,W) -> nplanes contiguous bitplanes (40 bytes/line, MSB left)."""
|
|
out = bytearray()
|
|
for p in range(nplanes):
|
|
bit = ((codes >> p) & 1).astype(np.uint8)
|
|
out += np.packbits(bit, axis=1).reshape(-1).tobytes() # (H,40)
|
|
return bytes(out)
|
|
|
|
|
|
def _kmeans_keys(img_lab, plab, k, iters=12):
|
|
rng = np.random.default_rng(0)
|
|
flat = img_lab.reshape(-1, 3)
|
|
pts = flat[rng.choice(len(flat), min(6000, len(flat)), replace=False)]
|
|
# k-means++ style seeding so every centroid is a real, distinct colour
|
|
cen = pts[rng.integers(len(pts))][None].copy()
|
|
for _ in range(k - 1):
|
|
d = np.min(((pts[:, None, :] - cen[None]) ** 2).sum(2), 1)
|
|
cen = np.vstack([cen, pts[int(d.argmax())]]) # farthest-point seed
|
|
for _ in range(iters):
|
|
lab = np.argmin(((pts[:, None, :] - cen[None]) ** 2).sum(2), 1)
|
|
for j in range(k):
|
|
m = pts[lab == j]
|
|
cen[j] = m.mean(0) if len(m) else pts[rng.integers(len(pts))]
|
|
# snap each centroid to the nearest 4096 palette key
|
|
keys = [int(np.argmin(((plab - c) ** 2).sum(1))) for c in cen]
|
|
return keys
|
|
|
|
|
|
def ham_encode(img_rgb, dither_mode):
|
|
plab = apal.palette_lab() # (4096,3)
|
|
img_lab = c64pal.srgb_to_lab(img_rgb)
|
|
base_keys = _kmeans_keys(img_lab, plab, 16)
|
|
base_lab = plab[base_keys]
|
|
base_rgb4 = [(k >> 8 & 15, k >> 4 & 15, k & 15) for k in base_keys]
|
|
|
|
t4 = np.clip(np.rint(img_rgb / 17.0), 0, 15).astype(np.int64) # 4-bit targets
|
|
# base option (independent of the held colour) precomputed for all pixels
|
|
bd = ((img_lab[:, :, None, :] - base_lab[None, None]) ** 2).sum(3) # (H,W,16)
|
|
base_best = bd.argmin(2)
|
|
base_cost = bd.min(2)
|
|
|
|
codes = np.zeros((H, W), np.uint8)
|
|
final = np.zeros((H, W, 3), np.uint8)
|
|
P = plab
|
|
for y in range(H):
|
|
tl = img_lab[y]; t4y = t4[y]
|
|
bb = base_best[y]; bc = base_cost[y]
|
|
pr = pg = pb = 0 # hardware holds black at line start
|
|
for x in range(W):
|
|
t0, t1, t2 = tl[x, 0], tl[x, 1], tl[x, 2]
|
|
bi = int(bb[x]); best = bc[x]
|
|
ctrl = 0; data = bi; nr, ng, nb = base_rgb4[bi]
|
|
# force an absolute "set" on the first pixel so the line establishes a
|
|
# colour regardless of the held-colour start (an all-modify run from
|
|
# the wrong start would otherwise stay dark -- HAM left-edge bug).
|
|
if x > 0:
|
|
tr, tg, tb = int(t4y[x, 0]), int(t4y[x, 1]), int(t4y[x, 2])
|
|
for c_ctrl, c_data, kr, kg, kb in (
|
|
(2, tr, tr, pg, pb), (3, tg, pr, tg, pb), (1, tb, pr, pg, tb)):
|
|
pk = P[(kr << 8) | (kg << 4) | kb]
|
|
dr = pk[0] - t0; dg = pk[1] - t1; db = pk[2] - t2
|
|
cc = dr * dr + dg * dg + db * db
|
|
if cc < best:
|
|
best = cc; ctrl = c_ctrl; data = c_data; nr, ng, nb = kr, kg, kb
|
|
codes[y, x] = (ctrl << 4) | data
|
|
pr, pg, pb = nr, ng, nb
|
|
final[y, x, 0] = pr * 17; final[y, x, 1] = pg * 17; final[y, x, 2] = pb * 17
|
|
|
|
planes = planar_split(codes, 6)
|
|
colors = [apal.color_word(k) for k in base_keys] # 16 base registers
|
|
return planes, colors, final, _perr(final, img_rgb)
|
|
|
|
|
|
def flat_encode(img_rgb, n_colors, dither_mode, mono=False, base_color=None):
|
|
plab = apal.palette_lab()
|
|
prgb = apal.get_palette().astype(np.uint8)
|
|
img_lab = c64pal.srgb_to_lab(img_rgb)
|
|
nplanes = (n_colors - 1).bit_length() # 32->5, 16->4
|
|
|
|
if mono:
|
|
keys = list(apal.GREYS) # 16 greys
|
|
if base_color in range(4096):
|
|
keys = sorted({keys[0], int(base_color), keys[-1]}, key=lambda i: plab[i, 0])
|
|
keys = (keys + keys)[:n_colors]
|
|
work = np.zeros_like(img_lab); work[..., 0] = img_lab[..., 0]
|
|
pw = np.zeros_like(plab); pw[:, 0] = plab[:, 0]
|
|
else:
|
|
keys = _kmeans_keys(img_lab, plab, n_colors)
|
|
work, pw = img_lab, plab
|
|
|
|
allowed = np.tile(np.array(keys[:n_colors]), (H, W, 1))
|
|
qidx = dither.quantize(work, allowed, pw, dither_mode).astype(np.int64)
|
|
# map palette key -> pen index
|
|
lut = {k: i for i, k in enumerate(keys[:n_colors])}
|
|
pen = np.vectorize(lut.get)(qidx).astype(np.uint8)
|
|
|
|
planes = planar_split(pen, nplanes)
|
|
colors = [apal.color_word(k) for k in keys[:n_colors]]
|
|
final = prgb[qidx]
|
|
if mono: # measure greyscale against luminance
|
|
g = img_rgb.mean(2, keepdims=True).repeat(3, 2)
|
|
err = _perr(final, g.astype(np.uint8))
|
|
else:
|
|
err = _perr(final, img_rgb)
|
|
return planes, colors, final, err
|