[論文レビュー] SUNet: Swin Transformer UNet for Image Denoising

Chi-Mao Fan, Tsung-Jung Liu+1 more|arXiv (Cornell University)|Feb 28, 2022

Image and Signal Denoising Methods被引用数 4

ひとこと要約

本論文では、長距離依存関係とグローバルコンテキストを捉えるために従来の畳み込み層の代わりにSwin Transformerブロックを組み込んだ、Swin TransformerベースのU-NetアーキテクチャであるSUNetを提案する。モデルはベンチマークデータセットで最先端の性能を達成し、PSNRとSSIMが向上している。また、チェス盤アーチファクトを低減し、再構築品質を向上させるために、新たなデュアルアップサンプリングブロックを導入している。

ABSTRACT

Image restoration is a challenging ill-posed problem which also has been a long-standing issue. In the past few years, the convolution neural networks (CNNs) almost dominated the computer vision and had achieved considerable success in different levels of vision tasks including image restoration. However, recently the Swin Transformer-based model also shows impressive performance, even surpasses the CNN-based methods to become the state-of-the-art on high-level vision tasks. In this paper, we proposed a restoration model called SUNet which uses the Swin Transformer layer as our basic block and then is applied to UNet architecture for image denoising. The source code and pre-trained models are available at https://github.com/FanChiMao/SUNet.

研究の動機と目的

• 畳み込みニューラルネットワーク（CNN）ベースの画像ノイズ除去モデルの限界、すなわち局所的特徴依存性と固定カーネルの適用を是正すること。
• 特にグローバルコンテキストモデリングが重要な画像ノイズ除去という低レベルビジョンタスクにおいて、Swin Transformerの潜在的可能性を検証すること。
• 特徴表現と修復忠実度を向上させるために、Swin Transformerブロックで強化されたU-Netアーキテクチャを設計すること。
• サブピクセルとバイリニアアップサンプリングを組み合わせたデュアルアップサンプリングブロックを導入することで、特徴アップサンプリングにおけるチェス盤アーチファクトを軽減すること。
• 標準ベンチマークで検証されたU-NetフレームワークにSwin Transformerバックボーンを組み込んだことで、画像ノイズ除去の新しい最先端水準を確立すること。

提案手法

• SUNetは、U-Netのエンコーダーとデコーダーにおける標準的な畳み込み層を、長距離依存関係をモデル化しグローバルコンテキストを捉えるためにSwin Transformerブロックに置き換える。
• 主なSwin TransformerベースのU-Netエンコーダー・デコーダーパathに渡す前に、初期特徴を抽出するために浅い3×3畳み込み層を用いる。
• チェス盤アーチファクトを低減し、特徴再構築を向上させるために、サブピクセルとバイリニアアップサンプリングを組み合わせたデュアルアップサンプリングブロックを導入する。
• Swin Transformerブロックはシフトされたウィンドウを用いることで計算複雑度を低減しつつ、局所的およびグローバルな受容 field における自己注意を維持する。
• ネットワークは、DIV2Kデータセットから抽出したパッチを用いて、さまざまなレベルの加法的白色ガウスノイズ（AWGN）を加えた状態で、平均二乗誤差（MSE）損失を用いてエンドツーエンドで訓練される。
• 最終出力は3×3畳み込み層を介して再構築され、ノイズ除去画像が得られる。

実験結果

リサーチクエスチョン

• RQ1 Swin TransformerベースのU-Netアーキテクチャは、既存のCNNベースおよびU-Netベースのモデルを上回る性能を達成できるか？
• RQ2 Swin TransformerをU-Netアーキテクチャに統合することで、特徴表現と修復品質にどのような影響を与えるか？
• RQ3 提案されたデュアルアップサンプリングブロックは、転置畳み込みによるアップサンプリングと比較して、チェス盤アーチファクトを効果的に低減できるか？
• RQ4 既存のU-Netバリアントと比較して、より少ないパラメータ数とFLOPsで競争力のある性能を達成できるか？
• RQ5 異なるノイズレベルおよび実世界の画像データセットにおいて、モデルの一般化性能はどのように評価されるか？

主な発見

• CBSD68データセット（σ=50）において、SUNetはPSNR 36.79 dB、SSIM 0.953を達成し、比較されたすべてのCNNベースおよびU-Netベースのモデルを上回った。
• Kodak24データセット（σ=50）において、SUNetはPSNR 29.54 dB、SSIM 0.810を達成し、多様な画像セットにわたる優れた一般化性能を示した。
• パラメータ数が9900万（多くのベースラインと比較して高い）であるにもかかわらず、DHDNおよびRDUNetと比較してFLOPsは3%削減、パラメータ数は60%削減されており、効率性の向上を示している。
• デュアルアップサンプリングブロックは、視覚的比較により、滑らかなテクスチャと鋭いエッジを示すことで、チェス盤アーチファクトを顕著に低減していることが裏付けられた。
• CBSD68データセット（σ=50）において、SUNetは最高のSSIM（0.958）を達成し、優れた知覚的品質と構造的忠実度を示した。
• すべてのノイズレベル（σ=10, 30, 50）において、SUNetは従来手法よりも一貫したPSNRおよびSSIMの向上を維持していた。