Un-squashing my family photos with a 2.6M-parameter network

Originally published at ludwina.top, where this post also has interactive charts and a tiling demo.

The story starts when I was roughly in middle school. Our family photo library was getting a bit obese - dozens of gigabytes felt like a lot back then. I started searching for a solution and found out about this little GUI around mozJPEG (probably this one mozJpegGUI), that could magically turn megabyte photos into kilobyte photos. So I ran the whole collection through it at something like quality=60. It took me years to understand just what I did at that time, and I've been trying to undo my mistakes since.

My first proper attempt was porting FBCNN to ComfyUI, which fixed the photos nicely, but ComfyUI was never meant to be a good batch processor for images. FBCNN also wasn't light. Even on my fairly beefy GPU it took seconds per picture, which meant days for the whole collection.

Recently, I had to pick my master's thesis topic and I wanted to pick this exact problem - JPEG artifact removal. My promoter unfortunately turned down the idea and wanted me to go into NR-IQA (no-reference image quality assessment) instead. This gave me a final push to just it by myself. dejpeg is the result of this effort. My own restoration network and a public website that runs the model locally, live at dejpeg.ludwina.top for everyone to use.

One of my photos, 512x512 crop. As PNG it is 394 KiB; as JPEG quality 10 with 4:2:0 subsampling, 10.2 KiB. Right side: DeJPEGNet's restoration, running in the browser.

What JPEG actually destroys

JPEG cuts an image into 8x8 blocks and keeps a handful of cosine frequencies per block. Quantization throws away the high frequencies, more aggressively at low quality, and once they're gone a sharp edge no longer fits inside a block. What's left is blocky steps and mosquito ringing around edges.

The same crop before and after JPEG quality 10 with 4:2:0 subsampling, at 4x zoom. Blue lines mark the 8x8 DCT grid the codec works on.

The useful part is that JPEG is a fixed algorithm, so the damage is structured. It always lands on the same 8x8 grid, always the same families of blur and ringing. None of that is perfectly invertible, the information really is gone, but a network can learn what the missing structure plausibly looked like without inventing new content. The plan naturally emerges from that: restore the codec's damage and don't repaint the photo. Having a small network gets you that almost for free: it can't hallucinate new details, it's simply too dumb to.

Why not just run FBCNN in the browser

The app actually started that way, running my ONNX port of FBCNN. Nothing wrong with FBCNN, apart from its size (a bit unfit for a light browser app). I just knew that with a bit of research and good old trial-and-error I can achieve better results than a 5-year-old network. I had a couple of assumptions:

• I need something small - if it's going to run directly on people's computers, it needs to be efficient
• The size also should bring the benefit of simplified hosting (e.g. Cloudflare pages size limits)
• I can make the model more universal - remove flexible sizing in favor of tiling, remove flexible quality in favor of auto-guessing (I used auto in FBCNN anyways), and remove separate models for color and grayscale - it should perform well in both scenarios

So: a single small model for everything, that can be compiled to native ONNX.

DeJPEGNet

DeJPEGNet is a small three-level U-Net (2.6M parameters) built from NAFNet-style blocks. Most of its design follows from one constraint: memory on the user's machine is limited, and you can't know how limited until you're running. Most commonly, this is solved by processing an image in tiles. This works best when the processing does not require global context of the image (e.g. composition, lighting, content) - a small crop of an image, which has some white texture in it can be both a piece of paper and a white wall. You do not know which one it is without looking at the rest of the image. But JPEG artifact removal is a local problem: the codec's damage is local, and the network can learn to fix it locally.

My original plan was to just accept the tiling: lock the input to a constant 512x512 and build the tiling around that. Turns out I didn't need to. The architecture I ended up with has no global operation anywhere, so it works on any input size by itself. I kept tiling anyway, because smaller inputs save memory.

What's inside the blocks:

• Local channel attention (LSCA) pools over a fixed 32x32-pixel window. It's NAFNet's simple channel attention with a window, and the window is what makes tiles not matter: a pixel a thousand pixels away has no way to affect the output. Two overlapping crops of the same image produce identical pixels on the overlap, as long as the overlap is wider than the network's bounded reach (at the deepest U-Net level a 32px window spans ~128 input pixels, so a few hundred pixels is a probable real value).
• There is no ReLU or GELU anywhere. Nonlinearity comes from SimpleGate (split the channels in half, multiply the halves) and from multiplicative attention. Fewer element-wise ops and kernel launches, and the ONNX export gets easier.
• Each block's residual branches are scaled by zero-initialized beta/gamma, and the output head is a zero-init conv. At step zero the whole network computes the identity function, so it starts out doing nothing and only learns corrections. If it doesn't understand something, it should leave it alone.
• Training uses parallel 3x3, 1x1, and identity branches; at export they are algebraically turned into a single depthwise 3x3, the same structural re-parameterization RepVGG uses for its deploy path.

Left: the full network - a 3-level U-Net, widths 40/80/160 channels, 2.6M parameters, output is input plus a zero-initialized residual. Right: one NAFBlock. The 32x32 attention window is what makes tiled inference seam-free.

The exported model is a 5.4 MB FP16 ONNX file with dynamic axes, so any image size that's a multiple of 16 px works with no re-export.

Training: 64 GB of photos and a degradation sampler

There is no public paired dataset of squashed and clean images, and you wouldn't want one anyway, since it would bake in a single encoder's idea of degradation. So training pairs are synthesized on the fly from clean sources: 604 WebDataset shards (~64 GB) built from DIV2K, Flickr2K, and LIU4K-class corpora, plus ~18 GB of my own raw photos straight from my camera.

