はじめに
先日、先輩エンジニアの方とお話する機会があり、超解像タスクに取り組まれてるとのことで、超解像に興味を持ちました。
調べたところ、超解像と言ってもたくさん手法がありますが、なんでも学習し始めには最初のモデルがいいだろうということで、CNNモデルを使った一番最初のモデル「SRCNN」を実際に自分でTensorFlow, kerasで実装してみたのでその記録です。
超解像とは
めちゃくちゃ簡単にいうと画像の「高画質化」。
低画質の画像（ぼやけた画像）から擬似的に補うことで高画質の画像（細かい部分がはっきりと見える画像）を作る技術。
論文によれば
(b),(c)のような画像を低画質（解像度が低い）と言う。
この低画質な画像から(a)のような解像度が高い高画質な画像を生成する技術。
SRCNNとは
公式論文はコチラ
名の通りで超解像モデルにCNNを用いた手法で、当時はCNNモデルは組み込まれていなかったので画期的とされた。
SRCNNモデルは低解像度の画像を入力として受け取り、それを高解像度の画像に変換する。この変換の仕方をモデルは学習する。
興味深いのはモデル構造が非常にシンプルなこと。
畳み込み層3層しかないです。どうやらこれ以上層を厚くしても精度はほとんど変わらず、学習コストだけ増えたらしい。
※もちろん別の手法でこれよりも層を厚くして、かつ精度も高いモデルもあります。
1層目でパッチ抽出、2層目で非線形マッピング、3層目で画像の再構築を行なっている。
実行環境
画像データを扱うのでGoogleColabのGPUをお借りします。
Tensorflow 2.6.0
データの用意
自分のGoogleフォトの中の田村保乃ちゃんフォルダを使います()
動画とかもあるので画像データだけ読み取らせます。
この訓練データ正直なんでも良さげな気がするのですが、、、
低画質な画像から高画質な画像にする方法を学習できればどんな低画質な画像でも対応できそうな気がするけどどうなんだろう
ただ超解像タスクのコンペなどではそのテーマに合った訓練データが与えられるし、一概には言えないのかも
import shutil
shutil.copytree('/content/drive/MyDrive/田村保乃/', '/content/images')
GoogleDrive上のファイルを読み込むのは時間がかかりすぎるので、shutilを使って、ドライブ上の画像フォルダを、content直下のimagesフォルダにコピーします。
import glob
image_path_list = glob.glob('/content/images/*')
len(image_path_list)
# -> 551
globで画像フォルダ内の画像パスを全て取得します。今回は551枚（ただし動画も含まれているので、実際の訓練データ数はこれより少なくなります）
def create_dataset(image_path_list, size, scale, method='bilinear'):
high_imgs = []
low_imgs = []
for image_path in image_path_list:
if image_path.split('.')[-1] in ['mp4', 'gif']:
continue
img = cv2.imread(image_path)
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
high_img = cv2.resize(img_rgb, (size, size))
low_img = cv2.resize(high_img, (int(size/scale), int(size/scale)))
if method == 'bilinear':
low_img = cv2.resize(low_img, (size, size))
elif method == 'bicubic':
low_img = cv2.resize(low_img, (size, size), interpolation=cv2.INTER_CUBIC)
high_imgs.append(np.array(high_img).astype('float32') / 255)
low_imgs.append(np.array(low_img).astype('float32')/ 255)
high_imgs = np.array(high_imgs)
low_imgs = np.array(low_imgs)
return high_imgs, low_imgs
データセットを作る関数を作ります。画像だけ取得したいので拡張子で除外処理を行なってます。また取得した画像は縦横の長さがバラバラなので縦と横がsizeの長さの正方形にリサイズします。scaleは低解像度の画像を作るときに元の画像の何倍にするかの引数です。（scaleが4なら画像サイズを1/4にしてから元のサイズに拡大することを意味する。大きければ大きいほど、低画質になる）
最後の引数のmethodはcv2のリサイズする時の補完方法の指定を行うためのものです。bilinearとbicubicに対応しています。補完方法に関してはコチラが詳しいです。
※よくよく考えたらhigh_imgのリサイズも2通り選択できるようにしておいた方がいいですね
high_imgs, low_imgs = create_dataset(image_path_list, size=640, scale=4)
high_imgs.shape, low_imgs.shape
# ->
訓練データの枚数が少ないのと、とりあえず動かしてみるのが目的で精度は2の次なので、検証データのバリデーションは行いません。
モデル構築
def build_SRCNN():
model = Sequential()
model.add(Conv2D(64, kernel_size=9, padding='same', input_shape=(None, None, 3)))
model.add(Activation('relu'))
model.add(Conv2D(32, kernel_size=1, padding='same'))
model.add(Activation('relu'))
model.add(Conv2D(3, kernel_size=5, padding='same'))
return model
元論文の通りに実装しました。こう見るとこんな貧弱なモデルで大丈夫なのかという感じですが、、、
損失関数、評価指標、コンパイル
def psnr(low, high):
return tf.image.psnr(low, high, max_val=1.0)
optimizer = tf.keras.optimizers.Adam()
loss = tf.keras.losses.MeanSquaredError()
metrics = [psnr]
アルゴリズムはAdam、lossはMSE, 評価指標はPSNRを使います。PSNRについてはコチラの記事が分かりやすいです。
以前はPSNRをスクラッチで実装する必要があったぽいですが、tensorflowにもPSNRを計算してくれるtf.image.psnr関数があったのでこれを使います。引数の中のmax_valは輝度の最大値を指定しますが、今回前処理として255で割って0~1に正規化しているので最大値は1です。そのためmax_valの値を1としています。
model = build_SRCNN()
model.compile(optimizer = optimizer,
loss = loss,
metrics = metrics)
モデルをインスタンス化してコンパイルします。
学習
model.fitで行えるのもSRCNNモデルの手軽さがわかるポイント。
EPOCHS = 200
BATCH_SIZE = 16
history = model.fit(low_imgs, high_imgs, epochs=EPOCHS, batch_size=BATCH_SIZE)
# Epoch 196/200
# 31/31 [==============================] - 11s 360ms/step - loss: 0.0024 - psnr: # 27.4958
# Epoch 197/200
# 31/31 [==============================] - 11s 360ms/step - loss: 0.0024 - psnr: # 27.6208
# Epoch 198/200
# 31/31 [==============================] - 11s 359ms/step - loss: 0.0027 - psnr: # 26.8651
# Epoch 199/200
# 31/31 [==============================] - 11s 360ms/step - loss: 0.0027 - psnr: # 26.8456
# Epoch 200/200
# 31/31 [==============================] - 11s 358ms/step - loss: 0.0024 - psnr: # 27.4956
説明変数に低画質画像、目的変数に高画質画像を与えて学習させます。
200エポックで大体40分くらい。
テスト画像で実験
田村保乃ちゃんの画像を使いたいですが、全部学習データに使っちゃったので、「Dead end」でのカッコよさがエグい森田ひかるちゃんの画像を使います（）
image_path = ['/content/drive/MyDrive/hikaru_test.jpg']
test_high_img, test_low_img = create_dataset(image_path, size = 640, scale=4)
_, test_low_img2 = create_dataset(image_path, size = 640, scale=8)
_, test_low_img3 = create_dataset(image_path, size = 640, scale=4, method='bicubic')
pred = model.predict(test_low_img)
pred2 = model.predict(test_low_img2)
pred3 = model.predict(test_low_img3)
比較用として
①scale=4でbilinear補完をした低画質画像
②scale=8でbilinear補完をした低画質画像
③scale=4でbicubic補完をした低画質画像
の3種類を用意し、それぞれの低画質画像を学習済みモデルに入力として与えて画像を生成させます。
その結果が以下。
正直分かりづらいので、2つずつ対応させて表示します。
①scale=4でbilinear補完をした画像とSRCNNモデルが生成した画像
パーカーのロゴが若干綺麗になっているかも、、、
②scale=8でbilinear補完をした画像とSRCNNモデルが生成した画像
8倍ともなるとかなり画質が落ちてます。ただSRCNNモデルは全体的に綺麗になっているような。ぱっと見変化がわかりやすいのはこれですかね。
③scale=4でbicubic補完をした画像とSRCNNモデルが生成した画像
変化がわからん、、、 SRCNNモデルの方が顔の部分に緑色の線が入っているのが気になる
ちなみに数値として元の高画質画像とのPSNR値を計算すると以下のように。
4倍bilinear : 28.329716
4倍SRCNN : 28.710464
8倍bilinear : 25.28858
8倍SRCNN : 25.338888
4倍bicubic : 28.370005
4倍bicubicのSRCNN : 27.949007
これを見ると4倍でのSRCNNモデルが一番値が高く、有益とされているbicubic補完よりも綺麗になっていることがわかります。（肉眼ではほぼ差がないけど）
PSNRの値が30を超えると高画質といえるらしいですから、まだまだですね
まとめ
超解像、やってみるまでは難しそうと思っていたが、実装自体は簡単で、データセットの用意も高画質画像を用意するだけでよく手軽に実装できるのは良かった。
コチラで実施されている超解像コンペに参加して、より実践的な超解像技術を身につけたいと思う。