LeapMindのいまむらです。

画像復元、してますか？

画像復元というと超解像やインペインティング、着色など様々なタスクがありますが、特に深層学習を使って画像からノイズを取り除きたくなること、長い人生で、たまにはありますよね。今日はそんな話をしたいと思います。

はじめに

そもそも画像におけるノイズ除去とは何かというと、(加法性ノイズの場合)以下の式で定式化される劣化モデルに従って観測されるノイズを含む画像 $\boldsymbol{y}$ からノイズのないクリーンな画像 $\boldsymbol{x}$ を復元(予測)することを目的としたタスクです。

\begin{equation} \boldsymbol{y=Dx+n} \tag{1} \end{equation}

ここで、 $D$ はボケや欠損などの劣化過程をあらわします。 $\boldsymbol{n}$ は加法性ノイズで、平均 $0$ かつ標準偏差 $\sigma$ のガウシアンノイズを仮定することが多いですが、必ずしもその限りではありません。

劣化モデル(1)に対し、画像復元モデルは以下のように定式化されます。

\begin{equation} \boldsymbol{x'}=\underset{\boldsymbol{x}}{\rm arg\hspace{0.17em}min}\hspace{0.17em}E(\boldsymbol{x}; \boldsymbol{y})+\lambda R(\boldsymbol{x}) \tag{2} \end{equation}

ここで、 $E(\boldsymbol{x}; \boldsymbol{y})$ はデータ忠実項であり、復元画像 $\boldsymbol{x'}$ が観測画像 $\boldsymbol{y}$ から離れすぎないようにするための役割を果たします。 $R(\cdot)$ は正則化項であり、元画像に対する事前知識(Prior)をここに組み込むことで、復元画像に対してその性質を反映させることができます。 $\lambda$ はデータ忠実項と正則化項に関するトレードオフパラメータです。正則化項を設定しない場合(もしくは $\lambda=0$ のとき)、当然ながら復元画像 $\boldsymbol{x'}$ は観測画像 $\boldsymbol{y}$ と同一のものとなってしまうため、精度良く復元するためにも元画像に対するPriorを考慮した正則化項の設計が重要です。伝統的には、全変動(TV)[1]などがよく用いられてきました。

こうしたモデルベースのアプローチ以外にも、様々な手法が提案されています。ノイズ除去界隈では、BM3D[2]という手法が長らく最先端の性能を誇る手法として知られていました。BM3Dは、画像の非局所領域に対して類似度の高い領域を集めて3次元のブロックを作り、そのブロックに対する直交変換および係数に対するシュリンケージによってノイズを除去します（図1）。これは、自然画像中には類似度の高い領域が複数存在する(自己相似性が高い)というPriorに基づいた手法であるといえます。

BM3D

図1: BM3Dによる非局所領域でのブロックマッチングの例（[2]より引用）

深層学習によるアプローチ

クリーン画像を用いた教師あり学習によるノイズ除去

BM3DがSoTAとして君臨する中、深層学習技術の発展に伴い、深層学習を用いたノイズ除去に関する研究が盛り上がってきました[3-5]。その中でも、ノイズ除去界でのブレイクスルーといえる手法が、Zhangらの提案したDnCNN[6]と呼ばれる手法です。

DnCNNは、畳み込み層とバッチノーマライゼーション、ReLUからなるネットワークを用いてノイズを含んだ画像とノイズのないクリーン画像ペアによる教師あり学習で訓練し、BM3Dを超える性能を発揮しました（図2）。ネットワーク構造的にはなんの変哲もないシンプルなものですが、綺麗な画像を出力するよう学習させるのではなく、画像に含まれるノイズを出力するよう学習させるという工夫によって性能を向上させています。これは、綺麗な画像を出力するために入力画像の複雑な構造を中間層内で保持するよりもむしろ、画像の構造を除去していき、潜在的なノイズだけを残していく方が学習器にとって容易だという主張に基づいています。

DnCNN

図2: DnCNNのアーキテクチャ（[6]より引用）

DnCNNは特定のノイズレベルのガウシアンノイズに対してだけではなく、ノイズレベルが未知(ブラインドノイズ)の画像に対しても良好な性能を示しました。また、ガウシアンノイズ除去だけでなく、JPEGノイズ除去や超解像といったタスクにも同じ学習方法を用いて当時の最先端の性能を発揮することが示されました。さらに、これら三つのタスクを一つのモデルで学習したとしてもなお最先端の性能を示すなど、ノイズ除去だけにとどまらない汎用的な画像復元手法となっています。

DnCNNの登場を皮切りに、次々と教師あり学習手法[7-8]が提案されては、ベンチマークセットに対するスコアを次々と塗り替えていったのですが、実際にノイズ除去を行いたいシチュエーションでは依然としてBM3Dが重宝されていました。既に述べた教師あり学習ベースの手法はノイズの種類や分布をPriorとして与えて学習しているとみなすと、実際に直面する分布のわからないノイズに対しては良好な復元性能が期待できなかったからです。また、実際の入力画像と教師画像のペアを用意するために短時間露光および長時間露光での撮影をする必要があり、これらの手法によるモデルを学習するために必要な数百枚の画像からなるデータセットを人手で作成するのは現実的でないという問題がありました。

クリーン画像なしでの学習によるノイズ除去

Deep Image Prior

Ulyanovら[9]は、CNNそのものが画像に対するPriorだという主張のもと、ノイズを含む画像一枚のみを用いて学習しノイズ除去を実現するDeep Image Priorとよばれる手法を提案しました。具体的には、乱数によって生成された画像 $\boldsymbol{z}$ 一枚を入力とし、ノイズを含む画像を教師画像として学習を行います。この学習に用いる損失関数 $l_{dip}$ を式であらわすと以下のようになります。

\begin{equation} l_{dip} = E(\phi(\boldsymbol{z}); \boldsymbol{y}) \tag{3} \end{equation}

ここで、 $\phi$ はCNNをあらわします。また、復元画像 $\boldsymbol{x'}$ は

\begin{equation} \boldsymbol{x'}=\phi(\boldsymbol{z}) \tag{4} \end{equation}

として得ることができます。

冒頭の画像復元モデルの話を踏まえると、 $l_{dip}$ にはデータ忠実項しか含まれないため、復元画像 $\boldsymbol{x'}$ は $\boldsymbol{y}$ と同じものになりそうです。実際に、 $l_{dip}$ を最小化する $\phi$ を用いた時には $\boldsymbol{x}'\approx \boldsymbol{y}$ となる画像 $\boldsymbol{x'}$ が得られます。しかし、 $l_{dip}$ を最小化しきってしまう前の $\phi$ を用いると、ノイズを含まない滑らかな画像が得られます（図3）。これは、CNNがノイズよりも滑らかな成分を学習しやすいことを利用しており、「ノイズ除去後の画像は滑らかである」というPriorに基づいた手法であるとみなせば、CNNは式(2)における正則化項の役割を果たしているといえます。学習が進むにつれてデータ忠実項の影響が大きくなるのは、ネットワークが元画像の細部を表現できるようになることで正則化が弱まる、すなわちトレードオフパラメータ $\lambda$ が0へと変化していると考えると感覚的にはわかりやすいかもしれません。

Deep Image Prior

図3: Deep Image Priorによる学習の経過（[9]より引用）

Deep Image Priorは式(2)に基づいたモデルベースアプローチであり、ノイズの分布に関する事前知識なしに劣化画像一枚から復元することが可能な手法ですが、他にもノイズ画像-クリーン画像のペアを必要としない手法が提案されています。

Noise2Noise

Noise2Noise[10]は、ノイズにより劣化した画像セットの期待値がクリーン画像と一致する時、その劣化画像セット中の任意のペアを入力-教師データとして用いて(つまりクリーンな画像を用いずに)学習できることを示しました。これは、 $l_2$ ロスの最小化によって訓練されたネットワークは、元画像として想定されうる画像セットからなる分布の期待値を復元画像として出力する、という洞察に基づいたものです。この手法は、ノイズ画像-クリーン画像のペアでの学習と比較して同等かそれを上回る復元性能を発揮しました。Noise2Noiseのアイデアは、同一シーンを写した独立したノイズを含む二枚以上の劣化画像があるとき、学習の際にクリーン画像が不要となるということを示しています。

著者らは、一例として、長時間露光によって撮影されたノイズのない写真の各画素値は、短時間露光で撮影された独立したノイズを含んだ画像セットの各画素の期待値である、と述べています。短時間露光画像と長時間露光画像のペアを用意するよりも、短時間露光画像を複数枚用意する方が遥かに簡単であるため、本手法は実際のノイズ除去のシチュエーションにおいて実用的なアイデアだといえます。

Noise2Void

Noise2Noiseの学習には同一シーンかつ独立したノイズがのった2枚以上の画像セット $(\boldsymbol{x}+\boldsymbol{n}_1, \boldsymbol{x}+\boldsymbol{n}_2, ...)$ が必要ですが、動的な対象物の撮影などで同一の $\boldsymbol{x}$ を得ることは難しい場合があります。Noise2Void[11]は、ノイズを含む画像一枚からノイズ除去を学習する自己教師あり学習手法として提案されました。Noise2Voidは、入力画像と教師画像に同一の劣化画像を用います。ただし、そのままだとネットワークは恒等写像を学習してしまうので、そうならないように受容野に死角を作るという工夫が施されています。具体的には、学習のために切り出されたパッチの中央の画素をマスクし、その周囲の画素値からマスクされた画素値を推定するといったアイデアが提案されています（図4）。実際の実装では、パッチ内の一定の割合の画素値をその近傍の画素値で置き換えることでマスク処理を実現しています。Noise2Voidは、各画素が隣接画素と高い相関を持つことと、ノイズが各画素に対して独立に発生しているというPriorを考慮したモデルであるとみなせます。

Noise2Void

図4: Noise2Voidのアイデア（[11]より引用）。(a)のように注目画素を含む領域をもとに注目画素を予測した場合は恒等写像を学習してしまうが、(b)のように注目画素についてマスクを行うことでそれを防ぎ、周囲の画素から注目画素の値を予測するモデルを学習する。

近年のトレンド

ここまで深層学習を用いた代表的なノイズ除去手法を紹介しましたが、この他にも新しい手法[12-13]が次々と提案されています。こうして手法の発展の流れを追ってみると、大量のデータセットを用いて画像の性質を学習するよりもむしろ、Priorを適切に設定して少数または一枚の劣化画像から復元を行う復元モデルベースのアプローチに回帰しているように感じます。ノイズ除去タスクでは、未知のノイズに対する復元性能の高さから復元モデルベースのアプローチが好まれてきましたが、こうした傾向に伴い、実問題へのアプローチにおいても深層学習を用いた手法の存在感が増してきているように思います。

まとめ

本記事では、画像復元、特に深層学習を用いたノイズ除去に関する手法について紹介しました。ノイズ除去は信号処理/画像処理における基本的なタスクであり、ノイズが全く載らない画期的なセンサが開発されるまでは必要とされ続ける技術であることは間違いありません。アプローチとしては、今後も深層学習ベースの手法が最先端の性能を達成し続けるのではないかと個人的には思っています。ノイズ除去の技術はデジカメや監視カメラのような多くのエッジデバイスに搭載されており、さらなる性能向上のためにこうした深層学習ベースの手法の需要が高まっていくことが大いに予想されます。しかし、エッジデバイスでは、画質だけでなく低消費電力かつリアルタイム性が求められるため、計算コストの高い深層学習ベースのモデルとの相性はあまりよくありません。

そこで、LeapMindでは、深層学習モデルの精度低下を抑えつつ軽量化する極小量子化技術を活用し、エッジデバイス上で低消費電力かつリアルタイムで動作する深層学習モデルの導入を推し進めています。極小量子化技術を使ってエッジデバイス上で画像からノイズを除去したい方もそうでない方も、様々な職種で採用活動を実施しておりますので、興味がある方はぜひ一度LeapMindの採用ページから応募してみてください。いきなり面接はちょっと、という人はカジュアル面談から始めることもできるとのことですので、お気軽にお問い合わせください！

参考文献

[1] Rudin, Leonid I., Stanley Osher, and Emad Fatemi. "Nonlinear total variation based noise removal algorithms." Physica D: nonlinear phenomena 60.1-4 (1992): 259-268.

[2] Dabov, Kostadin, et al. "Image denoising by sparse 3-D transform-domain collaborative filtering." IEEE Transactions on image processing 16.8 (2007): 2080-2095.

[3] Burger, Harold C., Christian J. Schuler, and Stefan Harmeling. "Image denoising: Can plain neural networks compete with BM3D?." 2012 IEEE conference on computer vision and pattern recognition. IEEE, 2012.

[4] Xie, Junyuan, Linli Xu, and Enhong Chen. "Image denoising and inpainting with deep neural networks." Advances in neural information processing systems. 2012.

[5] Chen, Yunjin, and Thomas Pock. "Trainable nonlinear reaction diffusion: A flexible framework for fast and effective image restoration." IEEE transactions on pattern analysis and machine intelligence 39.6 (2016): 1256-1272.

[6] Zhang, Kai, et al. "Beyond a gaussian denoiser: Residual learning of deep cnn for image denoising." IEEE transactions on image processing 26.7 (2017): 3142-3155.

[7] Bae, Woong, Jaejun Yoo, and Jong Chul Ye. "Beyond deep residual learning for image restoration: Persistent homology-guided manifold simplification." Proceedings of the IEEE conference on computer vision and pattern recognition workshops. 2017.

[8] Liu, Pengju, et al. "Multi-level wavelet-CNN for image restoration." Proceedings of the IEEE conference on computer vision and pattern recognition workshops. 2018.

[9] Ulyanov, Dmitry, Andrea Vedaldi, and Victor Lempitsky. "Deep image prior." Proceedings of the IEEE conference on computer vision and pattern recognition. 2018.

[10] Lehtinen, Jaakko, et al. "Noise2noise." International Conference on Machine Learning. PMLR, 2018.

[11] Krull, Alexander, Tim-Oliver Buchholz, and Florian Jug. "Noise2void-learning denoising from single noisy images." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2019.

[12] Quan, Yuhui, et al. "Self2self with dropout: Learning self-supervised denoising from single image." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2020.

[13] Jo, Yeonsik, Se Young Chun, and Jonghyun Choi. "Rethinking deep image prior for denoising." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2021.

Tech Blog