nagi_denoise / README.md
uniuyuni's picture
Add NagiV2-L weights, Core ML exports and model card (README.md)
c1349a6 verified
|
Raw
History Blame Contribute Delete
8.84 kB
---
license: cc-by-nc-4.0
library_name: pytorch
pipeline_tag: image-to-image
tags:
- denoising
- image-restoration
- photography
- hdr
- coreml
- apple-silicon
---
# Nagi Denoise β€” NagiV2-L
A blind denoiser for **real high-ISO photographs**. Float32 linear-light RGB
in, float32 linear-light RGB out. HDR-safe (values above 1.0 are preserved,
never clipped), seam-free at any resolution including 40MP+, and fully
deterministic β€” same bytes in, same bytes out.
* **Architecture**: NagiV2-L, **15.43M parameters** (width 64, encoder
(2,2,4), middle 6, decoder (1,2,2)).
* **Code**: <https://github.com/uniuyuni/nagi_denoise> (Apache-2.0)
* **Weights**: this repo (CC BY-NC 4.0 β€” see [Licence](#licence))
## Files in this repo
| file | size | what it is |
|---|---|---|
| `nagi_v2_l_ft2_final.pt` | 236 MB | the production PyTorch checkpoint |
| `nagi_v2_l_ft2_t768_fp16.mlpackage/` | 30 MB | Core ML export, 768px tile, fp16 β€” the fast path |
| `nagi_v2_l_ft2_t768_fp32.mlpackage/` | 59 MB | Core ML export, 768px tile, fp32 |
The `.mlpackage` assets are directories, stored here as folders of files.
## Usage
```python
import numpy as np
from nagi_denoise import denoise
# img: (H, W, 3) float32 numpy array, linear-light RGB. Values > 1.0 are fine.
out = denoise(img) # PyTorch reference path
out = denoise(img, backend="coreml") # Core ML fast path on Apple Silicon
```
The weights are resolved locally first and downloaded from this repo only if
they are not already present: an explicit `weights=` path wins, then
`$NAGI_DENOISE_WEIGHTS`, then the in-repo `runs/` copy, then an existing
Hugging Face cache, and only then a download. `allow_download=False` (or
`NAGI_DENOISE_OFFLINE=1`) forbids the network entirely.
To fetch deliberately, ahead of time:
```python
from nagi_denoise.assets import resolve_weights, resolve_coreml_package
resolve_weights() # -> Path to nagi_v2_l_ft2_final.pt
resolve_coreml_package() # -> Path to the fp16 .mlpackage
```
## ⚠️ Core ML: never use `compute_units="all"`
`ALL` lets Core ML dispatch to the Apple Neural Engine, and **the ANE computes
this fp16 graph wrongly**. On every one of the 23 validation tiles, peak
output values run **1.25x–4.8x higher** than the PyTorch reference β€” worst
absolute per-pixel error **8.7 on a scene whose true peak is 5.3**. Visually
this is blown-out garbage, not a subtle numerical difference.
The same graph on `cpu_and_gpu` agrees with PyTorch to **0.018** max, which is
just fp16 rounding. `cpu_and_gpu` is the default in this project and it must
stay that way.
## Training data
| source | weight | licence |
|---|---|---|
| SIDD (Smartphone Image Denoising Dataset) | 0.25 | MIT |
| PolyU Real-World Noisy Images Dataset | 0.35 | **non-commercial** |
| synthetic Poisson-Gaussian degradations, generated by this project | 0.40 | this project's code (Apache-2.0) |
The synthetic noise is **spatially correlated**, not white. Real demosaiced
sensor noise is correlated; training on white noise alone teaches the model
that "isolated pixel deviation = noise", so it preserves real correlated noise
blobs as if they were fine structure. Measured lag-1 autocorrelation of the
residual after a 3x3 median:
| source | lag-1 autocorrelation |
|---|---|
| Fujifilm X-T5 (real) | +0.235 |
| Pentax K-5 (real) | +0.139 |
| PolyU (real) | +0.27 … +0.48 |
| SIDD (real, but near-white) | +0.024 |
| naive white synthetic noise | βˆ’0.146 |
The synthetic noise field is therefore drawn as a Gaussian-blurred, unit-std
normal field, with the blur sigma sampled uniformly per image over
**[0.0, 1.0]** pixels (chroma field: [0.8, 2.0]). The sweep used to calibrate
it β€” sigma 0.5 / 0.7 / 0.9 β†’ lag-1 0.005 / 0.217 / 0.365 β€” puts sigma β‰ˆ 0.7 on
the real X-T5 figure and sigma β‰ˆ 0.9 in the PolyU range.
## Measured results
| benchmark | result |
|---|---|
| **SIDD Validation (sRGB PSNR)** | **39.030 dB** |
| β€” NAFNet-w64 teacher, for reference | 40.21 dB, at 116M parameters (7.5x the size) |
**Structure-vs-noise selectivity**, X-T5 Occi hair ROI β€” how much more
high-frequency energy survives on structure than on flat noise:
| pipeline | selectivity | retention on structure |
|---|---|---|
| **NagiV2 (this model)** | **+11.2 pt** | **90.5%** |
| legacy in-house v12 pipeline | +9.8 pt | 82.1% |
More detail retained *and* cleaner β€” not a trade.
**HDR highlight retention** β‰₯ 0.99 on every scene with true HDR content.
**Speed**, 39.8MP frame, end to end, on an M1 with 16GB:
| path | time |
|---|---|
| Core ML fp16 / `cpu_and_gpu` | **~83 s** |
| pure PyTorch / MPS | ~261 s |
## Known limitations
Stated plainly, because they are real:
* **Isolated small specular highlights on low-dynamic-range files get treated
as impulse noise.** On Z7 fix / Z7 bird, top-1% luma retention is
**0.62–0.92**. The conditional highlight guard does not help here by design:
it only arms when the image genuinely contains above-SDR content, because
arming it on low-range scenes blends noisy input back in over large areas
(measured 3.7x more flat-region noise) and reads as uneven denoising.
* **The confidence-gated detail head is inert on this checkpoint.** The gate
is closed, so `detail_strength` is a no-op by design. It is wired up so it
activates automatically if a future checkpoint opens the gate.
* **It removes noise; it does not generate texture.** This is a restoration
model, not a generative one. It will not match the synthesised detail of
DxO PhotoLab DeepPrime XD, and it is not trying to.
## Licence
**The weights in this repo are CC BY-NC 4.0 β€” non-commercial.**
The source code at <https://github.com/uniuyuni/nagi_denoise> is Apache-2.0.
Full text: <https://creativecommons.org/licenses/by-nc/4.0/legalcode>
The weights are not more restrictive by choice. They are a derived work of
their training data, and one of the sources carries a non-commercial
restriction that flows through to anything trained on it:
> PolyU Real-World Noisy Images Dataset
> Copyright (c) 2018, The Hong Kong Polytechnic University
> "Any redistribution, use, or modification is done solely for
> non-commercial purposes."
PolyU supplied **35%** of the training mixture, and it was not an incidental
ingredient: it was **the only source of real spatially-correlated camera
noise**, which is exactly what made the model work on real photographs at all.
The other sources are permissive β€” SIDD is MIT, and the synthetic degradations
are generated by this project's own code.
You may: use the weights for any non-commercial purpose; redistribute them,
modified or not; fine-tune, distil, quantise or convert them and redistribute
the result. You must: give attribution, state that the weights are CC BY-NC
4.0, and indicate whether you modified them. Derivative weights inherit these
terms.
### If you need commercially usable weights
**Retrain without PolyU.** The code to do so is Apache-2.0 and complete: drop
`polyu` from `data.mixture` in a training config, redistribute its weight
between `sidd` and `synthetic`, and train per the repo's README. Expect a
quality cost β€” removing the only real correlated-noise source is a material
change and the result needs re-validating against the project's gates. Weights
you produce that way are yours, and are not covered by this licence.
### No warranty
The weights are provided "as is", without warranty of any kind, express or
implied. See the CC BY-NC 4.0 text for the full disclaimer.
## Citation
The datasets and reference models this work depends on:
```bibtex
@inproceedings{abdelhamed2018sidd,
title = {A High-Quality Denoising Dataset for Smartphone Cameras},
author = {Abdelhamed, Abdelrahman and Lin, Stephen and Brown, Michael S.},
booktitle = {CVPR},
year = {2018}
}
@article{xu2018polyu,
title = {Real-world Noisy Image Denoising: A New Benchmark},
author = {Xu, Jun and Li, Hui and Liang, Zhetong and Zhang, David and Zhang, Lei},
year = {2018},
note = {PolyU Real-World Noisy Images Dataset. Non-commercial use only.}
}
@article{chen2022nafnet,
title = {Simple Baselines for Image Restoration},
author = {Chen, Liangyu and Chu, Xiaojie and Zhang, Xiangyu and Sun, Jian},
journal = {arXiv:2204.04676},
year = {2022}
}
@article{zhang2022scunet,
title = {Practical Blind Image Denoising via Swin-Conv-UNet and Data Synthesis},
author = {Zhang, Kai and Li, Yawei and Liang, Jingyun and Cao, Jiezhang and
Zhang, Yulun and Tang, Hao and Timofte, Radu and Van Gool, Luc},
year = {2022}
}
```
NAFNet (MIT) is vendored in the repo as the distillation teacher and
benchmark reference; SCUNet (Apache-2.0) as a benchmark reference. Full
third-party notices are in `NOTICE` in the source repo.