More than 3 years have passed since last update.

SRCNNで超解像をやってみた(TensorFlow)

Posted at 2021-09-20

はじめに

先日、先輩エンジニアの方とお話する機会があり、超解像タスクに取り組まれてるとのことで、超解像に興味を持ちました。

調べたところ、超解像と言ってもたくさん手法がありますが、なんでも学習し始めには最初のモデルがいいだろうということで、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)