jere1882/document-image-denoising

Document Image Denoising

Six trained models for removing JPEG/blur/noise artifacts from scanned document images. Three architectures (unet, residual_unet, plain_encoder_decoder) and a width ablation of residual_unet down to 1.6% of its original parameter count. Each is provided as a PyTorch Lightning checkpoint and as fp32/int8 ONNX.

Code, training notebooks, and the full writeup: github.com/jere1882/DocDenoising. Training data: jeremiasrodriguez/document-crops-50k on Kaggle (50k crops rendered from scanned PDFs).

Results

Trained on 2×T4 GPUs, 10 epochs, identical data/degradation/seed across runs. PSNR is the global metric (one PSNR from the summed squared error over the full held-out validation split, not an average of per-image PSNRs — those disagree by several dB and aren't the same quantity).

Model	Params	PSNR fp32 (dB)	PSNR int8 (dB)	Size fp32→int8 (ONNX)	CPU latency fp32→int8 (ms/img)*	Train time
unet (bc64)	7.70M	30.33	30.24	30.8MB → 7.8MB	288 → 243	6h33m
residual_unet (bc64)	7.70M	30.25	30.18	30.8MB → 7.8MB	286 → 239	7h06m
residual_unet (bc32)	1.93M	30.03	30.00	7.7MB → 2.0MB	90 → 76	3h26m
residual_unet (bc16)	0.48M	29.22	29.22	1.9MB → 0.5MB	44 → 34	2h44m
residual_unet (bc8)	0.12M	27.36	27.23	0.5MB → 0.16MB	18 → 18	2h35m
plain_encoder_decoder (bc64)	6.93M	28.28	28.20	27.7MB → 7.0MB	239 → 194	5h41m

* Single 512×512 tile, batch 1, 4 CPU threads (Apple Silicon, one machine — not a general hardware claim). int8 is static-quantized (QDQ, per-channel weights), calibrated on 500 training crops never seen by the validation set.

Takeaways (full analysis in the repo README):

• Residual learning (predicting the noise, not the clean image) converges faster early on but reaches the same final PSNR as direct prediction — no final-quality edge, contrary to the usual assumption.
• Skip connections matter far more than width: plain_encoder_decoder (no skips) loses ~2 dB versus unet/residual_unet (same depth, with skips); cutting residual_unet to 32 channels (25% of params) only loses 0.2 dB.
• int8 quantization costs ≤0.14 dB everywhere and a consistent 4× size reduction, down to 160KB for the smallest model.

Usage

import torch
from PIL import Image
# clone github.com/jere1882/DocDenoising for the model + inference code
from denoising.models.lit_module import DenoisingLitModule
from denoising.infer import make_predictor, tile_denoise

model = DenoisingLitModule.load_from_checkpoint("residual_unet_bc64/checkpoint.ckpt", map_location="cpu")
predict = make_predictor(checkpoint="residual_unet_bc64/checkpoint.ckpt", onnx_path=None, device=torch.device("cpu"))
denoised = tile_denoise(Image.open("page.png"), predict, tile_size=512, overlap=32, batch_size=1)

Or load model_int8.onnx directly with onnxruntime for a dependency-free CPU path.

Files

Each folder (unet_bc64, residual_unet_bc{64,32,16,8}, plain_encoder_decoder_bc64) contains:

• checkpoint.ckpt — PyTorch Lightning checkpoint (fp32 weights + optimizer state)
• model_fp32.onnx / model_int8.onnx — ONNX exports, dynamic height/width axes

License

MIT (code and weights). Training data is cited above, not redistributed here — see its own Kaggle page for terms.