DRUNet (colour) — ONNX export
An ONNX export of
DRUNet, the plug-and-play denoiser from Kai Zhang's
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, 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 — 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.
1from models.network_unet import UNetRes # from cszn/KAIR
2import torch
3
4net = UNetRes(in_nc=4, out_nc=3, nc=[64, 128, 256, 512], nb=4, act_mode="R",
5 downsample_mode="strideconv", upsample_mode="convtranspose", bias=False)
6net.load_state_dict(torch.load("drunet_color.pth", map_location="cpu", weights_only=True), strict=True)
7net.eval()
8
9torch.onnx.export(
10 net, torch.randn(1, 4, 256, 256), "drunet_color.onnx",
11 input_names=["input"], output_names=["output"],
12 dynamic_axes={"input": {0: "b", 2: "h", 3: "w"}, "output": {0: "b", 2: "h", 3: "w"}},
13 opset_version=17,
14)
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
