E2A2 — image#
E2A2 (E²A², Encoding Efficient Asymmetric Autoencoders, github.com/danjacobellis/E2A2) is the successor of FRAPPE v3: the same split trainer (frozen per-group encoders — one learned strided projection per latent channel — SC6 softsign companding onto the coder alphabet [−31, 31], one merged decoder per operating point) with the entropy stage replaced by GGDLPC (a learned FIR predictor per channel, two affine laws mapping a causal statistic to a folded-GGD Huffman table, plus a derived zero-run mode; the C engine builds itself on import). Every reported rate is real byte-padded stream bits, round-trip verified in this notebook by decoding the latents back from the truncated bytes.
E2A2I is the E2A2 image codec (Kodak eval; ladder ps 32/16/8/4/2 × 3/6/6/6/6 channels, operating points 3/9/15/21/27; macroregion 64 px).
Stream layout two_stream: one global stream (the first scale group over the whole plane, one byte-pad) plus one detail stream per macroregion (one byte-pad each). The stream for operating point n is the literal bit-prefix covering channels ≤ n.
The codec is loaded from the Hugging Face hub (danjacobellis/E2A2, image/ subdirectory) through the compressors.e2a2 public API — there is no local checkpoint or training code in this notebook.
import json, glob
import torch, numpy as np, matplotlib.pyplot as plt
import PIL.Image, datasets
from torchvision.transforms.v2.functional import pil_to_tensor, to_pil_image
from huggingface_hub import hf_hub_download
from compressors import e2a2
Load the codec and its entropy coder#
One MergedAutoencoder per operating point, all from the hub. Each operating point has its own merged decoder (multidec); the encoders froze at merge, so every operating point’s encoder is a channel prefix of the full one. entropy.json is the GGDLPC blob (per-channel predictor taps + law scalars); the Huffman tables are generated from it when the coder loads.
NAME = 'image'
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
config, models, coder = e2a2.load_all_codecs(NAME, device=device)
ops = list(config.cumulative_channels)
model = models[ops[-1]]
rom = e2a2.rom_facts(coder.blob)
print(f'codec = {config.codec}/{NAME} (hub format {config.format}, source checkpoint {config.source_checkpoint})')
print(f'modality = {config.modality} spatial rank = {config.dim} input channels = {config.input_channels}')
print(f'ps ladder = {list(config.ps)} x {list(config.group_sizes)} channels')
print(f'boundaries = {ops} (cumulative channels = operating points)')
for n in ops:
d = config.decoder_configs[ops.index(n)]
print(f' op {n:>2}: decoder ps={d["decoder_ps"]} dim={d["decoder_dim"]} depth={d["decoder_depth"]}')
for s, (ps_s, start, end) in enumerate(model.scale_groups):
print(f' scale {s+1}: ps={ps_s:>4} channels {start+1}-{end}')
print(f'\nentropy coder: GGDLPC, stream_mode={coder.stream_mode}, macroregion={coder.macroregion} px'
+ (f', global stream = channels 1-{coder.n_coarse}' if coder.n_coarse else ', no global stream'))
print(f' alphabet [-31, 31]; persistent model = {rom["law_scalars"]} law scalars + '
f'{rom["predictor_ints"]} predictor ints for {rom["n_channels"]} channels')
print(f' {rom["note"]}')
codec = E2A2/image (hub format e2a2_hub1, source checkpoint e2a2i_g5.pth)
modality = image spatial rank = 2 input channels = 3
ps ladder = [32, 16, 8, 4, 2] x [3, 6, 6, 6, 6] channels
boundaries = [3, 9, 15, 21, 27] (cumulative channels = operating points)
op 3: decoder ps=8 dim=768 depth=12
op 9: decoder ps=8 dim=768 depth=12
op 15: decoder ps=8 dim=768 depth=24
op 21: decoder ps=8 dim=768 depth=32
op 27: decoder ps=8 dim=768 depth=32
scale 1: ps= 32 channels 1-3
scale 2: ps= 16 channels 4-9
scale 3: ps= 8 channels 10-15
scale 4: ps= 4 channels 16-21
scale 5: ps= 2 channels 22-27
entropy coder: GGDLPC, stream_mode=two_stream, macroregion=64 px, global stream = channels 1-3
alphabet [-31, 31]; persistent model = 162 law scalars + 162 predictor ints for 27 channels
edges, folded GGD pmfs, quotient split, canonical Huffman tables and the run mode are GENERATED from the per-channel scalars at load (GGDLPC.tables)
Analysis filterbank#
One learned ps×ps projection per latent channel. The top row of each panel is the impulse response (the RGB filter taps, normalized to ±4σ per scale); the bottom row is the magnitude response (2-D DFT magnitude, averaged over the three input channels, zero frequency at the center).
N = 65
for s, (ps_s, start, end) in enumerate(model.scale_groups):
W = model.encoders[s][0].weight.data.cpu()
n_g = end - start
sigma = W.std().item()
fig, axes = plt.subplots(2, n_g, figsize=(1.4 * n_g, 3.0), dpi=120)
axes = axes.reshape(2, n_g)
for j in range(n_g):
f = W[j]
axes[0, j].imshow((f.permute(1, 2, 0) / (4 * sigma) + 0.5).clamp(0, 1).numpy())
H = np.fft.fftshift(np.abs(np.fft.fft2(f.mean(0).numpy(), s=(N, N))))
axes[1, j].imshow(H, cmap='inferno')
axes[0, j].set_title(f'ch{start + j + 1}', fontsize=8)
for ax in (axes[0, j], axes[1, j]):
ax.axis('off')
axes[0, 0].set_title(f'ch{start + 1}\nimpulse', fontsize=8)
axes[1, 0].text(-0.25, 0.5, '|H(f)|', transform=axes[1, 0].transAxes,
rotation=90, va='center', fontsize=8)
fig.suptitle(f'scale ps={ps_s} (channels {start + 1}-{end})', fontsize=10)
plt.tight_layout()
plt.show()
Load an example image#
A Kodak image, mapped to the codec convention [-1, 1], cropped to a multiple of the macroregion size (a no-op at Kodak’s native sizes). linear_input=false, so pixels are fed in sRGB. Kodak is eval-only — nothing in the codec or coder was trained or fit on it.
dataset = datasets.load_dataset('danjacobellis/kodak', split='validation')
img = dataset[22]['image'].convert('RGB')
x = pil_to_tensor(img).to(torch.float).to(device).unsqueeze(0) / 127.5 - 1.0
unit = coder.macroregion
x = x[..., :unit * (x.shape[2] // unit), :unit * (x.shape[3] // unit)]
n_pixels = x.shape[2] * x.shape[3]
x_01 = x / 2 + 0.5
print(f'input {tuple(x.shape)} ({n_pixels} pixels = '
f'{x.shape[2] // unit} x {x.shape[3] // unit} macroregions of {unit}x{unit}), '
f'range [{x.min():.2f}, {x.max():.2f}]')
display(to_pil_image(x_01[0].cpu().clamp(0, 1)))
input (1, 3, 512, 768) (393216 pixels = 8 x 12 macroregions of 64x64), range [-1.00, 1.00]
Analysis transform and rounding#
Encode to the per-scale latents (companding included), then round to the coder’s [-31, 31] grid. encode_to_latents returns the int8 latents; the softsign asymptote makes the ±31 bound structural, so nothing is clipped.
with torch.no_grad():
latents = model.encode(x)
n_latent_values = sum(z.numel() for z in latents)
print(f'{x.numel()} input values -> {n_latent_values} latent values ({n_latent_values / x.numel():.3f}x)')
for s, z in enumerate(latents):
print(f' scale {s+1} (ps={model.scale_groups[s][0]:>2}): latent {tuple(z.shape)} range [{z.min():.1f}, {z.max():.1f}]')
n_clip = sum(int((z.round().abs() > 31).sum()) for z in latents)
latents_q = e2a2.encode_to_latents(model, x) # [(1, C_s, H_s, W_s) int8 on CPU, ...]
latents_np = [z[0].numpy() for z in latents_q]
print(f'clipped: {n_clip}/{n_latent_values}')
plt.figure(figsize=(5, 2), dpi=120)
plt.hist(np.concatenate([z.ravel() for z in latents_np]), range=(-31.5, 31.5), bins=63, width=0.85)
plt.xlim([-15, 15]); plt.title('histogram of integer latents (zoomed)')
plt.xlabel('value'); plt.ylabel('count'); plt.tight_layout(); plt.show()
1179648 input values -> 784512 latent values (0.665x)
scale 1 (ps=32): latent (1, 3, 16, 24) range [-15.4, 13.9]
scale 2 (ps=16): latent (1, 6, 32, 48) range [-8.2, 8.7]
scale 3 (ps= 8): latent (1, 6, 64, 96) range [-9.7, 8.7]
scale 4 (ps= 4): latent (1, 6, 128, 192) range [-8.2, 7.9]
scale 5 (ps= 2): latent (1, 6, 256, 384) range [-16.4, 16.1]
clipped: 0/784512
Quantized latents (prior to entropy coding)#
Every channel group is a multiple of 3 channels, so consecutive channel triples display directly as RGB images at the latent’s native resolution — the coder alphabet [-31, 31] mapped linearly to [0, 255] (0 → mid-gray).
for s, (ps_s, start, end) in enumerate(model.scale_groups):
z = latents_np[s]
for j in range(0, z.shape[0], 3):
rgb = ((z[j:j+3].astype(np.int16) + 31) * 255 // 62).astype(np.uint8)
print(f'ps={ps_s} channels {start + j + 1}-{start + j + 3} ({rgb.shape[2]}x{rgb.shape[1]} native)')
display(PIL.Image.fromarray(rgb.transpose(1, 2, 0)))
ps=32 channels 1-3 (24x16 native)
ps=16 channels 4-6 (48x32 native)
ps=16 channels 7-9 (48x32 native)
ps=8 channels 10-12 (96x64 native)
ps=8 channels 13-15 (96x64 native)
ps=4 channels 16-18 (192x128 native)
ps=4 channels 19-21 (192x128 native)
ps=2 channels 22-24 (384x256 native)
ps=2 channels 25-27 (384x256 native)
GGDLPC two-stream entropy coding#
encode_latents produces one global stream (channels 1–3, whole plane, one byte-pad) plus one detail stream per 64×64 macroregion (one byte-pad each). Per channel a learned FIR predictor forms the residual; two affine laws map the causal statistic (pooled template + last-16 |residual|) to a folded-GGD Huffman table; the decoder regenerates every table choice from already-decoded values. The stream for operating point n is the literal bit-prefix covering channels ≤ n. CR is against 24-bit RGB.
enc = e2a2.encode_latents(coder, latents_q)
print(f'global stream: {len(enc.gblob):,} bytes (channels 1-{coder.n_coarse}, whole plane, one byte-pad)')
region_bytes = (enc.prefix[:, -1] + 7) // 8
print(f'detail streams: {len(enc.streams)} macroregions, one byte-pad each; at max detail '
f'{int(region_bytes.sum()):,} bytes total, per region min/median/max = '
f'{int(region_bytes.min())}/{int(np.median(region_bytes))}/{int(region_bytes.max())}')
print(f'\n{"op":>4} {"stream bits":>12} {"bpp":>9} {"CR":>8} {"bits/latent":>12}')
for n in ops:
bits = e2a2.op_bits(coder, enc, n)
n_lat_n = sum(min(max(n - start, 0), end - start) * z.shape[1] * z.shape[2]
for z, (ps_s, start, end) in zip(latents_np, model.scale_groups))
print(f'{n:>4} {bits:>12,} {bits / n_pixels:>9.4f} {24 * n_pixels / bits:>7.1f}x {bits / n_lat_n:>12.3f}')
global stream: 571 bytes (channels 1-3, whole plane, one byte-pad)
detail streams: 96 macroregions, one byte-pad each; at max detail 100,711 bytes total, per region min/median/max = 271/955/2806
op stream bits bpp CR bits/latent
3 4,568 0.0116 2065.9x 3.965
9 23,288 0.0592 405.2x 2.246
15 78,184 0.1988 120.7x 1.655
21 178,648 0.4543 52.8x 0.918
27 810,256 2.0606 11.6x 1.033
Round-trip verification#
verify_latents decodes the latents back from the truncated bytes alone at every operating point’s literal bit-prefix (the global stream from its own bytes; every macroregion from its own bytes, no cross-region input) and asserts them bit-exact against the encoder’s. The reconstructions below are decoded from those byte-decoded latents (decode_latents), not from the encoder’s copy.
e2a2.verify_latents(coder, enc, latents_q, ops)
print(f'round trip verified at ops {ops}: every op decoded from its truncated bytes alone, bit-exact')
round trip verified at ops [3, 9, 15, 21, 27]: every op decoded from its truncated bytes alone, bit-exact
Reconstructions at each operating point#
One encode serves every operating point: the latents for op n are decoded from the truncated bytes (decode_latents) and fed to that point’s own merged decoder; the rate is the same encode’s literal bit-prefix. A fixed 256×256 crop of the original (left) and the reconstruction (right) per point.
ch_, cw_ = 256, 256
ct = (x.shape[2] - ch_) // 2
cl = (x.shape[3] - cw_) // 2
orig_crop = to_pil_image(x_01[0, :, ct:ct+ch_, cl:cl+cw_].cpu().clamp(0, 1))
rd_img = []
for n in ops:
lts = e2a2.decode_latents(coder, enc, n) # from the truncated bytes
with torch.no_grad():
xh = models[n].decode([z.to(device) for z in lts]).clamp(-1, 1)
bpp_n = e2a2.op_bits(coder, enc, n) / n_pixels
psnr_n = -10 * torch.nn.functional.mse_loss(x_01, xh / 2 + 0.5).log10().item()
rd_img.append((bpp_n, psnr_n, n))
print(f'{n:>2} ch: bpp={bpp_n:.4f} CR={24 / bpp_n:7.1f}x PSNR={psnr_n:.2f} dB')
recon_crop = to_pil_image((xh[0, :, ct:ct+ch_, cl:cl+cw_] / 2 + 0.5).cpu().clamp(0, 1))
pair = PIL.Image.new('RGB', (2 * cw_ + 4, ch_), 'white')
pair.paste(orig_crop, (0, 0)); pair.paste(recon_crop, (cw_ + 4, 0))
display(pair)
3 ch: bpp=0.0116 CR= 2065.9x PSNR=23.79 dB
9 ch: bpp=0.0592 CR= 405.2x PSNR=27.64 dB
15 ch: bpp=0.1988 CR= 120.7x PSNR=32.15 dB
21 ch: bpp=0.4543 CR= 52.8x PSNR=35.64 dB
27 ch: bpp=2.0606 CR= 11.6x PSNR=42.15 dB
Per-macroregion rate map#
Each macroregion carries its own byte-padded detail stream, so the spatial rate allocation is directly visible — and any macroregion could be transmitted or upgraded independently.
nuh, nuw = x.shape[2] // unit, x.shape[3] // unit
fig, axes = plt.subplots(1, 2, figsize=(9, 3.6), dpi=120)
axes[0].imshow(x_01[0].permute(1, 2, 0).cpu().clamp(0, 1).numpy())
axes[0].set_title('input'); axes[0].axis('off')
im = axes[1].imshow(region_bytes.reshape(nuh, nuw), cmap='viridis')
axes[1].set_title(f'detail-stream bytes per {unit}x{unit} macroregion')
axes[1].axis('off')
fig.colorbar(im, ax=axes[1], fraction=0.046)
plt.tight_layout(); plt.show()
Rate–distortion#
This image’s operating points over the whole-Kodak curve from results/e2a2_image/ (the compressors.e2a2 harness: one top-op encode per image, every op decoded from its byte-prefix; falls back to the hub rd.json when no harness run is present), against FRAPPE v3 (hub rd.json) and the FRAPPE v1 / v2 baselines from results/ (JPEG-LS accounting).
def latest(pattern):
paths = sorted(glob.glob(pattern))
return paths[-1] if paths else None
rd_path = latest(f'results/e2a2_{NAME}/rate_distortion_*.json')
if rd_path is not None:
rd = json.load(open(rd_path))
key = 'channel_counts' if 'channel_counts' in rd else 'quality_values'
rd_ops = list(rd[key])
rd_label = f"E2A2 {NAME} — whole validation split ({rd.get('n_images', rd.get('n_samples'))} samples, compressors harness, verified streams)"
else:
rd = None
hub_rd = e2a2.hub_rd_curve(NAME)
rd_ops = [p['n_ch'] for p in hub_rd['rd_curve']]
rd_label = f"E2A2 {NAME} — whole validation split (hub rd.json, trainer's verified streams)"
if rd is not None:
val_bpp = [rd['results'][str(n)]['mean']['bpp'] for n in rd_ops]
val_psnr = [rd['results'][str(n)]['mean']['PSNR_dB'] for n in rd_ops]
else:
val_bpp = [p['bpp'] for p in hub_rd['rd_curve']]; val_psnr = [p['psnr'] for p in hub_rd['rd_curve']]
rd3 = json.load(open(hf_hub_download(repo_id='danjacobellis/FRAPPEv3', filename='image/rd.json')))
v3_bpp = [p['bpp'] for p in rd3['rd_curve']]; v3_psnr = [p['psnr'] for p in rd3['rd_curve']]
v3_ops = rd3.get('ops') or [p.get('n_ch') for p in rd3['rd_curve']]
def load_baseline(pattern):
d = json.load(open(latest(pattern)))
key = 'channel_counts' if 'channel_counts' in d else 'quality_values'
pts = [(d['results'][str(q)]['mean']['bpp'], d['results'][str(q)]['mean']['PSNR_dB']) for q in d[key]]
return [p[0] for p in pts], [p[1] for p in pts]
v1_bpp, v1_psnr = load_baseline('results/frappe/rate_distortion_*.json')
v2_bpp, v2_psnr = load_baseline('results/frappe_v2_image/rate_distortion_*.json')
plt.figure(figsize=(6.5, 4), dpi=120)
plt.semilogx(v1_bpp, v1_psnr, '-', color='0.75', lw=1.0, label='FRAPPE v1 + JPEG-LS')
plt.semilogx(v2_bpp, v2_psnr, '^-', color='black', ms=5, lw=1.2, alpha=0.8, label='FRAPPE v2 + JPEG-LS')
plt.semilogx(v3_bpp, v3_psnr, 'D-', color='tab:gray', ms=4, lw=1.0, label='FRAPPE v3 (lawcoder, hub rd.json)')
plt.semilogx(val_bpp, val_psnr, 'o-', color='tab:blue', label=rd_label)
plt.semilogx([p[0] for p in rd_img], [p[1] for p in rd_img], 's--', color='tab:blue', alpha=0.5, label='E2A2 image — this image')
for n, bv, psv in zip(rd_ops, val_bpp, val_psnr):
plt.annotate(f'{n}ch', (bv, psv), fontsize=7, color='tab:blue', xytext=(3, -9), textcoords='offset points')
plt.xlabel('bpp'); plt.ylabel('PSNR (dB)')
plt.title('E2A2 image vs FRAPPE v1/v2/v3 — Kodak')
plt.grid(True, which='both', alpha=0.3); plt.legend(fontsize=7); plt.tight_layout(); plt.show()
print('E2A2 image, Kodak average (verified streams):')
for n, bv, psv in zip(rd_ops, val_bpp, val_psnr):
print(f' {n:>2} ch: bpp={bv:.4f} CR={24 / bv:7.1f}x PSNR={psv:.2f} dB')
E2A2 image, Kodak average (verified streams):
3 ch: bpp=0.0107 CR= 2246.6x PSNR=21.53 dB
9 ch: bpp=0.0624 CR= 384.3x PSNR=25.08 dB
15 ch: bpp=0.2489 CR= 96.4x PSNR=28.70 dB
21 ch: bpp=0.7541 CR= 31.8x PSNR=32.94 dB
27 ch: bpp=3.0311 CR= 7.9x PSNR=40.76 dB