File size: 4,673 Bytes
d9aae02 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | ---
license: openrail
tags:
- image-restoration
- image-denoising
- document-denoising
- onnx
- restormer
- cvpr
- transformer
- computer-vision
datasets: []
language: []
metrics:
- psnr
- ssim
pipeline_tag: image-to-image
---
# Restormer ONNX — Document Image Denoising
**ONNX-exported Restormer model for document image denoising and reconstruction.**
This repository contains pre-exported [ONNX](https://onnx.ai/) models of [Restormer](https://arxiv.org/abs/2111.09881) (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
```bash
pip install -r requirements.txt
```
### Python API
```python
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
```bash
# 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)
```python
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:
```bash
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
```bibtex
@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](https://huggingface.co/spaces/CompVis/stable-diffusion-license) license.
## Related
- [Restormer Official Repository](https://github.com/swz30/Restormer)
- [Paper (arXiv)](https://arxiv.org/abs/2111.09881)
- [Gradio Demo](https://huggingface.co/spaces/) — also in this repo (`app.py`)
|