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 – 無乾淨數據去噪

相關