8bitlenser/lenser/a2600/convert/pf.py
2026-07-03 19:35:35 -07:00

144 lines
6.6 KiB
Python

"""Atari 2600 playfield image: 40x192, 3 colours per scanline.
No framebuffer -- the picture is a table the "racing the beam" kernel feeds to
the TIA one scanline at a time: a 40-pixel asymmetric playfield (left 20 + right
20) plus a shared playfield colour (COLUPF) and TWO background colours -- COLUBK
is rewritten mid-line, so the left and right 20px halves each get their own
background. Per line we jointly pick the shared foreground and the two
backgrounds that minimise dithered error, then dither each half between its two
colours.
"""
from __future__ import annotations
import numpy as np
from ... import dither
from ...convert.base import Conversion, perceptual_error, _box_blur
from ...palette import srgb_to_lab
from .. import palette as tpal
WIDTH, HEIGHT = 40, 192
HALF = WIDTH // 2
PIXEL_ASPECT = (4 / 3) / (WIDTH / HEIGHT) # very wide pixels
def _best_triples(img_lab, plab, split_gain=0.12, cand=None):
"""Per scanline choose a shared foreground + a background for each 20px half.
For a fixed foreground f, each half's best error is min over its background
bg of sum_px min(d[px,f], d[px,bg]); we pick the f minimising left+right.
To avoid a gratuitous vertical seam down the centre, the two halves keep a
UNIFIED background unless splitting it reduces the line's error by more than
``split_gain`` (so a seam only appears where the image really differs L/R).
``cand`` (palette indices) restricts the usable colours, e.g. to the greys.
Returns (fg, bgL, bgR) index arrays."""
fg = np.zeros(HEIGHT, np.int64)
bgL = np.zeros(HEIGHT, np.int64)
bgR = np.zeros(HEIGHT, np.int64)
forbid = None
if cand is not None:
forbid = np.full(plab.shape[0], np.inf)
forbid[np.asarray(cand)] = 0.0
for y in range(HEIGHT):
d = np.sum((img_lab[y][:, None, :] - plab[None, :, :]) ** 2, axis=-1) # (40,P)
if forbid is not None:
d = d + forbid[None, :] # forbid non-candidate colours
dl = d[:HALF].T # (P, 20)
dr = d[HALF:].T
# G?[f,bg] = sum_px min(d[f,px], d[bg,px])
GL = np.minimum(dl[:, None, :], dl[None, :, :]).sum(-1) # (P,P)
GR = np.minimum(dr[:, None, :], dr[None, :, :]).sum(-1)
total = GL.min(1) + GR.min(1)
f = int(total.argmin())
fg[y] = f
split_err = total[f]
uni = GL[f] + GR[f] # same bg for both halves
bg_uni = int(uni.argmin())
uni_err = float(uni[bg_uni])
if split_err < uni_err * (1.0 - split_gain):
bgL[y] = int(GL[f].argmin())
bgR[y] = int(GR[f].argmin())
else:
bgL[y] = bgR[y] = bg_uni # unified -> no seam on this line
return fg, bgL, bgR
def _encode_frame(img_rgb, plab, dither_mode, cand=None):
"""Encode one frame -> (9-table data bytes, idx image (192x40 palette indices))."""
img_lab = srgb_to_lab(img_rgb)
fg, bgL, bgR = _best_triples(img_lab, plab, cand=cand)
allowed = np.empty((HEIGHT, WIDTH, 2), np.int64)
allowed[:, :HALF, 0] = bgL[:, None]; allowed[:, :HALF, 1] = fg[:, None]
allowed[:, HALF:, 0] = bgR[:, None]; allowed[:, HALF:, 1] = fg[:, None]
idx = dither.quantize(img_lab, allowed, plab, dither_mode).astype(np.int64)
bit = (idx == fg[:, None]).astype(np.uint8) # 1 -> playfield
pf0L = np.zeros(HEIGHT, np.uint8); pf1L = np.zeros(HEIGHT, np.uint8)
pf2L = np.zeros(HEIGHT, np.uint8); pf0R = np.zeros(HEIGHT, np.uint8)
pf1R = np.zeros(HEIGHT, np.uint8); pf2R = np.zeros(HEIGHT, np.uint8)
colubkL = np.zeros(HEIGHT, np.uint8); colubkR = np.zeros(HEIGHT, np.uint8)
colupf = np.zeros(HEIGHT, np.uint8)
for y in range(HEIGHT):
pf0L[y], pf1L[y], pf2L[y] = tpal.pack20(bit[y, :HALF])
pf0R[y], pf1R[y], pf2R[y] = tpal.pack20(bit[y, HALF:])
colubkL[y] = tpal.color_byte(int(bgL[y]))
colubkR[y] = tpal.color_byte(int(bgR[y]))
colupf[y] = tpal.color_byte(int(fg[y]))
data = b"".join(bytes(t) for t in
(pf0L, pf1L, pf2L, pf0R, pf1R, pf2R, colubkL, colubkR, colupf))
return data, idx
def _widen(rgb):
disp_w = int(round(WIDTH * PIXEL_ASPECT))
xs = (np.arange(disp_w) * WIDTH) // disp_w
return rgb[:, xs]
def convert(img_rgb, palette_name="tia", dither_mode="floyd",
intensive=False, base_color=None, interlace=False, gray=False):
plab = tpal.palette_lab()
pal = tpal.PALETTE.astype(np.float64)
img_lab = srgb_to_lab(img_rgb)
# mono mode restricts the TIA to one hue's 8 luminances (hue 0 = greys, or
# the hue of base_color for a tinted mono); colour index = hue*8 + lum.
cand = None
if gray:
hue = 0 if base_color is None else (int(base_color) // 8)
cand = list(range(hue * 8, hue * 8 + 8))
if interlace:
# Frame A approximates the image; frame B is encoded so the per-pixel
# blend (averaged in linear light, the way 60Hz flicker is perceived)
# lands on the target -- giving ~4-6 perceived colours per scanline.
from ...palette import srgb_to_linear, linear_to_srgb
dataA, idxA = _encode_frame(img_rgb, plab, dither_mode, cand=cand)
renA = pal[idxA]
linT, linA = srgb_to_linear(img_rgb.astype(np.float64)), srgb_to_linear(renA)
targetB = linear_to_srgb(np.clip(2 * linT - linA, 0, 1))
dataB, idxB = _encode_frame(targetB, plab, dither_mode, cand=cand)
renB = pal[idxB]
blend = np.clip(linear_to_srgb((srgb_to_linear(renA) +
srgb_to_linear(renB)) / 2), 0, 255)
blend_lab = srgb_to_lab(blend)
# perceptual (blurred) error -- credits the 60Hz flicker blend the eye averages
err = float(np.sqrt(((_box_blur(blend_lab) - _box_blur(img_lab)) ** 2)
.sum(-1)).mean())
return Conversion(
mode="pf_il", width=WIDTH, height=HEIGHT, pixel_aspect=PIXEL_ASPECT,
index_image=idxA.astype(np.uint16), data=dataA + dataB, data_addr=0,
viewer="a2600", preview_rgb=_widen(blend.astype(np.uint8)),
error=err, meta={"palette": "tia", "dither": dither_mode, "interlace": True},
)
data, idx = _encode_frame(img_rgb, plab, dither_mode, cand=cand)
return Conversion(
mode="mono" if gray else "pf", width=WIDTH, height=HEIGHT,
pixel_aspect=PIXEL_ASPECT,
index_image=idx.astype(np.uint16), data=data, data_addr=0,
viewer="a2600", preview_rgb=_widen(tpal.PALETTE.astype(np.uint8)[idx]),
error=perceptual_error(idx, img_lab, plab),
meta={"palette": "tia", "dither": dither_mode},
)