scunet-onnx

SCUNet — Image Denoising (ONNX, full 8-variant bundle)

ONNX exports of SCUNet (Swin-Conv-UNet) — Kai Zhang et al., 2022. Hybrid CNN + Swin Transformer architecture for image denoising. This repo bundles all 8 published checkpoints from the upstream model_zoo/ so you get the full size / variant ladder in a single download.

Re-exported from upstream PyTorch weights. Provenance trail: Zhang et al. → cszn/SCUNet model_zoo/*.pth → torch.onnx.export (one pass per checkpoint) → these files.

Toolchain: torch 2.4.x (CUDA 12.4), timm latest, einops latest, thop latest, onnx latest, onnxruntime>=1.17, opset 17, do_constant_folding=True. Full conversion script: scripts/export-kair.ps1 in the Heliosoph repo (runs once per .pth checkpoint via -Model scunet-color or -Model scunet-gray).

Credit: Kai Zhang, Yawei Li, Jingyun Liang, Jiezhang Cao, Yulun Zhang, Hao Tang, Deng-Ping Fan, Radu Timofte, Luc Van Gool. Paper: "Practical Blind Image Denoising via Swin-Conv-UNet and Data Synthesis", 2022.

What this repo contains

Each variant ships as an .onnx (small graph file) + .onnx.data (~70 MB of external tensor data) sibling pair — torch's ONNX exporter externalizes weights at the size/opset combination used here. Both files must be present in the same directory at load time — the .onnx references the .data by relative filename.

Color variants (5)

Grayscale variants (3)

The grayscale variants are ~3× cheaper to run than the color variants on grayscale inputs (they accept 1-channel input directly; the color variants need the gray channel replicated across RGB).

Input / output (all variants)

All variants share the same forward-pass shape; the only differences are the input channel count and the trained weights.

How to use

import onnxruntime as ort
import numpy as np
from PIL import Image

# Pick a variant. Both the .onnx and .onnx.data must be present in
# the same directory — ORT resolves the external data automatically.
sess = ort.InferenceSession("scunet_color_real_psnr.onnx")

img = Image.open("noisy.jpg").convert("RGB")
W, H = img.size
W8, H8 = (W // 8) * 8, (H // 8) * 8                          # 8-align
img = img.crop((0, 0, W8, H8))

arr = np.asarray(img, dtype=np.float32) / 255.0              # HWC, [0,1]
arr = arr.transpose(2, 0, 1)[None, ...]                      # 1x3xHxW

result = sess.run(None, {"image": arr.astype(np.float32)})[0][0]
result = np.clip(result, 0.0, 1.0).transpose(1, 2, 0)
Image.fromarray((result * 255).astype(np.uint8)).save("denoised.jpg")

Which one should I use?

• General-purpose photo denoising: scunet_color_real_psnr — blind, faithful, no guesswork required.
• Consumer photo cleanup (subjectively pretty matters more than ground truth): scunet_color_real_gan.
• Matched-σ benchmark or known-noise scenario: pick the _15, _25, or _50 variant that matches your noise level.
• Grayscale (medical / document / B&W): use scunet_gray_* directly — ~3× faster than the color variant on gray inputs.
• Comparison demos: the σ-specialist variants are great for showing matched-vs-mismatched specialist behavior. Run the same noisy image through scunet_color_{15,25,50} and the differences are visually obvious.

For fixed-σ Gaussian denoising in a research-benchmark context, SwinIR's swinir_denoising_color_25 is the apples-to-apples transformer counterpart. For denoising + sharpening as one step, look at NAFNet (opencv/deblurring_nafnet) — different task (deblur) but adjacent.

License

Apache-2.0 — same as the upstream cszn/SCUNet repo. LICENSE file included.