---
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: **<your-github-url>**.

## 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

```python
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.