A JPEG-style lossy codec generalised to N channels

klt_jpeg.py

"""KLT-JPEG: a JPEG-style lossy codec generalised to N channels.

JPEG's pipeline, with the one 3-channel-specific step (RGB->YCbCr) replaced by
its proper N-channel form (a learned KLT / PCA across channels):

  encode:  mean-centre -> PCA across C dims -> 8x8 block DCT per channel
           -> quantise (JPEG luminance table x quality) -> zlib the integers
  decode:  reverse.

So it does BOTH decorrelations: cross-channel (PCA, what VQ exploits) and
spatial (DCT, what per-plane JPEG exploits). Entropy stage is zlib instead of
JPEG's bespoke Huffman (same pragmatic shortcut PNG128 took).
"""
import struct, zlib, numpy as np
from scipy.fft import dctn, idctn

# Standard JPEG luminance quantisation table (the spatial-frequency weighting).
_Q8 = np.array([
    [16,11,10,16,24,40,51,61],[12,12,14,19,26,58,60,55],
    [14,13,16,24,40,57,69,56],[14,17,22,29,51,87,80,62],
    [18,22,37,56,68,109,103,77],[24,35,55,64,81,104,113,92],
    [49,64,78,87,103,121,120,101],[72,92,95,98,112,100,103,99]], np.float32)
MAGIC = b"KJPG\r\n\x1a\n"


def _qtable(quality: float) -> np.ndarray:
    """JPEG quality (1-100) -> 8x8 quantisation step matrix."""
    s = 5000.0 / quality if quality < 50 else 200.0 - 2.0 * quality
    return np.maximum(np.floor((_Q8 * s + 50) / 100), 1).astype(np.float32)


def _blocks(a):                              # (Hp,Wp,C) -> (Hb,8,Wb,8,C)
    hp, wp, c = a.shape
    return a.reshape(hp // 8, 8, wp // 8, 8, c)


def encode(tile: np.ndarray, quality: float = 50.0, n_components: int | None = None) -> bytes:
    """tile: (H,W,C) float32 -> KLT-JPEG bytes."""
    h, w, c = tile.shape
    x = tile.reshape(-1, c).astype(np.float32)
    mean = x.mean(0)
    xc = x - mean
    # KLT across channels: principal axes from the (subsampled) covariance.
    samp = xc if xc.shape[0] <= 100_000 else xc[np.linspace(0, xc.shape[0] - 1, 100_000).astype(int)]
    _, _, vt = np.linalg.svd(samp, full_matrices=False)   # vt: (C,C) rows = components
    nc = n_components or c
    comp = vt[:nc].astype(np.float32)                     # (nc, C)
    coords = (xc @ comp.T).reshape(h, w, nc)              # (H,W,nc) decorrelated channels

    # the JPEG quant table is calibrated for 0-255 pixels; our coords are O(1),
    # so rescale into a comparable range (store the gain to undo on decode).
    gain = 64.0 / (float(coords.std()) + 1e-12)
    coords = coords * gain

    # pad spatial dims to multiples of 8 (edge-replicate) and block-DCT per channel
    ph, pw = (-h) % 8, (-w) % 8
    padded = np.pad(coords, ((0, ph), (0, pw), (0, 0)), mode="edge")
    dct = dctn(_blocks(padded), axes=(1, 3), norm="ortho")

    qt = _qtable(quality)[None, :, None, :, None]         # broadcast over (Hb,8,Wb,8,nc)
    q = np.round(dct / qt).astype(np.int16)

    body = zlib.compress(q.tobytes(), 6)
    hdr = struct.pack(">IIIIff", h, w, c, nc, float(quality), gain)
    return (MAGIC + struct.pack(">I", len(hdr)) + hdr
            + struct.pack(">I", mean.size * 4) + mean.astype("<f4").tobytes()
            + struct.pack(">I", comp.size * 4) + comp.astype("<f4").tobytes()
            + struct.pack(">I", len(body)) + body)


def decode(blob: bytes) -> np.ndarray:
    assert blob[:8] == MAGIC, "bad signature"
    off = 8
    def take():
        nonlocal off
        (n,) = struct.unpack(">I", blob[off:off + 4]); off += 4
        d = blob[off:off + n]; off += n
        return d
    h, w, c, nc, quality, gain = struct.unpack(">IIIIff", take())
    mean = np.frombuffer(take(), "<f4")
    comp = np.frombuffer(take(), "<f4").reshape(nc, c)
    q = np.frombuffer(zlib.decompress(take()), np.int16)

    ph, pw = (-h) % 8, (-w) % 8
    hp, wp = h + ph, w + pw
    q = q.reshape(hp // 8, 8, wp // 8, 8, nc).astype(np.float32)
    qt = _qtable(quality)[None, :, None, :, None]
    coords = idctn(q * qt, axes=(1, 3), norm="ortho").reshape(hp, wp, nc)[:h, :w]
    coords = coords / gain                                # undo the range rescale
    x = coords.reshape(-1, nc) @ comp + mean              # inverse KLT
    return x.reshape(h, w, c).astype(np.float32)


if __name__ == "__main__":
    import sys

    # Operates on DEQUANTISED float embeddings (int8 * scale) -> no scales file needed.
    if len(sys.argv) > 2:                                 # python klt_jpeg.py cube.npy scales.npy [quality]
        cube = np.load(sys.argv[1]).astype(np.float32)
        cube *= np.load(sys.argv[2])[:, :, None]
        quality = float(sys.argv[3]) if len(sys.argv) > 3 else 50.0
    else:                                                 # else a correlated synthetic cube
        rng = np.random.default_rng(0)
        latent = rng.standard_normal((256, 256, 16)).astype(np.float32)
        mix = rng.standard_normal((16, 128)).astype(np.float32)         # 128 dims from 16
        cube = (latent.reshape(-1, 16) @ mix).reshape(256, 256, 128)
        quality = 50.0

    blob = encode(cube, quality=quality)
    recon = decode(blob)
    rmse = float(np.sqrt(np.mean((recon - cube) ** 2)))
    rel = 100 * rmse / float(np.sqrt(np.mean(cube ** 2)))
    raw = cube.size * 4                                   # float32 reference
    print(f"shape {cube.shape}  quality {quality:g}")
    print(f"{len(blob):,} B   rel-RMSE {rel:.1f}%   ({raw / len(blob):.1f}x vs float32 raw)")