[NAFNet] 機械学習で画像のノイズ除去、ブラー除去、超解像 [Denoise, Deblur, Super Resolution]

2022年5月7日土曜日

Artificial Intelligence

本記事では、NAFNetと呼ばれる機械学習手法を用いて、画像のノイズ除去、ブラー除去、超解像を行う方法をご紹介します。

アイキャッチ
出典: megvii-research/NAFNet

NAFNet

概要

近年、画像復元(Image Resolution)技術は様々な進歩がみられていますが、システムの複雑さも増加しています。

NAFNetは、SOTAを達成しつつ、計算効率の高い単純なベースラインで構成される画像復元技術です。

NAFNetは、Sigmoid, ReLU, GELU, Softmaxなどの非線形活性化関数を使用せず、乗算で置き換えるか、削除することでベースラインを単純化しています。

このため、非線形活性化関数フリーなネットワークであることから、Nonlinear Activation Free NetworkでNAFNetと呼ばれています。

NAFNetはblur除去のデータセットなどでSOTAを達成しながら、従来技術より計算コストを大きく削減しています。

アイキャッチ
出典: megvii-research/NAFNet

詳細はこちらの論文をご参照ください。

本記事では上記手法を用いて、ノイズ除去、ブラー除去、超解像を行います。

デモ(Colaboratory)

それでは、実際に動かしながらImage Reconstructionを行っていきます。
ソースコードは本記事にも記載していますが、下記のGitHubでも取得可能です。
GitHub - Colaboratory demo

また、下記から直接Google Colaboratoryで開くこともできます。
Open In Colab

なお、このデモはPythonで実装しています。
Pythonの実装に不安がある方、Pythonを使った機械学習について詳しく勉強したい方は、以下の書籍などがおすすめです。

環境セットアップ

それではセットアップしていきます。 Colaboratoryを開いたら下記を設定しGPUを使用するようにしてください。

「ランタイムのタイプを変更」→「ハードウェアアクセラレータ」をGPUに変更

初めに、論文発表元のGithubからソースコードを取得します

%cd /content

!git clone https://github.com/megvii-research/NAFNet

次にライブラリをインストールします。

%cd /content/NAFNet

!pip install -r requirements.txt
!pip install --upgrade --no-cache-dir gdown
!python3 setup.py develop --no_cuda_ext

環境セットアップの最後にライブラリをインポートします。

%cd /content/NAFNet

import gdown

import torch

from basicsr.models import create_model
from basicsr.utils import img2tensor as _img2tensor, tensor2img, imwrite
from basicsr.utils.options import parse
import numpy as np
import cv2
import matplotlib.pyplot as plt

import os
from google.colab import files
import shutil
import glob

以上で環境セットアップは完了です。

学習済みモデルのダウンロード

続いてgdownを用いてGoogle Driveから学習済みモデルをダウンロードします。

%cd /content/NAFNet

# Denoise
if not os.path.exists("./experiments/pretrained_models/NAFNet-SIDD-width64.pth"):
  gdown.download('https://drive.google.com/uc?id=14Fht1QQJ2gMlk4N1ERCRuElg8JfjrWWR', "./experiments/pretrained_models/", quiet=False)
# Deblur
if not os.path.exists("./experiments/pretrained_models/NAFNet-REDS-width64.pth"):
  gdown.download('https://drive.google.com/uc?id=14D4V4raNYIOhETfcuuLI3bGLB-OYIv6X', "./experiments/pretrained_models/", quiet=False)
# Super Resolution
if not os.path.exists("./experiments/pretrained_models/NAFSSR-L_4x.pth"):
  gdown.download('https://drive.google.com/uc?id=1TIdQhPtBrZb2wrBdAp9l8NHINLeExOwb', "./experiments/pretrained_models/", quiet=False)

ノイズ除去(Denoise)

それでは、まずノイズ除去から行います。

はじめにモデルをビルドします。

%cd /content/NAFNet

opt_path = 'options/test/SIDD/NAFNet-width64.yml'
opt = parse(opt_path, is_train=False)
opt['dist'] = False
NAFNet = create_model(opt)

ノイズ除去を実施します。

%cd /content/NAFNet

