First public commit.

lenser/a7800/convert/c160.py

@@ -0,0 +1,127 @@
+"""Atari 7800 MARIA 160A colour mode -- 160x192, 4 px/byte (2bpp).
+
+MARIA gives 8 palettes of 3 colours each plus one shared background = up to 25
+colours on screen.  The line is split into 8 objects of 20 px (5 bytes) each, and
+every object may use a different palette, so each 20 px segment can pick the
+3-colour palette (+ shared background) that fits it best -- chosen globally by
+clustering the segments into 8 palettes.  The 2bpp bitmap packing is identical to
+Atari GR.15, so the Atari encoder helpers are reused.
+
+Conversion.data = bitmap(7680) + seg_palettes(192*8) + colours(25):
+    bitmap        192 lines x 40 bytes, 2bpp, value 0..3 per pixel
+    seg_palettes  one palette index (0..7) per 20px segment (8 per line)
+    colours       [BACKGRND, P0C1,P0C2,P0C3, P1C1..P7C3]  (MARIA register order)
+"""
+from __future__ import annotations
+
+import numpy as np
+
+from ... import palette as c64pal
+from ...convert.base import Conversion, perceptual_error, DIFFUSION_DITHERS
+from ...atari import palette as apal
+from ...atari.convert import _common
+
+WIDTH, HEIGHT = 160, 192
+PIXEL_ASPECT = 2.0
+SEG_W = 40                       # pixels per object (10 bytes, 2bpp)
+N_SEG = WIDTH // SEG_W           # 4 objects per line (MARIA DMA budget)
+N_PAL = 8                        # MARIA palettes
+
+
+def _pack_bitmap(val_image):
+    return b"".join(bytes(b) for b in _common.pack_2bpp(val_image))
+
+
+def convert(img_rgb, palette_name="ntsc", dither_mode="floyd",
+            intensive=False, base_color=None, candidates=None, _mode="c160"):
+    plab = apal.palette_lab("ntsc")
+    prgb = apal.get_palette("ntsc").astype(np.uint8)
+    img_lab = c64pal.srgb_to_lab(img_rgb)
+
+    bg, palettes, seg_pal, idx = _solve(img_lab, plab, dither_mode, intensive,
+                                        candidates)
+
+    # value image (0..3) per pixel: 0 = background, 1..3 = the segment palette's
+    # three colours; build it from idx (palette indices) + the per-segment palette.
+    val = np.zeros((HEIGHT, WIDTH), np.uint8)
+    for row in range(HEIGHT):
+        for s in range(N_SEG):
+            x0 = s * SEG_W
+            pal = [bg] + palettes[seg_pal[row, s]]
+            lut = {c: v for v, c in enumerate(pal)}
+            block = idx[row, x0:x0 + SEG_W]
+            val[row, x0:x0 + SEG_W] = [lut[int(c)] for c in block]
+
+    bitmap = _pack_bitmap(val)
+    seg_bytes = bytes(seg_pal.reshape(-1).astype(np.uint8))
+    colours = bytearray([bg])
+    for p in range(N_PAL):
+        colours += bytes(palettes[p])
+    data = bitmap + seg_bytes + bytes(colours)
+
+    err = perceptual_error(idx, img_lab, plab)
+    return Conversion(
+        mode=_mode, width=WIDTH, height=HEIGHT, pixel_aspect=PIXEL_ASPECT,
+        index_image=idx.astype(np.uint16), data=data, data_addr=0,
+        viewer="a7800", preview_rgb=np.repeat(prgb[idx], int(PIXEL_ASPECT), axis=1),
+        error=err, meta={"palette": "ntsc", "dither": dither_mode, "bg": bg},
+    )
+
+
+def _pick(pixels_lab, plab, k, diff, candidates):
+    """Pick k palette colours best representing the given pixels."""
+    if diff:
+        return _common.choose_palette_dither(pixels_lab, plab, k=k,
+                                             candidates=candidates)
+    if candidates is not None:
+        sub = plab[candidates]
+        return [candidates[i] for i in _common.choose_palette(pixels_lab, sub, k=k)]
+    return _common.choose_palette(pixels_lab, plab, k=k)
+
+
+def _solve(img_lab, plab, dither_mode, intensive, candidates=None):
+    """Choose a shared background + 8 three-colour palettes by clustering the
+    image's 40px segments by colour (so each palette is tuned to a group of
+    similar segments), assign each segment its cluster's palette, and dither.
+    Returns (bg, palettes[8][3], seg_pal, idx)."""
+    from ... import dither as _dith
+    diff = dither_mode in DIFFUSION_DITHERS
+    H, W, _ = img_lab.shape
+
+    # shared background = darkest of a small global palette
+    gp = _pick(img_lab, plab, 8, diff, candidates)
+    bg = min(gp, key=lambda c: plab[c, 0])
+
+    # cluster the H*N_SEG segments by their mean colour into N_PAL groups
+    seg_means = img_lab.reshape(H, N_SEG, SEG_W, 3).mean(axis=2)      # (H,N_SEG,3)
+    feats = seg_means.reshape(-1, 3)
+    rng = np.random.default_rng(0)
+    cent = feats[rng.choice(len(feats), N_PAL, replace=False)].copy()
+    labels = np.zeros(len(feats), np.int64)
+    for _ in range(12):
+        labels = ((feats[:, None] - cent[None]) ** 2).sum(-1).argmin(1)
+        for g in range(N_PAL):
+            m = labels == g
+            if m.any():
+                cent[g] = feats[m].mean(0)
+    seg_pal = labels.reshape(H, N_SEG)
+
+    # each palette = 3 colours tuned to the pixels of its cluster's segments
+    seg_px = img_lab.reshape(H, N_SEG, SEG_W, 3)
+    palettes = []
+    for g in range(N_PAL):
+        mask = seg_pal == g                                          # (H,N_SEG)
+        if not mask.any():
+            palettes.append([bg, bg, bg])
+            continue
+        px = seg_px[mask].reshape(-1, 1, 3)                          # (Npix,1,3)
+        palettes.append(_pick(px, plab, 3, diff, candidates))
+
+    # dither the whole image with each segment restricted to {bg} + its palette
+    sets = [[bg] + palettes[g] for g in range(N_PAL)]
+    allowed = np.zeros((H, W, 4), np.int64)
+    for r in range(H):
+        for s in range(N_SEG):
+            allowed[r, s * SEG_W:s * SEG_W + SEG_W] = sets[seg_pal[r, s]]
+    idx = _dith.quantize(img_lab, allowed, plab, dither_mode).astype(np.int64)
+    return bg, palettes, seg_pal, idx