joeylitalien/noise2noise-pytorch

noise2noise‑pytorch – 无干净数据去噪

简介

Noise2Noise 论文（Lehtinen et al., 2018）的非官方 PyTorch 重新实现。该方法训练一个图像到图像的网络，将含噪输入映射到含噪目标，从而在无法获得干净参考图像的情况下学习恢复图像。

主要组件

U‑Net 风格的卷积网络（代码中存在一个小 bug，即编码器和解码器共享权重，但仍可正常工作）。
支持多种合成噪声模型的训练脚本（train.py）：
- 高斯噪声（通过 --noise-param 控制 σ）
- 泊松噪声（通过 --noise-param 指定 λ）
- 文本叠加（通过 --noise-param 指定概率 p）
- 蒙特卡洛渲染噪声（参见 MonteCarlo.md）。
测试脚本（test.py），用于加载检查点并可视化含噪/去噪/干净图像。
可选的训练过程中损失/PSNR 曲线绘图（--plot-stats）。

入门指南

安装所需的 Python 包（PyTorch 0.4.1、torchvision 0.2.0、NumPy、Matplotlib、Pillow、OpenEXR），使用 pip3 install -r requirements.txt。
准备数据 – 任何图像集合均可使用。README 展示了如何将 COCO 2017 验证集拆分为 train、valid 和 test 文件夹，或者您可以将脚本指向您自己的目录。

训练模型，例如针对高斯噪声：

python3 train.py \
  --train-dir data/train --train-size 1000 \
  --valid-dir data/valid --valid-size 200 \
  --ckpt-save-path ckpts \
  --nb-epochs 10 \
  --batch-size 4 \
  --loss l2 \
  --noise-type gaussian \
  --noise-param 50 \
  --crop-size 64 \
  --plot-stats \
  --cuda

对于其他损坏类型，请替换 --noise-type 和 --noise-param。

测试保存的检查点：

python3 test.py \
  --data data/test \
  --load-ckpt ckpts/gaussian/n2n.pt \
  --noise-type gaussian \
  --noise-param 50 \
  --crop-size 256 \
  --show-output 3 \
  --cuda

该脚本将显示含噪输入、去噪输出以及（如果可用）干净目标的并排蒙太奇图像。

仓库中展示的结果

高斯噪声（σ = 25）去噪图像的 PSNR 达到约 32 dB，接近干净目标基线。
文本叠加（p = 0.25）去噪图像的 PSNR 达到约 28 dB。
训练是在旧的 GTX 780 上进行了几百个 epoch 完成的。

已知限制 / 未解决问题

U‑Net 实现错误地让编码器与解码器共享权重；修复它需要少量的代码更改。
激活函数是反的（论文建议使用 ReLU 的地方使用了 LeakyReLU，反之亦然）；当前设置可以工作，但可能不是最优的。
由于泊松噪声是数据相关的，泊松噪声的处理与原始论文不完全一致；用户可能需要适应 TensorFlow 参考实现以获得适当的处理。

适用人群

实验自监督图像恢复的研究人员或爱好者。
需要在无干净真值数据的去噪任务中使用轻量级 PyTorch 基线的开发人员。

参考文献

Lehtinen, J. et al. “Noise2Noise: Learning Image Restoration without Clean Data”, ICML 2018.
示例数据的 COCO 数据集引用。

该仓库是一个用于机器学习图像恢复技术的真实、面向研究的代码库。

joeylitalien/noise2noise-pytorch

noise2noise‑pytorch – 无干净数据去噪

相关