【Denoising】SCUNetによるブラインドノイズ除去を試してみる

2022年3月に発表されたSwin-Conv-UNet（SCUNet）によるブラインドノイズ除去を紹介します。

実際に画像からブラインドノイズ除去を試してみましょう。

Google colabを使用して簡単に実装することができますので、ぜひ最後までご覧ください。

SCUNetは

Swin-Conv-UNet (SCUNet)はノイズ除去ネットワークのアーキテクチャ（Swin-Convブロック）と新しいデータ合成方法により、ノイズ除去を実現します。

まずは以下の画像をご覧ください。

SCUNetにり、ノイズの多い画像に対して適用できることがわかります。

局所的なモデリング能力を取り入れた、新しいSWINCONVブロックを採用しています。
残差畳み込みブロックと非局所的なモデリング能力を持つスウィントランスフォーマーブロックをメインビルディングとしてUNetに組み込むことで、局所・非局所性をさらに強化することができます。

データ合成の設計を行い、異なる種類のノイズを考慮したパイプラインにより、学習されたディープブラインドを実際のノイズが多い画像に対して適用しています。

SCUNetの導入

早速、SCUNetを使用していきましょう。

以下、Google colab環境で進めていきます。

まずはGPUを使用できるように設定をします。

今回紹介するコードは以下のボタンからコピーして使用していただくことも可能です。

from google.colab import drive
drive.mount('/content/drive')
%cd ./drive/MyDrive

公式よりcloneしてきます。

!git clone https://github.com/cszn/scunet
%cd scunet

必要なライブラリをインストールします。

!pip install thop
!pip install einops
!pip install timm

以上で導入が完了しました。

デモ

モデルをダウンロードする

まずは学習済モデルをダウンロードします。

ダウンロードのためのファイルが用意されていますので、実行します。

!python main_download_pretrained_models.py --models "SCUNet" --model_dir "model_zoo"

Gaussian denoising（グレースケール画像）

まずはガウスノイズの除去から見ていきます。

# grayscale images
!python main_test_scunet_gray_gaussian.py --model_name scunet_gray_25 --noise_level_img 25 --testset_name set12

Gaussian denoising（カラー画像）

# color images
!python main_test_scunet_color_gaussian.py --model_name scunet_color_25 --noise_level_img 25 --testset_name bsd68

Blind real image denoising

次に実際の画像でノイズ除去を行ってみます。

# Blind real image denoising
!python main_test_scunet_real_application.py --model_name scunet_color_real_psnr --testset_name real3

処理前後の画像を比較すると結果は一目瞭然です。

オリジナル動画でテスト

次にオリジナル画像でも試してみましょう。

「scunet/testsets/set_main」というフォルダを作成して、画像を格納します。

画像サイズが大きくメモリ不足となる場合には、サイズダウンをしておきましょう。

%cd testsets/set_main

import glob
from PIL import Image

files = glob.glob("*.JPG")
for f in files:
    im = Image.open(f)
    w,h = im.size
    im = im.resize((int(w/4),int(h/4)))
    im.save(f)

ノイズ除去を実行してみましょう。

!python main_test_scunet_real_application.py --model_name scunet_color_real_psnr --testset_name set_main

こちらの画像でもノイズ除去ができました。

まとめ

最後までご覧いただきありがとうございました。

今回はSwin-Conv-UNet（SCUNet）によるノイズ除去を紹介しました。

ブラインドノイズ除去は実用性も高いことから、更なる性能の向上が期待されます。