IDR: Self-Supervised Image Denoising via Iterative Data Refinement

Paper

GitHub - zhangyi-3/IDR: Self-Supervised Image Denoising via Iterative Data Refinement (CVPR2022)

Introduction

Image denoising task에 있어서, noise image와 clean image의 pair를 취득하는것은 매우 까다롭고 high-cost인 작업입니다.

그렇기 때문에 noise image들만 사용하여 pair를 구성하여 학습하는 noise2noise와 같은 방법론이 소개되었습니다.

하지만 noise image pair를 얻는 것 또한 완벽하게 align된 scene을 새롭게 취득해야하기 때문에 쉬운 일이 아닙니다.

이를 해결하고자 noise2self, noise2void, neighbor2neighbor 등 single noisy image만을 사용하여 self-supervised learning을 하려는 연구가 계속해서 존재하였지만, 만족스럽지 못한 결과를 보여주었습니다.

최근에는 noise image에 synthetic noise를 더하여 noisier - noise image를 pair로 학습하는 Noisier2Noise 라는 방법이 소개되었고, 이전의 연구보다 성능이 뛰어나며 noise2noise나 noise-clean image로 학습한 결과에도 뒤떨어지지 않는 결과를 보여주었습니다.

해당 방법을 사용하여 학습하면 synthetic noise와 비슷한 분포의 noise는 매우 잘 없어지지만, 실제 noisy-clean image와 noisier-noise data 사이에 존재하는 일종의 “data bias” 때문에 actual noise에 대해서는 매우 떨어지는 성능을 보이게 됩니다.

따라서 IDR 논문에서는, 기존 noisier2noise 방법론에 더해서 학습에 있어서 noisy-clean image와 noisier-noisy pair 사이의 data bias를 줄이기 위한 방법을 제시합니다.

Methods

Pilot Study on Data Bias

noisier2noise 방법으로 모델을 학습했을 때 actual noise $x$ 에 대해서는 성능이 좋지 않은 이유가 data bias 때문이라고 가정하고 이를 확인하기 위해 선행연구를 진행하였습니다.

clean image $y$ 에 대해서 noise image $x:=y+n$ 과 noisier image $x+n$ 을 생성하였고, $\{{y_i+n_i, y_i}\}, \{{x_i+n_i, x_i}\}$ 에 대해서 U-net을 학습하고, 학습을 마친 모델에 대해서 noise image를 inference하여 결과를 분석하였습니다.

우선 noisier-noisy dataset 으로 학습했을 때 actual noise에 대한 denoising 성능을 확인하기 위해 Gaussian, Raw, Correlated noise에 대해서 각각 학습을 진행하였습니다.

모든 noise type에 대해서 clean image와의 PSNR을 구해봤을 때, 모델을 통과하지 않은 noisy image보다 높은 PSNR를 보여주었기 때문에 다음과 같은 결론을 얻을 수 있습니다.

Finding (1): noisier-noisy dataset으로 학습한 모델이 actual noise에 대해서도 denoising을 할 수 있다.

다음으로는 data bias가 denoising에 어떠한 영향을 미치는지 알아보기 위해서 기존 $\{{y_i+n_i, y_i}\},$ 로 학습한 모델에 임의로 gaussian noise, blur를 주어 data bias를 만든 후 inference를 진행하였습니다.

모든 noise type에 대해서 gaussian의 $\sigma$ 값이 커지는것에 비례하여 기존의 결과보다 PSNR Drop 이 일어나는것을 확인할 수 있었고, 따라서 다음과 같은 결론을 얻을 수 있습니다.

Finding (2): ideal noisy-clean dataset과 비교하였을때 더 적은 biased dataset에 대해서 학습된 모델이 더 좋은 denoising 성능을 갖는다.

이 두가지 findings를 통해서, noisier-noisy 와 noisy-clean dataset의 data bias 줄이는것을 통해 denoising 성능의 향상을 이끌어낼 수 있다는것을 알 수 있습니다.

Iterative Data Refinement

논문에서는 data bias를 줄이기 위한 방법으로 Iterative Data Refinement (IDR) 을 제시합니다.

우선, noisier-noise dataset $\{{x_i+n_i, x_i}\}$ 으로 model $\mathcal{F_0}$ 를 학습합니다.

Finding(1) 에 따르면, 이렇게 학습된 모델도 actual noisy image에 대해서도 어느 정도의 denoising을 할 수 있습니다.

