圧縮ノイズ除去機能付き超解像AIを作成した。

今回は、JPEGやWebPなどの画像圧縮ノイズを除去する能力を備えた超解像AIをGANで作成致しました。
” SRGAN-R ” です。
GitHubの “SRGAN-R” のページ: https://github.com/ImpactCrater/SRGAN-R

因みに前回は、AutoencoderのEncoder出力を損失関数に組み込んだ人工ニューラルネットワークの超解像AI(GANではない)を実験致しました。

今回も董豪(Hao Dong)様のsrganのコードを元に私が改造致しました。

以下は私がAWS (Amazon Web Services)のEC2 (Amazon Elastic Compute Cloud)の ” p3.2xlarge ” GPU インスタンスとローカル PCで学習させたSRGAN-Rによるノイズ除去4倍超解像と無補間の4倍拡大とを比較した画像です。

無補間の4倍拡大と改造SRGAN-Rによる超解像との比較画像。左右に同じ画像を拡大したものが配置されており、左半分が無補間で右半分がSRGAN-Rである。ジーンズパンツを履いた若いモデルの女子が微笑んでいる。 — 左: 無補間 / 右: SRGAN-R

左は無補間の4倍拡大画像、右はSRGAN-Rで4倍に拡大した画像です。
SRGAN-Rでは圧縮ノイズが綺麗に除去されています。

無補間の4倍拡大と改造SRGAN-Rによる超解像との比較画像。左右に同じ画像を拡大したものが配置されており、左半分が無補間で右半分がSRGAN-Rである。ジーンズパンツのポケットに親指を掛けた手が写っている。 — 左: 無補間 / 右: SRGAN-R

無補間の4倍拡大と改造SRGAN-Rによる超解像との比較画像。左右に同じ画像を拡大したものが配置されており、左半分が無補間で右半分がSRGAN-Rである。若い女子の顔が写っている。 — 左: 無補間 / 右: SRGAN-R

やや滑らかになり過ぎている感じも致しますが、ノイズはほとんど除去出来ており、拡大も綺麗です。

これらの画像はSRGAN-Rには学習させていない未学習の画像です。
加えて、人間の全身像については全く学習していないモデルでの超解像の結果です。
元の画像はiStockより使用ライセンスを購入致しました。

[GitHub リポジトリー]
この ” SRGAN-R ” はTensorFlowとTensorLayerによる実装となっています。
以下に私のSRGAN-RのGitHub リポジトリーへのリンクを載せます。

GitHubの ” SRGAN-R ” のリポジトリー: https://github.com/ImpactCrater/SRGAN-R

[学習済みモデルデータ]
私がAWSのEC2のp3.2xlarge GPU インスタンスを用いて36時間ほど学習させ、60 epoch目の段階の学習済みモデルデータを~~公開致します~~。
(残念ながら、Yahoo!ボックスのサーヴィスは終了してしまいました。)
” ~~Yahoo!ボックス ” サーヴィス~~からZIP ファイルをダウンロードして解凍し、 ” ./SRGAN-R/checkpoint/ ” に ” d.npz, g.npz ” の2つのファイルを入れます。
学習用の大量のPNG画像を ” ./SRGAN-R/HRImage_Training ” のディレクトリーに用意し、端末から

python ./SRGAN-R/main.py

として続きから再学習させる事が可能です。
もちろん拡大だけの実行も可能です。

[SRGAN-Rの特徴]

[1: Residual of Residual Network]
勾配消失に強い、残差学習を行います。
Residual Blockで、shortcutに対して全てのConvolution層の出力との要素和を取っております。

[2: Generator vs Discriminator]
今回は学習の安定性が高い、 “Relativistic LSGAN” で学習させました。

[3: Noise Reduction]
JPEG画像や動画像のように圧縮ノイズが含まれた画像を拡大する際、ノイズを綺麗に除去する為、敢えて画像の縮小時にランダムな強さで圧縮を施して学習させました。
JPEGでは固定の大きさのブロック単位で圧縮しますが、これがノイズ除去の柔軟性を失わせてしまうと考え、ブロックが目立ち難いWebP画像形式で圧縮を施しました。
WebPは不可逆圧縮でJPEGより軽量、高品質であり、可逆圧縮でもPNGよりもファイルサイズを小さく出来る優れた圧縮画像形式です。
TensorFlow, TensorLayer, Python 3, PillowでWebPを使用する際は、端末で

sudo apt install libwebp-dev

としてWebPのライブラリーをインストールした上で、Pillowを再インストールする必要があります。
尚、結果は上出来で、強目の圧縮ノイズも綺麗に除去出来るようになりました。

[4: Swish活性化関数]
活性化関数には性質の優れたSwish関数を用いております。

Swish(x)=x * Sigmoid(a * x)
Sigmoid(x)=1 / (1 + e^(-a * x))

[5: Group Normalization]
ミニバッチサイズが少ない場合に問題が起こり易いBatch Normalizationではなく、channelをグループ化して正規化するGroup Normalization層を用いています。

[6: 画像の読み込み]
学習させる画像群について、指定したディレクトリー以下にあるサブディレクトリーまで再帰的に検索して読み込みます。
画像はepoch毎にランダムにシャッフルされ、ランダムな位置でクロップされます。
全画像ファイル数がミニバッチの数で割り切れない場合には、不足分は既存の学習用画像の中からランダムに再度読み込まれます。
また、学習中にも画像を次々に追加出来ます。
指定したディレクトリー以下に、画像ファイルもしくは画像ファイルが入ったフォルダーを放り込むだけで、次のepoch開始時に読み込まれます。

[プログラムの実行について]
このプログラムで学習を実行するには12GB以上のメモリーが必要です。
CPUまたはNVIDIAのGPUによるCUDAで実行出来ます。
低性能なCPUで実行する場合、充分な学習には数か月以上の時間が掛かります。

私は、AWSのEC2の ” p3.2xlarge ” GPU インスタンスをスポットリクエストで利用させて頂きました。
p3.2xlargeは1つのNVIDIA Tesla V100 GPUが利用可能です。
このGPUは64bit精度(倍精度)で7TFLOPS、詰まり最大で毎秒7兆回もの浮動小数点数演算が可能な性能を誇ります。
私の自宅のデスクトップ PCのIntel Core-i7 3770Tよりも約110倍高速に処理出来ました。

AWS EC2 p3.2xlarge GPU インスタンスはオンデマンド価格が私が実行終了した時点で $3.06 /1 時間でスポット価格が $1.0049 /1 時間 (US東部Virginia北部, Linux)でした。
スポットインスタンスの最高価格は初期設定であるオンデマンド価格とし、36時間ほど回しましたが、ほぼ最初の価格のまま最後まで停止無く使えました。
非常にお得です。