本記事では、NAFNetと呼ばれる機械学習手法を用いて、画像のノイズ除去、ブラー除去、超解像を行う方法をご紹介します。
NAFNet
概要
近年、画像復元(Image
Resolution)技術は様々な進歩がみられていますが、システムの複雑さも増加しています。
NAFNetは、SOTAを達成しつつ、計算効率の高い単純なベースラインで構成される画像復元技術です。
NAFNetは、Sigmoid, ReLU, GELU,
Softmaxなどの非線形活性化関数を使用せず、乗算で置き換えるか、削除することでベースラインを単純化しています。
このため、非線形活性化関数フリーなネットワークであることから、Nonlinear Activation Free NetworkでNAFNetと呼ばれています。
NAFNetはblur除去のデータセットなどでSOTAを達成しながら、従来技術より計算コストを大きく削減しています。
詳細はこちらの論文をご参照ください。
なお、以下の記事では有償にはなりますが、詳細な技術解説、モデルのトレーニング、推論方法などを記載しております。
よろしければご参照ください。
NAFNetで画像からノイズを除去するレシピ
本レシピでは、NAFNetの技術概要を解説し、画像のノイズ除去を行う方法や、モデルのトレーニング方法をご紹介します。
本記事では上記手法を用いて、ノイズ除去、ブラー除去、超解像を行います。
デモ(Colaboratory)
それでは、実際に動かしながらImage Reconstructionを行っていきます。
ソースコードは本記事にも記載していますが、下記のGitHubでも取得可能です。
GitHub - Colaboratory demo
また、下記から直接Google Colaboratoryで開くこともできます。
また、このデモは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ノイズではうまくノイズ除去できていません。
現状、ノイズの得手不得手がありそうです。
ブラー除去(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)
出力結果は以下の通りです。
こちらは自然にブレが除去されています。
超解像(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)
出力結果は以下の通りです。
顔周辺などは自然な仕上がりです。
まとめ
本記事では、NAFNetを用いた画像のノイズ除去、ブラー除去、超解像をご紹介しました。
ネットワークの単純さからか推論時間が短い点もメリットの一つと言えそうです。
また本記事では、機械学習を動かすことにフォーカスしてご紹介しました。
もう少し学術的に体系立てて学びたいという方には以下の書籍などがお勧めです。ぜひご一読下さい。
リンク
リンク
また動かせるだけから理解して応用できるエンジニアの足掛かりに下記のUdemyなどもお勧めです。
参考文献
1.
論文 - Simple Baselines for Image Restoration
2. GitHub - megvii-research/NAFNet
論文を見たら全部英語で読めませんでした。せめて読めるちゃんとした正式な論文がアレばたすかります。
返信削除申し訳ございませんが、正式な論文がリンクに記載しています英語の論文です。
削除現状日本語の論文はありません。
解説ありがとうございます、Googlecolabで実行すると「RuntimeError: CUDA out of memory. 」となり、GPUメモリがオーバーしてしまうようです。実行は有料版でされていたりしますでしょうか。
返信削除動作確認は無料版のGoogle Colabで確認しています。
削除おそらくですが、入力の画像サイズが大きく、GPUメモリがオーバーしている可能性があります。お手数ですが、一度画像サイズを小さくしてお試しください。