metadata
license: mit
library_name: pytorch
pipeline_tag: image-to-image
tags:
- image-restoration
- denoising
- image-denoising
- pytorch
- low-level-vision
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.