| --- |
| license: mit |
| library_name: onnx |
| tags: |
| - image-denoising |
| - onnx |
| - drunet |
| - raw-photography |
| --- |
| |
| # DRUNet (colour) — ONNX export |
|
|
| An ONNX export of **DRUNet**, the plug-and-play denoiser from Kai Zhang's |
| [KAIR](https://github.com/cszn/KAIR) toolbox. Nothing here is retrained or modified: this is the |
| original `drunet_color.pth` weights, exported to ONNX so they can be run from ONNX Runtime. |
|
|
| Published for use in [PxlMonk](https://pxlmonk.com), a desktop raw photo editor, whose build fetches |
| this file at build time rather than committing a 130 MB binary. |
|
|
| ## Provenance |
|
|
| | | | |
| |---|---| |
| | Architecture | `UNetRes` from [cszn/KAIR](https://github.com/cszn/KAIR) — MIT | |
| | Weights | `drunet_color.pth`, from the KAIR v1.0 release — unmodified | |
| | Export | `torch.onnx.export`, opset 17, dynamic batch/height/width | |
| | Paper | Zhang et al., *Plug-and-Play Image Restoration with Deep Denoiser Prior*, TPAMI 2021 | |
|
|
| The weights were trained by the original author and are redistributed unchanged, in a different file |
| format. No training data is redistributed here. |
|
|
| ### A note on the training corpora |
|
|
| DRUNet was trained on `train400`, **DIV2K**, **Flickr2K** and the Waterloo Exploration Database. Some of |
| those carry their own terms — DIV2K states: |
|
|
| > This dataset is made available for academic research purpose only. All the images are collected from |
| > the Internet, and the copyright belongs to the original owners. |
|
|
| Whether such dataset terms reach the resulting model weights is unsettled, and this repository takes no |
| position on it. What it does do is state the provenance plainly, so anyone using these weights can make |
| that call with the facts in front of them rather than discovering them later. The MIT tag above reflects |
| the licence under which the original author published the weights. |
|
|
| ## Interface |
|
|
| Input `input`: `float32[batch, 4, height, width]`, NCHW. |
|
|
| - Channels 0–2: RGB, **sRGB display-referred**, range 0…1. The model was trained on display images — |
| feeding scene-linear data gives poor results. |
| - Channel 3: the **noise level map**, a constant `sigma / 255` across the plane. This is DRUNet's |
| strength control: sigma is in 0…255 units, so `sigma = 20` means a plane filled with `20/255`. |
|
|
| Output `output`: `float32[batch, 3, height, width]`, same convention, denoised. |
|
|
| Height and width must be divisible by 8 (three downsampling stages). Pad or tile accordingly. |
|
|
| ## Choosing sigma |
|
|
| There is no single right value — it is the strength dial. Measured on a 24 MP frame shot at ISO 16000, |
| using mean deviation from the local 3×3 mean as a noise figure (lower is cleaner): |
|
|
| | sigma | noise figure | | |
| |---|---|---| |
| | — (untouched) | 7.41 | | |
| | 10 | 1.67 | visible noise remains | |
| | **20** | **0.23** | clean, fine detail intact | |
| | 35 | 0.17 | over-smoothed — eyelashes go waxy | |
|
|
| ## Reproducing this export |
|
|
| The weights are stored **without bias terms**, so the network must be constructed with `bias=False`. |
| Loading with `strict=False` silently leaves 64 randomly-initialised bias tensors in place and produces |
| a model that runs but denoises badly — always load strictly. |
|
|
| ```python |
| from models.network_unet import UNetRes # from cszn/KAIR |
| import torch |
| |
| net = UNetRes(in_nc=4, out_nc=3, nc=[64, 128, 256, 512], nb=4, act_mode="R", |
| downsample_mode="strideconv", upsample_mode="convtranspose", bias=False) |
| net.load_state_dict(torch.load("drunet_color.pth", map_location="cpu", weights_only=True), strict=True) |
| net.eval() |
| |
| torch.onnx.export( |
| net, torch.randn(1, 4, 256, 256), "drunet_color.onnx", |
| input_names=["input"], output_names=["output"], |
| dynamic_axes={"input": {0: "b", 2: "h", 3: "w"}, "output": {0: "b", 2: "h", 3: "w"}}, |
| opset_version=17, |
| ) |
| ``` |
|
|
| ## Execution providers |
|
|
| Measured on Apple Silicon (10 cores), 256×256 tile, best of three after a warm-up: |
|
|
| | | | |
| |---|---| |
| | CPU | 560 ms | |
| | CoreML | **117 ms** | |
| | WebGPU | 154 ms | |
|
|
| DRUNet is a plain CNN, so it maps cleanly onto accelerators — unlike transformer-based denoisers, where |
| the graph fragments into hundreds of partitions and the copying between them costs more than the |
| compute saves. |
|
|
| ## Licence |
|
|
| MIT, matching KAIR. Include the original copyright notice when redistributing: |
|
|
| > Copyright (c) 2019 Kai Zhang |
|
|