따라서 $\mathcal{F_0}$ 에 noisy image $x_i$ 를 통과시켜 refined noisy image $\mathcal{F_0}(x_i)$ 를 얻습니다.

얻은 refined noisy image에 다시 임의의 synthetic noise $n_i$ 를 더해 새로운 noisier image $\mathcal{F_0}(x_i)+n_i$ 를 얻어 최종적으로 refined trainning dataset:

$\{\mathcal{F_0}(x_i) + n_i, \mathcal{F_0}(x_i)\}$

를 얻게됩니다.

이렇게 새롭게 구성된 dataset으로 다음 stage의 모델 $\mathcal{F_1}$ 을 학습하게 되면, ideal noisy-clean dataset과의 bias가 줄어들었기 때문에 더 좋은 denoising 성능을 가진 model을 학습할 수 있게 됩니다.

$\mathcal{F_{m+1}}←\{\mathcal{F_m}(x_i)+n_i,\mathcal{F_m}(x_i)\}$

이 과정을 반복적으로 수행하면 점점 더 data bias가 줄어든 새로운 dataset을 구성할 수 있고, denoising 성능도 따라서 향상되게 됩니다.

기존에 존재하는 iterative method와 달리, IDR은 inference 과정에서 한 번만 denoising을 하기 때문에 detail이 심하게 사라지는 문제점또한 없는것이 장점이라고 할 수 있습니다.

Fast Iterative Data Refinement

위에서 소개한 학습 방법론대로라면, 매 round 마다 초기화된 model을 새로 학습해야하기 떄문에 training time이 매우 오래걸린다는 단점이 있습니다.

논문에서는 이를 해결하기 위한 대안 2가지를 제시합니다.

1) One-epoch training

Refined dataset $\{\mathcal{F_m}(x_i)+n_i,\mathcal{F_m}(x_i)\}$ 에 대해서 1 epoch만 학습하는 방식입니다. dataset의 한 cycle이 끝나면, 그 당시의 model state를 고정시키고 next stage의 refined dataset 를 생성하고, $\mathcal{F_{m+1}}$ 를 학습할때 사용합니다.

이 방법을 사용하면 새로운 refined dataset을 생성하는 시간의 overhead가 생기지만, 이 시간은 전체 training time 대비 5% 정도만 차지하기 때문에 무시 가능한 수준입니다.

해당 방법을 사용하여 학습할 시 기존 full training 에서의 한 round 시간만에 전체 학습을 끝낼 수 있습니다.

2) Share model wieghts

두번째 방법은 model weight를 공유하는것입니다. 기존의 방법에서는 round가 새로 시작될때마다 model weight를 초기화하고 학습을 시작합니다.

따라서 각 round 마다 model이 안정적으로 학습되기 위한 epoch수가 많이 필요하게 됩니다.

그렇기 때문에 weight를 초기화하는 대신, 이전 round에서 model의 weight를 불러와 학습을 시작하면 각 round에서 학습이 완료되는데 필요한 epoch가 크게 줄어들게됩니다.

특히 one-epoch training에 대한 실험 결과 table을 보면, 1 round의 시간만에 full training과 비슷한 denoising 성능을 내는것을 확인할 수 있습니다.

Experiments

Gaussian, Binomial, Impulse, Correlated noise, Real noise에 대해서 실험을 하였습니다.

Gaussian

Binomial, Impulse (non zero-mean)

Correlated kernel

Real noise

Custom Data

gaussian noise를 더해서 학습하였을 때 train noise인 gaussian noise 뿐만 아니라 ISP noise와 JPEG compression artifact 등 다른 type의 noise 또한 denoising 되는것을 확인할 수 있습니다.

Conclusions

Denoising에서 큰 cost로 존재하는 dataset에 있어서 single noisy image만 가지고 학습을 진행할 수 있다는 점에서 큰 이점을 줍니다.

기존의 self-supervised denoising과 비교하여, n2n이나 n2c 학습과 비교해도 뒤떨어지지 않는 퍼포먼스를 보여주는것에서 의의가 있다고 볼 수 있습니다.

또한 self-supervised 방식으로 인해 black box로 학습이 되기 때문에, 의도하지 않았던 추가적인 noise pattern 또한 denoising이 되는 효과도 직접 진행한 실험에서 관찰되었습니다.

다만 Real noise에 있어서, 더해주는 synthetic noise와 real noise의 distribution이 많이 다를 경우 실제 target인 actual real noise에 대해서는 denoising 능력이 잘 학습되지 않는 경향이 있다는 한계가 있습니다.