joeylitalien/noise2noise-pytorch

PyTorch Implementation of Noise2Noise (Lehtinen et al., 2018)

noise2noise‑pytorch – 깨끗한 데이터 없이 노이즈 제거

개요

  • Noise2Noise 논문(Lehtinen et al., 2018)의 비공식 PyTorch 재구현입니다. 이 방법은 깨끗한 참조 이미지가 없는 경우에도 이미지를 복원할 수 있도록 노이즈 입력을 노이즈 대상으로 매핑하는 이미지-이미지 네트워크를 학습합니다.

주요 구성 요소

  • U‑Net 스타일의 합성곱 네트워크(인코더와 디코더가 가중치를 공유하는 작은 버그가 있지만 여전히 작동합니다).
  • 여러 합성 노이즈 모델을 지원하는 학습 스크립트(train.py):
    • 가우시안(--noise-param으로 σ 제어)
    • 푸아송(--noise-param으로 λ 지정)
    • 텍스트 오버레이(--noise-param으로 확률 p 지정)
    • 몬테카를로 렌더링 노이즈(MonteCarlo.md 참조).
  • 체크포인트를 로드하고 노이즈/노이즈 제거/깨끗한 이미지를 시각화하는 테스트 스크립트(test.py).
  • 학습 중 손실/PSNR 곡선의 선택적 플롯(--plot-stats).

시작하기

  1. 필요한 Python 패키지(PyTorch 0.4.1, torchvision 0.2.0, NumPy, Matplotlib, Pillow, OpenEXR)를 pip3 install -r requirements.txt를 통해 설치합니다.
  2. 데이터 준비 – 어떤 이미지 컬렉션이든 작동합니다. README에서는 COCO 2017 검증 세트를 train, valid, test 폴더로 나누는 방법을 보여주지만, 스크립트를 사용자 자신의 디렉토리로 지정할 수도 있습니다.
  3. 모델을 학습합니다. 예를 들어 가우시안 노이즈의 경우:
    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을 교체하세요.
  4. 저장된 체크포인트를 테스트합니다:
    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) 노이즈 제거 이미지는 약 32 dB PSNR에 도달하여 깨끗한 대상 기준선에 근접합니다.
  • 텍스트 오버레이(p = 0.25) 노이즈 제거 이미지는 약 28 dB PSNR에 도달합니다.
  • 학습은 오래된 GTX 780에서 몇백 에포크 동안 수행되었습니다.

알려진 제한 사항 / 미해결 문제

  • U‑Net 구현에서 인코더 가중치가 디코더와 잘못 공유되고 있습니다. 수정하려면 작은 코드 변경이 필요합니다.
  • 활성화 함수가 반대로 되어 있습니다(논문에서는 ReLU를 제안하지만 LeakyReLU가 사용되고 있으며 그 반대도 마찬가지입니다). 현재 설정은 작동하지만 최적일 수 없습니다.
  • 푸아송 노이즈는 데이터에 의존적이기 때문에 푸아송 노이즈 처리는 원래 논문과 완전히 일치하지 않습니다. 적절한 처리를 위해 사용자가 TensorFlow 참조 구현을 적응시켜야 할 수 있습니다.

사용 대상

  • 자기 지도 학습 이미지 복원을 실험하는 연구자나 취미 개발자.
  • 깨끗한 정답 데이터가 없는 노이즈 제거 작업에 경량 PyTorch 베이스라인이 필요한 개발자.

참고 문헌

  • Lehtinen, J. et al. “Noise2Noise: Learning Image Restoration without Clean Data”, ICML 2018.
  • 예시 데이터용 COCO 데이터셋 인용.

이 리포지토리는 기계 학습 이미지 복원 기법을 위한 진정한 연구 지향 코드베이스입니다.

관련

  • 프로젝트
  • 프로젝트
  • 프로젝트
  • 프로젝트