本記事では2021年に論文発表されたImage Super-Resolution via Iterative Refinement(SR3)と呼ばれる超解像の手法をGoogle Colaboratoryで動かしながら、極力簡単に解説していきます。

超解像とは

概要

難しい言葉なしにざっくりとした理解としては、画像を表現する格子の細かさ(解像度)が低い画像(低解像度画像)を、格子がより細かい画像(高解像度画像)に変換する技術です。

もっと簡単な理解としては、ピンボケした画像を鮮明にする技術です。

左が低解像度の画像で、右が高解像度の画像です。左の画像は、解像度が小さく、高解像度画像と同じサイズまで引き延ばすと、格子一つあたりの面積が大きくなり、画像全体はぼやけてしまいます。

AIによる超解像では、このぼやけた画像から右の画像の精細な画像を生成します。

Image Super-Resolution via Iterative Refinement(SR3)とは

概要

AIによって超解像を実現する場合、一般的に2つの手法があります。
一つはGAN(敵対的生成ネットワーク)と、もう一方が、Diffusion Models(拡散モデル)です。GANの簡単な説明は、正誤判定を行うDiscreminatorと、Discreminatorに自分の作ったデータが正しいと思ってもらえるように頑張ってデータを作るGeneratorが戦って学習していく方式です。

出典: ディープラーニングによる画像の拡大技術(NTTPC COMMUNICATIONS)

上図の例でいえば、Generatorは自分が作った画像を、Discreminatorに正解データの高解像度画像だと思ってもらえるように頑張るわけですね。実に健気です。

このようなGANですが、トレーニングが不安定であるというデメリットが存在します。

このデメリットを解消する方式としてDiffusion Modelsがあります。このDiffusion Modelsは、高解像度画像にノイズを付与していきながら、徐々に画像を消去していきます。そして、画像がノイズのみになるまでの過程をトレーニングします。

そしてプロセスを逆転しモデルはノイズのみの画像から高解像度画像になるまでノイズを除去していくことを学習します。

こうして学習したモデルは、入力した低解像度画像を参考にしながら、ノイズから高解像度画像を生成します。
Image Super-Resolution via Iterative Refinementこちらの動画を見ていただくと、ノイズから高解像度画像を生成するというイメージをつかんでいただけるかと思います。

デモ(Colaboratory)

なかなか文章だけではイメージが掴みにくいものです。動かしてSR3を見ていきます。
ソースコードは全てColaboratoryで動作するようになっているので、ご自身でお試しいただけます。
ソースコードはこちらに記載しております。
動かす際にはColaboratoryで「ランタイムのタイプを変更」→「ハードウェアアクセラレータ」をGPUに変更してください。

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

それでは動かしていきます。

環境セットアップ

以下は、GPUの設定を確認、SR3のソースコードの取得、動作に必要なライブラリのインストールを行っています。

# カレントディレクトリ確認
!pwd
# GPU確認
!nvidia-smi --query-gpu=gpu_name,driver_version,memory.total --format=csv

# SR3のソースコードをGitHubから取得
%cd /content/
!git clone https://github.com/Janspiry/Image-Super-Resolution-via-Iterative-Refinement.git

# ライブラリインストール
!pip install tensorboardx wandb
!pip install lmdb
!pip install --upgrade gdown

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

次に学習済みモデルをダウンロードします。
一点注意点として、gdownを使ったダウンロードがサーバーの影響で失敗する場合があります。
その際には、https://drive.google.com/drive/folders/1mCiWhFqHyjt5zE4IdA41fjFwCYdqDzSFからI830000_E32_gen.pthとI830000_E32_opt.pthを直接ダウンロードして、/content/Image-Super-Resolution-via-Iterative-Refinementにアップロードしてください。

%cd /content/Image-Super-Resolution-via-Iterative-Refinement/
!gdown --id 1M015uQG8cq0mqwhEfc7klFJwDr2jsYSr
!gdown --id 1Vwe1Raa-Dpop75o5v7O5qvGoOdFDlyc-

# resume_stateにモデルを設定
!sed -i ./config/sr_sr3_64_512.json -e "s/\"resume_state\": null/\"resume_state\": \"I830000_E32\"/g"

デモファイルのアップロード

次に、超解像を行う低解像度画像をアップロードします。
低解像度画像は64ピクセルx64ピクセルである必要があるので、PEKO STEPやペイントなどで画像を作成してください。

今回はフリー素材の画像を切り抜いた以下の画像を使用しました。

# テストデータ用ディレクトリ作成
!mkdir input
%cd "/content/Image-Super-Resolution-via-Iterative-Refinement/input"

# ファイルアップロード
from google.colab import files
uploaded = files.upload()
uploaded = list(uploaded.keys())
print(uploaded)
file_name = uploaded[0]

%cd "/content/Image-Super-Resolution-via-Iterative-Refinement"

# データ準備
!mkdir -p ./dataset/celebahq

# ./dataset/celebahqにdata.mdbとlock.mdbを作成
!python data/prepare_data.py  --path ./input/  --out ./dataset/celebahq --size 64,512

超解像

最後に超解像を実行します。超解像を実行する前に、トレーニングを行うので約20分ほど時間を要します。

# 超解像実行
!python infer.py -c ./config/sr_sr3_64_512.json

実行結果は以下の通りです。

やや目元などに違和感があるものの、髪の毛の質感などは見事に精細になっています。
また、ノイズから生成された過程は以下の通りです。

まとめ

本記事では、超解像の概要やその実現手法、また、拡散モデルの超解像であるImage Super-Resolution via Iterative Refinement(SR3)によるデモを行いました。超解像は解像度の低い昔の写真や映像、医療用画像システムなどに利用されています。身近な例としてはモザイク除去に転用されたりしています。

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

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

参考文献

1.  論文 - Image Super-Resolution via Iterative Refinement

2. GitHub - Janspiry/Image-Super-Resolution-via-Iterative-Refinement