input_list = sorted(glob.glob(os.path.join(denoise_upload_folder, '*')))

for input_path in input_list:
  img_input = imread(input_path)
  inp = img2tensor(img_input)
  output_path = os.path.join(denoise_result_folder, os.path.basename(input_path))
  single_image_inference(NAFNet, inp, output_path)

output_list = sorted(glob.glob(os.path.join(denoise_result_folder, '*')))
for input_path, output_path in zip(input_list, output_list):
  img_input = imread(input_path)
  img_output = imread(output_path)
  display(img_input, img_output)

出力結果は以下の通りです。
カラーノイズ画像のノイズ除去結果

カラーノイズ結果

luminanceノイズのノイズ除去結果

luminance noise結果

トレーニングデータの影響か、カラーノイズ画像のノイズ除去は良好ですが、luminanceノイズではうまくノイズ除去できていません。
現状、ノイズの得手不得手がありそうです。

ブラー除去(Deblur)

続いて、動いている被写体などを撮影した際に発生するブレを軽減するブラー除去を試してみます。

%cd /content/NAFNet

opt_path = 'options/test/REDS/NAFNet-width64.yml'
opt = parse(opt_path, is_train=False)
opt['dist'] = False
NAFNet = create_model(opt)

モデルをビルド後ブラー除去を実行します。

%cd /content/NAFNet

input_list = sorted(glob.glob(os.path.join(deblur_upload_folder, '*')))

for input_path in input_list:
  img_input = imread(input_path)
  inp = img2tensor(img_input)
  output_path = os.path.join(deblur_result_folder, os.path.basename(input_path))
  single_image_inference(NAFNet, inp, output_path)

output_list = sorted(glob.glob(os.path.join(deblur_result_folder, '*')))
for input_path, output_path in zip(input_list, output_list):
  img_input = imread(input_path)
  img_output = imread(output_path)
  display(img_input, img_output)

出力結果は以下の通りです。

blur除去結果

こちらは自然にブレが除去されています。

超解像(Super Resolution)

最後に超解像を試してみます。

本来学習済みモデルとしては、ステレオ超解像用にトレーニングされているモデルですが、単一画像を超解像します。

opt_path = 'options/test/NAFSSR/NAFSSR-L_4x.yml'
opt = parse(opt_path, is_train=False)
opt['dist'] = False
NAFSSR = create_model(opt)

モデルのビルド後超解像を実施します。

%cd /content/NAFNet

input_list = sorted(glob.glob(os.path.join(sr_upload_folder, '*')))

for input_path in input_list:
  img_l = imread(input_path)
  inp_l = img2tensor(img_l)
  img_r = imread(input_path)
  inp_r = img2tensor(img_r)
  stereo_image_inference(NAFSSR, inp_l, inp_r, os.path.join(sr_result_folder, "sr_result{}.png"))

output_list = sorted(glob.glob(os.path.join(sr_result_folder, '*')))
for input_path, output_path in zip(input_list, output_list):
  img_input = imread(input_path)
  img_output = imread(output_path)
  display(img_input, img_output)

出力結果は以下の通りです。

SR結果

顔周辺などは自然な仕上がりです。

まとめ

本記事では、NAFNetを用いた画像のノイズ除去、ブラー除去、超解像をご紹介しました。
ネットワークの単純さからか推論時間が短い点もメリットの一つと言えそうです。

また本記事では、機械学習を動かすことにフォーカスしてご紹介しました。
もう少し学術的に体系立てて学びたいという方には以下の書籍などがお勧めです。ぜひご一読下さい。


また動かせるだけから理解して応用できるエンジニアの足掛かりに下記のUdemyなどもお勧めです。

参考文献

1.  論文 - Simple Baselines for Image Restoration

2. GitHub - megvii-research/NAFNet

プロフィール

自分の写真
製造業に勤務する傍ら、日々AIの技術動向を調査しブログにアウトプットしています。 AIに関するご相談やお仕事のご依頼はブログのお問い合わせフォームか以下のアドレスまでお気軽にお問い合わせください。 bhupb13511@yahoo.co.jp

カテゴリ

このブログを検索

ブログ アーカイブ

TeDokology