本記事では、SR3(Image Super-Resolution via Iterative Refinement)と呼ばれる機械学習手法を用いて、画像の解像度を上げる方法を紹介します。

SR3とは

SR3はRepeated Refinementによる超解像手法です。

SR3は、画像生成時にノイズ除去プロセスを適用しています。
推論時には、ガウスノイズなど様々なノイズ除去に関してトレーニングされたU-Netモデルを使用して、ノイズの多い出力を繰り返し学習しています。
SR3は顔や、自然画像など、様々な倍率での超解像タスクで強力なパフォーマンスを示すとされています。

左:低解像度画像(64x64)、右:超解像後画像(512x512) 出典: Image-Super-Resolution-via-Iterative-Refinement

左の解像度の低い入力画像を、自然に高解像度化していますね。

SR3の導入手順

セットアップ1: conda環境構築

それでは早速、開発環境にSR3をセットアップしていきます。
動作確認は下記の環境で行っています。

GPUをお持ちでない場合は以下のColaboratory(ブラウザから利用する無料GPU環境)のセットアップ手順をご確認ください。

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

それでは、SR3をインストールしていきます。他の機械学習環境に影響を与えないためにMinicondaの仮想環境上に構築していきます。Minicondaのインストール手順は公式ドキュメントをご参照ください。

conda環境構築手順を下記に記載します。

# python version 3.7のconda環境を作成
$ conda create -n sr3 python=3.7 -y

# 作成した環境をアクティベート
$ conda activate sr3
$ cd sr3

# Proxy配下の場合は事前に下記を設定
$ export http_proxy="http://"username":"password"@proxy:port"
$ export https_proxy="http://"username":"password"@proxy:port"

# gitからcodeをclone
$ git clone https://github.com/Janspiry/Image-Super-Resolution-via-Iterative-Refinement.git
$ cd Image-Super-Resolution-via-Iterative-Refinement/

# 関連するライブラリをインストール
$ conda install -c conda-forge tensorboardx
$ conda install -c conda-forge opencv
$ conda install python-lmdb tqdm
# https://pytorch.org/get-started/previous-versions/ からcudaバージョンが一致するものをインストール
# conda install pytorch==1.10 torchvision torchaudio cudatoolkit=11.3 -c pytorch

以上でconda環境構築は完了です。

セットアップ2: 学習済みモデルとデータ準備

続いて学習済みモデルと、超解像を試すためのデータを準備していきます。
学習済みモデルは公式のGithubから、確認用のデータは任意の512ピクセルx512ピクセル顔画像を使用します。
始めに、学習済みモデルをセットアップしていきます。

$ cd Image-Super-Resolution-via-Iterative-Refinement

$ mkdir pretrain
# https://drive.google.com/drive/folders/1mCiWhFqHyjt5zE4IdA41fjFwCYdqDzSF?usp=sharing からモデルダウンロード
# pretrainにI830000_E32_gen.pthとI830000_E32_opt.pthを配置

# config/sr_sr3_64_512.jsonのresumeに設定
E.g. "resume_state" : "./pretrain/I830000_E32"

以上で学習済みモデルのセットアップは完了です。
続いてデータのセットアップを行います。
本記事では、下記のロイヤリティフリーの画像を使用させていただきます。

上記の画像の顔部分をペイントなどで切り取ります。
その後一度64ピクセルx64ピクセル程度まで縮小させた後、512ピクセルx512ピクセルまで拡大し
意図的に下記のような低解像度の画像を生成します。

このような手順で生成した512ピクセルの低解像度画像を先ほどセットアップした環境に配置します。
手順は以下の通りです。

$ mkdir input
# inputに64x64ピクセルの画像を配置

# データ準備
$ python data/prepare_data.py --path ./input/ --size 64,512 --out ./dataset/celebahq

以上で、データ準備は完了です。

超解像の実行

いよいよ超解像画像を生成します。
下記のコマンドを実行すると超解像処理が実行されます。

# 超解像実行
$ python infer.py -c ./config/sr_sr3_64_512.json
# experiments/"distributed_high_sr..."/results に結果画像出力

上記を実行すると下記のような出力結果が得られます。

左:入力画像(64x64)、右:超解像後画像(512x512)

左目などにやや違和感の残る結果となりましたね。
出力結果が優れない場合は個別に追加学習を行うことで解消されるかと思います。

まとめ

本記事では、SR3で超解像度画像を生成する方法を紹介しました。
画像のサイズが小さく、拡大に耐えられない画像などに適用したい技術ですね。

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

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

参考文献

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

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