synthscript/drunet-color-onnx

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.

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