JPEG XL: lossless encoding like cjxl -e 3

Author: matmen

CHANGELOG.md

@@ -1,6 +1,16 @@
 # Changelog
 
-## 3.13.3
+## 3.14.0
+
+### New
+
+- Lossless `encodeJXL()` compresses like libjxl's `cjxl -e 3` (effort 3) by default: about 15% smaller files for
+  photos (a 1920x1080 photo: 1.17 MB instead of 1.39 MB) and 13% for graphics, within 1-2% of cjxl's sizes. It
+  uses a palette of up to 1024 colors where it fits (else palettes of the channels that use few of their values),
+  YCoCg and the weighted predictor in contexts of its error, and codes the groups in worker threads
+  (`JXLEncodeOptions.threads`, the same files): the photo takes 0.24 s with 8 threads, 0.65 s with one. `effort: 0`
+  keeps the fast mode, byte for byte what `cjxl -e 1` writes (0.11 s for the photo); `effort` applies to lossless
+  compression now too.
 
 ### Performance
 

scripts/jxl-lossless-references.py

@@ -1,5 +1,5 @@
 # Prints the SHA-256 of what cjxl -e 1 (libjxl's fast lossless mode) writes for the images of tests/jxl-encode.js:
-# Image.encodeJXL writes the same. The channels are those encodeJXL picks (grey if every pixel is, alpha if one isn't
+# Image.encodeJXL writes the same at effort 0. Then the sizes of what cjxl -e 3 writes (effort 1 is within a few percent). The channels are those encodeJXL picks (grey if every pixel is, alpha if one isn't
 # opaque). libjxl's output depends on the CPU: with AVX2 or NEON it works in chunks of 16 pixels (what encodeJXL does),
 # without SIMD in chunks of 8 (the test checks these too), with AVX-512 in chunks of 32. Needs Pillow and cjxl:
 #   CJXL=path/to/cjxl python scripts/jxl-lossless-references.py
@@ -21,6 +21,15 @@ def images():
     yield 'external.png grey', np.dstack([red, red, red, np.full_like(red, 255)])
 
 
+def save(rgba):
+    grey = (rgba[..., 0] == rgba[..., 1]).all() and (rgba[..., 0] == rgba[..., 2]).all()
+    opaque = (rgba[..., 3] == 255).all()
+    channels = [0] if grey else [0, 1, 2]
+    if not opaque: channels.append(3)
+    a = rgba[..., channels]
+    Image.fromarray(a[..., 0] if len(channels) == 1 else a, {1: 'L', 2: 'LA', 3: 'RGB', 4: 'RGBA'}[len(channels)]).save(tmp / 'in.png')
+
+
 for name, rgba in images():
     grey = (rgba[..., 0] == rgba[..., 1]).all() and (rgba[..., 0] == rgba[..., 2]).all()
     opaque = (rgba[..., 3] == 255).all()
@@ -30,3 +39,8 @@ def images():
     Image.fromarray(a[..., 0] if len(channels) == 1 else a, {1: 'L', 2: 'LA', 3: 'RGB', 4: 'RGBA'}[len(channels)]).save(tmp / 'in.png')
     subprocess.run([CJXL, '--quiet', str(tmp / 'in.png'), str(tmp / 'out.jxl'), '-d', '0', '-e', '1'], check=True)
     print(f"    '{name}': '{hashlib.sha256((tmp / 'out.jxl').read_bytes()).hexdigest()}',")
+
+for name, rgba in images():
+    save(rgba)
+    subprocess.run([CJXL, '--quiet', str(tmp / 'in.png'), str(tmp / 'out.jxl'), '-d', '0', '-e', '3'], check=True)
+    print(f"    '{name}': {(tmp / 'out.jxl').stat().st_size},")

src/codecs/jxl.ts

@@ -16,6 +16,7 @@ import { decodeFrame, DecoderState, type FrameThreads } from './jxl/frame.ts';
 import { ColorOutput } from './jxl/color.ts';
 import { ToneMapping } from './hdr.ts';
 import { encodeRGBA } from './jxl/fast-lossless.ts';