Each sample takes a clean crop, applies a random degradation, and emits the (degraded, clean) pair. The sampler mimics what actually happens to photos in the wild:

• quality factors sampled across 1–100, weighted toward the common 30–75 range
• chroma subsampling mixed 4:2:0 / 4:2:2 / 4:4:4, progressive JPEG 10% of the time
• clean PNGs re-saved as JPEG at random qualities between 1 and 100 (generation loss, since real images get re-saved all the time)
• random resizes before and after compression, mild noise, and a small crop jitter so the 8x8 grid lands at random phases like real crops do

Additionally, 10% of pairs skip degradation entirely, because without that the model got too eager on near-clean inputs and 25% get converted to grayscale-replicated RGB, because otherwise the model learns that "JPEG artifacts are colorful" and starts hallucinating color into grayscale photos.

The loss is L1 plus LPIPS, a learned perceptual distance: roughly, how different two images look to a human, rather than how many pixels differ. LPIPS runs on a random 128x128 crop rather than the full 256px patch, because full-patch LPIPS under bf16 autocast often returned NaNs and the crop stays numerically stable. Training uses the VGG flavor of LPIPS, while evaluation uses LPIPS-Alex and DISTS. Grading a network with the same feature extractor it trained against isn't a great way to get honest numbers.

The rest is standard: AdamW at lr 1e-3, warmup then cosine decay, gradient accumulation to an effective batch of 16, an EMA of the weights (a slowly-updated smoothed copy, and the smoothed copy is what ships), bf16 plus torch.compile, 150k iterations over 256px patches. I trained it locally on my Windows machine through WSL (much better ML library support than Windows), with the dataset sitting on an old 2 TB spinning disk.

The numbers: perceptually good, not really pixel-perfect

Restorations are scored against pristine originals on the standard Classic5 and LIVE1 sets, re-encoded at quality factors (QF) 10–40 with 4:2:0, with FBCNN re-run under the identical protocol. DISTS is a learned perceptual distance: lower is closer to how the difference looks.

QF	Input	FBCNN (35.7M)	DeJPEGNet (2.6M)
10	0.2412	0.1938	0.1191
20	0.1536	0.1597	0.0911
30	0.1215	0.1405	0.0743
40	0.1054	0.1268	0.0642

Classic5 benchmark, DISTS. Lower is closer to the original.

Look at the FBCNN column: from QF 20 up it scores worse than the compressed input it was given. It's trained to get pixels close, and chasing that smooths texture in a way the perceptual metric doesn't like. LPIPS agrees with DISTS here (Classic5 QF10: 0.1804 for DeJPEGNet vs 0.2594 for FBCNN).

FBCNN does win where you'd expect, on PSNR, by about 1 dB (29.80 vs 28.81 dB at QF10). Decibels are logarithmic, so a 1 dB lead is easier to read as a ratio of error power relative to the compressed input. That's the price of picking a perceptual metric as the main one. A mild unsharp filter applied after restoration (default 0.10, tunable, off-able) buys back a good chunk of the sharpness at a negligible perceptual cost.

The app: tiles that shrink themselves

The inference stack lives in inference-core, a framework-free TypeScript package, so the worker, codecs, and tiling also work outside the Astro/Svelte app. ONNX Runtime Web runs the model in a worker, on WebGPU with WASM as the fallback.

Large images are sliced into configurable tiles with tile/8 overlap and cosine-feathered back together. If a device doesn't have the memory for 512px tiles, the app doesn't crash, it tries again with a smaller tile size.

One thing that surprised me is that batching multiple tiles per one iteration looked faster but actually was considerably slower on weak GPUs, so the shipped default is one tile for one iteration.

There is an interactive demo of the tiling in the original post: you can move a tile across a damaged photo and shrink the tile size to watch the pass re-plan around more, smaller tiles. Try it at ludwina.top.

I expected deployment to be the annoying part. Serving models (or fairly large files in general) is a pain without a CDN (classically paid), and I wanted to avoid that. I looked into a few ways to serve the model: Cloudflare's R2, GitHub Releases, HuggingFace, jsDelivr and some more. In the end the 5.4 MB model just fits in the static bundle, well under Cloudflare's 25 MB per-asset limit, so there's nothing to host separately and no CORS to fight. Additionally, everything is cached, so one download of assets makes the website work offline.

Where it falls short

• If you want pixel-exact restoration measured in dB, FBCNN wins by about 1 dB. This model is tuned for how the results look, not for averaging pixels.
• FBCNN has a "quality" knob. Is this useful? In my experience not really, since you normally just want the "auto" mode.
• The degradation sampler covers the common real-world mess (multi-pass saves, resizes, subsampling mixes), but it's still synthetic. Some encoder somewhere will surprise it.
• Compression can remove detail and this model does not hallucinate it back. That's mostly by design and cuts both ways: restored photos look softer than the original, but the model can stay small, because it never has to learn how to invent detail.

Steal it

Everything is open (Apache-2.0): the app and training code, the live tool, and the model weights via GitHub Releases. There's also a plain Python package if browser inference isn't your thing:

from dejpeg import load_model, restore_image

model = load_model() # weights auto-download
restore_image("input.jpg", "out.png", model)

Training your own variant from any folder of clean images is one command, and since degradations are synthesized on the fly, your photo library is a dataset.

If you manage to train a better model, please let me know. I would love to be able to use it.