DenoiseGAN — detail-preserving single-step image denoiser

A single-step image denoiser tuned to remove noise while preserving fine detail, plus a small noise-translator front-end that adapts it to noise it was not trained on. Code: .

Files

File	What
psnr_final.pt	The denoiser (PSNR stage). Use the ema weights.
translator_0020000.pt	Noise translator front-end (~0.37M params). Enable for Gaussian / OOD noise.

Usage

import torch
from models import DenoiseGenerator, NoiseTranslator   # from the GitHub repo

D = DenoiseGenerator(channels=(48,96,192,320,448), use_checkpoint=False).eval().cuda()
D.load_state_dict(torch.load('psnr_final.pt')['ema'], strict=False)

T = NoiseTranslator().eval().cuda()                      # optional
T.load_state_dict(torch.load('translator_0020000.pt')['ema'])

with torch.no_grad():
    out = D(T(noisy))      # noisy normalized to [-1, 1]; drop T for in-distribution noise

Model

• Generator (~21M params): NAFNet/Restormer-hybrid U-Net, single forward pass, residual noise prediction, attention-gated skips.
• Translator (~0.37M params): bias-free, scale-equivariant residual CNN trained through the frozen denoiser so D(T(noisy)) ≈ clean.

Notes & limitations

• Shipped model is the PSNR (reconstruction) stage. The adversarial/GAN stage was tried and did not improve results (it traded fidelity for hallucination).
• Optimized for detail preservation on a specific noise family rather than for topping a single PSNR benchmark.
• For Gaussian / out-of-distribution noise, enable the translator.

License

MIT.