efwfe/restomer

Restormer ONNX — Document Image Denoising

ONNX-exported Restormer model for document image denoising and reconstruction.

This repository contains pre-exported ONNX models of Restormer (CVPR 2022 Oral), trained for document denoising. Run inference with ONNX Runtime — no PyTorch required.

Model	Precision	Size	Input Shape
restormer_denoise_dynamic.onnx	FP32	105 MB	[B, 3, H, W] (dynamic)
restormer_fp16_converted.onnx	FP16	55 MB	[B, 3, H, W] (dynamic)

Quick Start

Install

pip install -r requirements.txt

Python API

import cv2
from inference import denoise

# Read an image
img_bgr = cv2.imread('noisy_document.png')
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

# Denoise
result = denoise(img_rgb, model='fp32')  # or model='fp16'

# Save
cv2.imwrite('clean_document.png', cv2.cvtColor(result, cv2.COLOR_RGB2BGR))

Command Line

# Single image
python inference.py -i noisy_doc.png -o clean_doc.png

# FP16 model (faster, smaller)
python inference.py -i noisy_doc.png -o clean_doc.png --model fp16

# CPU only
python inference.py -i noisy_doc.png -o clean_doc.png --cpu

# Batch process a folder
python inference.py -i ./noisy/ -o ./clean/ --batch

ONNX Runtime (Minimal)

import onnxruntime as ort
import numpy as np
import cv2

session = ort.InferenceSession('models/restormer_fp16_converted.onnx')

img = cv2.imread('noisy.png')
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
img = img.transpose(2, 0, 1)[np.newaxis, ...]  # [1, 3, H, W]

# Pad to multiple of 8
h, w = img.shape[2], img.shape[3]
H, W = ((h + 7) // 8) * 8, ((w + 7) // 8) * 8
img_padded = np.pad(img, ((0, 0), (0, 0), (0, H - h), (0, W - w)), mode='reflect')

output = session.run(['output'], {'input': img_padded})[0]
output = (output[0].transpose(1, 2, 0) * 255).clip(0, 255).astype(np.uint8)
cv2.imwrite('clean.png', cv2.cvtColor(output, cv2.COLOR_RGB2BGR))

Model Details

Architecture

Restormer is an efficient Transformer architecture for high-resolution image restoration. Key innovations:

• MDTA (Multi-DConv Head Transposed Attention): depth-wise convolutions in the attention mechanism for local context
• GDFN (Gated-Dconv Feed-Forward Network): gated mechanism with depth-wise convolutions
• UNet-style encoder-decoder with skip connections for multi-scale feature aggregation

Property	Value
Params	~16.08 M
Input	RGB image, any size (padded to multiple of 8)
Output	RGB image, same resolution
Training data	Document / scanned document images
ONNX opset	18

Precision

The FP32 and FP16 models produce visually identical results. FP16 is recommended for deployment:

• FP32: 105 MB, maximum fidelity
• FP16: 55 MB, ~2× smaller, faster on compatible hardware

ONNX Export

The models were exported from PyTorch using:

python export_onnx.py

This produces both fixed-size (512×512) and dynamic-size ONNX graphs. The dynamic models in this repo support arbitrary input resolutions.

Tile Processing

For large images (> 512×512), the inference script automatically splits the image into overlapping tiles, processes each independently, and blends them seamlessly with a feathered weight map:

Tile size: 512×512, Overlap: 64 pixels, Multiple-of: 8

Benchmarks

Input Size	PyTorch (GPU)	ONNX FP32 (CPU)	ONNX FP16 (CPU)
256×256	~52 ms	~480 ms	~310 ms
512×512	~95 ms	~1,700 ms	~1,100 ms

Measured on NVIDIA RTX 3090 (GPU) / AMD EPYC 64-core (CPU). Actual performance depends on hardware.

Citation

@inproceedings{zamir2022restormer,
  title={Restormer: Efficient Transformer for High-Resolution Image Restoration},
  author={Zamir, Syed Waqas and Arora, Aditya and Khan, Salman and Hayat, Munawar
          and Khan, Fahad Shahbaz and Yang, Ming-Hsuan},
  booktitle={CVPR},
  year={2022}
}

License

This model is released under the OpenRAIL license.

Related

• Restormer Official Repository
• Paper (arXiv)
• Gradio Demo — also in this repo (app.py)