+import { encodeModular } from './jxl/modular-lossless.ts';
 import { encodeVarDCT } from './jxl/vardct-enc.ts';
 import { concat } from '../util/mem.ts';
 import { srgbDecode, type Curve } from '../util/icc.ts';
@@ -33,7 +34,8 @@ export function isJXL(d: Uint8Array): boolean {
  * codestream, in the container if there is Exif data (a TIFF header and IFDs); lossy with up to threads worker threads
  */
 export async function encode(rgba: Uint8Array, width: number, height: number, exif?: Uint8Array, distance = 0, effort = 1, threads = 1): Promise<Uint8Array> {
-    const stream = distance > 0 ? await encodeVarDCT(rgba, width, height, { distance, effort, threads }) : encodeRGBA(rgba, width, height);
+    const stream = distance > 0 ? await encodeVarDCT(rgba, width, height, { distance, effort, threads })
+        : effort ? await encodeModular(rgba, width, height, threads) : encodeRGBA(rgba, width, height);
     if (!exif?.length) return stream;
     // (ftyp: brand jxl, minor version 0, compatible jxl; Exif: the offset of the TIFF header, 0, then the TIFF data)
     const jxl = [0x6a, 0x78, 0x6c, 0x20];

src/codecs/jxl/entropy-enc.ts

@@ -39,6 +39,14 @@ const token = (v: number): number => {
     return SPLIT + ((n - SPLIT_EXPONENT) << (MSB + LSB)) + ((m >>> (n - MSB)) << LSB) + (m & ((1 << LSB) - 1));
 };
 
+/** The number of tokens (of 32-bit values) */
+export const MAX_TOKEN = token(0xffffffff) + 1;
+
+/** Counts the tokens of values (from to to) by context, into counts (MAX_TOKEN per context) */
+export function countTokens(counts: Float64Array, ctx: Uint32Array, value: Uint32Array, from: number, to: number): void {
+    for (let i = from; i < to; i++) counts[ctx[i] * MAX_TOKEN + token(value[i])]++;
+}
+
 // ---- ANS
 
 const ANS_LOG_TAB_SIZE = 12, ANS_TAB_SIZE = 1 << ANS_LOG_TAB_SIZE, LOG_ALPHABET_SIZE = 8;
@@ -250,22 +258,19 @@ export class EntropyCode {
     private readonly ctxMap: Uint8Array;
     private readonly ans: AnsCode[];
 
-    constructor(streams: Tokens[], readonly numContexts: number, { maxClusters = 64 }: EntropyOptions = {}) {
-        // the histograms of the contexts' tokens
-        let alphabet = 1;
-        for (const t of streams) {
-            for (let i = 0; i < t.length; i++) {
-                const tk = token(t.value[i]);
-                if (tk + 1 > alphabet) alphabet = tk + 1;
-            }
-        }
-        const hist = Array.from({ length: numContexts }, () => new Float64Array(alphabet)), totals = new Float64Array(numContexts);
-        for (const t of streams) {
-            for (let i = 0; i < t.length; i++) {
-                hist[t.ctx[i]][token(t.value[i])]++;
-                totals[t.ctx[i]]++;
-            }
+    /** the code of streams' tokens, or of their counts (see countTokens) */
+    constructor(streams: Tokens[] | Float64Array, readonly numContexts: number, { maxClusters = 64 }: EntropyOptions = {}) {
+        // the histograms of the contexts' tokens (up to the largest token used)
+        let counts: Float64Array;
+        if (streams instanceof Float64Array) counts = streams;
+        else {
+            counts = new Float64Array(numContexts * MAX_TOKEN);
+            for (const t of streams) countTokens(counts, t.ctx, t.value, 0, t.length);
         }
+        let alphabet = 1;
+        for (let i = 0; i < counts.length; i++) if (counts[i] && i % MAX_TOKEN + 1 > alphabet) alphabet = i % MAX_TOKEN + 1;
+        const hist = Array.from({ length: numContexts }, (_, c) => counts.slice(c * MAX_TOKEN, c * MAX_TOKEN + alphabet));
+        const totals = Float64Array.from(hist, h => h.reduce((sum, v) => sum + v, 0));
         this.ctxMap = new Uint8Array(numContexts);
         const clusters = cluster(hist, totals, Math.min(maxClusters, 255, numContexts), this.ctxMap).map(h => {
             let size = h.length;
@@ -293,29 +298,45 @@ export class EntropyCode {
 
     /** writes the tokens of a stream */
     writeTokens(w: BitWriter, t: Tokens): void {
-        const map = this.ctxMap;
-        // ANS: the tokens in reverse, the 16 bits of each renormalization read after the token's symbol (before its
-        // raw bits); the final state first
-        const codes = this.ans, n = t.length, chunks = new Int32Array(n);
-        let state = 0x130000;
-        for (let i = n - 1; i >= 0; i--) {
-            const code = codes[map[t.ctx[i]]], tk = token(t.value[i]), freq = code.freq[tk];
-            chunks[i] = -1;
-            if (state >>> 20 >= freq) {
-                chunks[i] = state & 0xffff;
-                state >>>= 16;
-            }
-            const v = Math.floor(state / freq);
-            state = v * ANS_TAB_SIZE + code.reverse[tk][state - v * freq];
+        writeTokens(w, this.tables(), t.ctx, t.value, 0, t.length);
+    }
+
+    /** what writing tokens needs (for threads, which get copies) */
+    tables(): CodeTables {
+        return { ctxMap: this.ctxMap, freq: this.ans.map(c => c.freq), reverse: this.ans.map(c => c.reverse) };
+    }
+}
+
+/** An entropy code's tables for writing tokens: the context map, per cluster the symbols' frequencies and table slots */
+export interface CodeTables {
+    ctxMap: Uint8Array;
+    freq: Int32Array[];
+    reverse: Int32Array[][];
+}
+
+/** Writes tokens (from to to) with a code */
+export function writeTokens(w: BitWriter, { ctxMap: map, freq: freqs, reverse }: CodeTables, ctx: Uint32Array, value: Uint32Array, from: number, to: number): void {
+    // ANS: the tokens in reverse, the 16 bits of each renormalization read after the token's symbol (before its raw
+    // bits); the final state first
+    const chunks = new Int32Array(to - from);
+    let state = 0x130000;
+    for (let i = to - 1; i >= from; i--) {
+        const c = map[ctx[i]], tk = token(value[i]), freq = freqs[c][tk];
+        chunks[i - from] = -1;
+        if (state >>> 20 >= freq) {
+            chunks[i - from] = state & 0xffff;
+            state >>>= 16;
         }
-        w.write(32, state);
-        for (let i = 0; i < n; i++) {
-            if (chunks[i] >= 0) w.write(16, chunks[i]);
-            const v = t.value[i];
-            if (v >= SPLIT) {
-                const nbits = 31 - Math.clz32(v) - MSB - LSB;
-                w.write(nbits, (v >>> LSB) & ((1 << nbits) - 1));
-            }
+        const v = Math.floor(state / freq);
+        state = v * ANS_TAB_SIZE + reverse[c][tk][state - v * freq];
+    }
+    w.write(32, state);
+    for (let i = from; i < to; i++) {
+        if (chunks[i - from] >= 0) w.write(16, chunks[i - from]);
+        const v = value[i];
+        if (v >= SPLIT) {
+            const nbits = 31 - Math.clz32(v) - MSB - LSB;
+            w.write(nbits, (v >>> LSB) & ((1 << nbits) - 1));
         }
     }
 }

src/codecs/jxl/fast-lossless.ts

@@ -694,8 +694,8 @@ export function encode(px: Uint8Array, width: number, height: number, numChans:
 
 // ---- headers
 
-// a regular modular frame: one pass, no filters, replacing, the last one
-function frameHeader(w: BitWriter, alpha: boolean): void {
+/** A regular modular frame: one pass, no filters, replacing, the last one */
+export function frameHeader(w: BitWriter, alpha: boolean): void {
     w.write(1, 0);
     w.write(2, 0b00);
     w.write(1, 1);