InkUNet gated UNet++ — handwriting erasure
Scan/photo of an answered exam page in, printed page without handwriting out. Single RGB input (no clean GT at test time).
$$\hat{A}=\mathrm{clip}(B+m\odot\tanh(r),0,1)$$
- \(B\): answered page
- \(m=\sigma(\mathrm{logit})\): ink gate
- \(r\): residual (paper stays identity when \(m\approx 0\))
Best checkpoint: epoch 53, 2.43 M params, 3× downsample (8× bottleneck), channels [32, 64, 128, 256].
Hub: https://huggingface.co/felixwang1024/InkUNet-erase-gated
Demo (Space): Gradio UI matching local erase_service (gated UNet++ only). Code is in the sibling folder InkUNet-Space/. Hosting a Gradio Space on Hugging Face currently needs HF PRO for even cpu-basic; after you create felixwang1024/InkUNet-erase-demo, the app can be pushed there. Local: cd InkUNet-Space && python app.py.
Results (same 115 SCUT pages, Baidu 1st-place tiling)
Protocol = public mix test.py: reflect pad 106, window 512, step 300, keep interior, hflip TTA. Test set = InkEraser/data/val (SCUT-EnsExam test). Not the hidden Baidu A-board.
| Model | PSNR ↑ | SSIM ↑ | L1 ↓ | ink L1 ↓ |
|---|---|---|---|---|
| Baidu 1st mix (Paddle) | 28.78 | 0.935 | 0.0122 | 0.232 |
| This model | 33.69 | 0.956 | 0.00656 | 0.0339 |
Center-512 val (training monitor, not full-page): PSNR 35.13, SSIM 0.953, L1 0.00673, ink L1 0.0396.
JSON: results/metrics_baidu_scut115.json, results/metrics_center512.json.
Qualitative triples (inked | pred | GT): examples/val_stem122_triple.png, examples/val_stem126_triple.png.
Files
| Path | What |
|---|---|
weights/gated_unetpp_erase.pt | Inference (EMA only, ~9 MB) |
weights/best.pt | Full training ckpt (optimizer + EMA, 38 MB) |
config.yaml | Architecture + tiling |
infer.py | CLI |
src/unetpp.py | Gated UNet++ |
Do not downsample a full A4 to 256/512. 512 is a tile at native resolution.
Install
pip install -r requirements.txt
# GPU: install a CUDA build of torch yourself
From this folder, or after download:
huggingface-cli download felixwang1024/InkUNet-erase-gated --local-dir ./InkUNet-erase-gated
cd InkUNet-erase-gated
python infer.py --input answered.jpg --output clean.png
# gaussian tiles (no Baidu interior crop):
python infer.py --input answered.jpg --output clean.png --protocol gauss
Load in code:
import sys, torch
from pathlib import Path
ROOT = Path(".") # this repo after download
sys.path.insert(0, str(ROOT))
from src.model import build_model
from src.utils import load_config
cfg = load_config(ROOT / "config.yaml")
model = build_model(cfg)
ckpt = torch.load(ROOT / "weights/gated_unetpp_erase.pt", map_location="cpu", weights_only=False)
model.load_state_dict(ckpt["ema"], strict=False)
model.eval()
Input tensor: 1×3×H×W in [0, 1]. Call model(x, None)["pred"][-1].
Method (short)
- Data: stroke width aligned to SCUT (~2.8 px): dehw A4 scaled so max side = 1500; drop coarse-dpi
dehw_train*from clear; add real SCUT-A4 + painted ink colors. n = 11287. Val = 115 SCUT test pages, never in train. - Arch: UNet++ nested skips at full res (
X^{0,*}); drop 4th stride-2 (8× not 16×); gate so paper is identity. - Loss: full L1 + 0.2(1−SSIM) + 4× ink L1 + paper Sobel + mask BCE/Dice + paper
|r|. - Select:
score = PSNR − 50·inkL1 − 40·paperGrad − 20·(1−Dice), not raw PSNR.
Generalization
This checkpoint is fit to the datasets we already have, not to “any document photo.” In-domain numbers (SCUT-EnsExam test, 115 pages) should not be read as out-of-domain (OOD) performance.
What it was trained / selected on
- Train: scaled dehw A4 (max side 1500), dehw crops at 640, non-dehw pages from
dehw_train_clear, real SCUT-A4, plus painted ink on those nets. Stroke width is aligned to SCUT (~2.8 px, ~140 dpi). - Val / the numbers in the table: SCUT test only (same corpus, held-out pages). That is in-domain for this model.
What we already measured as OOD
- UNet++ mask (same family, ink segmentation): md5 book-scan val F1 0.991, then full-page dehw exam test F1 0.736 (recall 0.99 → 0.65). Nested skips help in-domain; domain shift dominates architecture.
- Baidu 1st-place mix: trained on synthetic dehw, ink L1 0.232 on real SCUT (almost copy-input). The reverse of our shift: high-dpi synthetic → low-dpi scans.
- This erase model on high-dpi A-board pages: the gate often stays shut (looks like identity). Training squeezed pens to SCUT scale; a ~7 px stroke at ~360 dpi is a different visual class.
Expect degradation when
- dpi / stroke width is far from ~2.8 px (phone photos, 300–400 dpi archives, strong JPEG zoom);
- paper, print font, ruling, or layout is not exam-scan-like;
- ink color / tool is rare in the paint mix (e.g. highlighter, fountain-pen sheen);
- large circles, boxes, and ink-on-print — already weak in-domain (gate Dice ≈ 0.67).
Do not treat SCUT PSNR 33.69 or mix’s hidden-set 38.68 as a score on an unseen domain. For a new domain, collect paired pages in that dpi and finetune; do not only stack a bigger UNet++.
Limits
Large circles/boxes and ink-on-print leave faint ghosts. Curriculum fill head did not beat this ckpt. Baidu A-board has no public GT; PSNR vs mix pseudo-GT (24.71) is not a leaderboard score.
Citation
@misc{inkunet-erase-gated-2026,
title = {Gated residual UNet++ for exam handwriting erasure},
year = {2026},
note = {HW-Erase / InkUNet, compared to Baidu handwriting-erasure 1st-place mix on SCUT-EnsExam test}
}
UNet++: Zhou et al., TMI 2019. Tiling: zdyshine Baidu 1st-place test.py.