DeepLearningを用いた超解像手法/SRCNNの実装

概要

深層学習を用いた、単一画像における超解像手法であるSRCNNの実装したので、それのまとめの記事です。
Python + Tensorflow(Keras)で実装を行いました。
今回は、2倍拡大の超解像にチャレンジしました。

以前、3つに分けて記事を上げていたのですが、コードが汚かった＋1つにまとめたかったので再投稿しています。ご了承ください。

今回紹介するコードはGithubにも載せています。

1. 超解像のおさらい

超解像について簡単に説明をします。

超解像とは解像度の低い画像に対して、解像度を向上させる技術のことです。
ここでいう解像度が低いとは、画素数が少なかったり、高周波成分（輪郭などの鮮鋭な部分を表す成分）がないような画像のことです。
以下の図で例を示します。(図は[論文]より引用)

(a)は原画像、(b)は画素数の少ない画像を見やすいように原画像と同じ大きさにした画像、(c)は高周波成分を含まない画像の例です。
(b)と(c)は、荒かったりぼやけていたりしていると思います。
このような状態を解像度が低い画像といいます。

そして、超解像はこのような解像度が低い画像に処理を行い、(a)のような精細な画像を出力することを目的としています。

2. SRCNNの超解像アルゴリズム

SRCNNは初めて超解像分野に深層学習を導入したモデルです。
モデルとしては単純な構造にはなっていて、Convolutionを3層組み合わせたモデルとなっています。

SRCNNのアルゴリズムの概要図は以下の通りです。(図は論文から引用)

今回実装したSRCNNは、事前にbicubicで拡大処理を行います。

3. 実装したアルゴリズム

今回、実装したSRCNNのモデルは以下のように組みました。(コードの一部を抽出)

def SRCNN(input_channels = 1):
    #input single image
    input_shape = Input((None, None, input_channels))

    #convolution
    conv2d_0 = Conv2D(filters = 64, kernel_size = (9, 9), padding = "same", activation = "relu")(input_shape)
    conv2d_1 = Conv2D(filters = 32, kernel_size = (1, 1), padding = "same", activation = "relu")(conv2d_0)
    conv2d_2 = Conv2D(filters = input_channels, kernel_size = (5, 5), padding = "same")(conv2d_1)

    model = Model(inputs = input_shape, outputs = [conv2d_2])
    model.summary()

    return model

前述の通り、3層のConvolution層を組み合わせています。

4. 使用したデータセット

今回は、データセットにDIV2K datasetを使用しました。
このデータセットは、単一画像のデータセットで、学習用が800種、検証用とテスト用が100種類ずつのデータセットです。
今回は、学習用データと検証用データを使用しました。
パスの構造はこんな感じです。

train_sharp - 0001.png 
            - 0002.png 
            - ...
            - 0800.png

val_sharp   - 0801.png
            - 0802.png 
            - ...
            - 0900.png

このデータをbicubicで縮小したりしてデータセットを生成しました。

5. 画像評価指標PSNR

今回は、画像評価指標としてPSNRを使用しました。
PSNR とは Peak Signal-to-Noise Ratio(ピーク信号対雑音比) の略で、単位はデジベル (dB) で表せます。
PSNR は信号の理論ピーク値と誤差の2乗平均を用いて評価しており、8bit画像の場合、255(最大濃淡値)を誤差の標準偏差で割った値です。
今回は、8bit画像を使用しましたが、計算量を減らすため、全画素値を255で割って使用しました。
そのため、最小濃淡値が0で最大濃淡値が1です。
dB値が高いほど拡大した画像が元画像に近いことを表します。
PSNRの式は以下のとおりです。

PSNR = 10\log_{10} \frac{1^2 * w * h}{\sum_{x=0}^{w-1}\sum_{y=0}^{h-1}(p_1(x,y) - p_2(x,y))^2 }

なお、$w$は画像の幅、$h$は画像の高さを表しており、$p_1$は元画像、$p_2$はPSNRを計測する画像を示しています。

6. コードの使用方法

このコード使用方法は、自分が執筆した別の実装記事とほとんど同じです。

① 学習データ生成

まず、Githubからコードを一式ダウンロードして、カレントディレクトリにします。
Windowsのコマンドでいうとこんな感じ。

C:~/keras_SRCNN>

次に、main.pyから生成するデータセットのサイズ・大きさ・切り取る枚数、ファイルのパスなどを指定します。
main.pyの12~21行目です。
使うPCのメモリ数などに応じで、画像サイズや学習データ数の調整が必要です。

main.py

    parser.add_argument('--train_height', type = int, default = 33, help = "Train data size(height)")
    parser.add_argument('--train_width', type = int, default = 33, help = "Train data size(width)")
    parser.add_argument('--test_height', type = int, default = 700, help = "Test data size(height)")
    parser.add_argument('--test_width', type = int, default = 700, help = "Test data size(width)")
    parser.add_argument('--train_dataset_num', type = int, default = 10000, help = "Number of train datasets to generate")
    parser.add_argument('--test_dataset_num', type = int, default = 5, help = "Number of test datasets to generate")
    parser.add_argument('--train_cut_num', type = int, default = 10, help = "Number of train data to be generated from a single image")
    parser.add_argument('--test_cut_num', type = int, default = 1, help = "Number of test data to be generated from a single image")
    parser.add_argument('--train_path', type = str, default = "../../dataset/DIV2K_train_HR", help = "The path containing the train image")
    parser.add_argument('--test_path', type = str, default = "../../dataset/DIV2K_valid_HR", help = "The path containing the test image")

指定したら、コマンドでデータセットの生成をします。

C:~/keras_SRCNN>python main.py --mode train_datacreate

これで、train_data_list.npzというファイルのデータセットが生成されます。

ついでにテストデータも同じようにコマンドで生成します。コマンドはこれです。

C:~/keras_SRCNN>python main.py --mode test_datacreate

② 学習

次に学習を行います。
設定するパラメータの箇所は、epoch数と学習率、今回のモデルの層の数です。
まずは、main.pyの22~25行目

main.py

    parser.add_argument('--learning_rate', type = float, default = 1e-4, help = "Learning_rate")
    parser.add_argument('--BATCH_SIZE', type = int, default = 32, help = "Training batch size")
    parser.add_argument('--EPOCHS', type = int, default = 500, help = "Number of epochs to train for")

後は、学習のパラメータをあれこれ好きな値に設定します。73~86行目です。

main.py

        train_model = model.SRCNN()

        optimizers = tf.keras.optimizers.Adam(lr = args.learning_rate)
        train_model.compile(loss = "mean_squared_error",
                        optimizer = optimizers,
                        metrics = [psnr])

        train_model.fit(train_x,
                        train_y,
                        epochs = args.EPOCHS,
                        verbose = 2,
                        batch_size = args.BATCH_SIZE)

        train_model.save("SRCNN_model.h5")

optimizerはmomentum、損失関数はmean_squared_errorを使用しています。

学習はデータ生成と同じようにコマンドで行います。

C:~/keras_SRCNN>python main.py --mode train_model

これで、学習が終わるとモデルが出力されます。

③ 評価

最後にモデルを使用してテストデータで評価を行います。
これも同様にコマンドで行いますが、事前に①でテストデータも生成しておいてください。

C:~/keras_SRCNN>python main.py --mode evaluate

このコマンドで、画像を出力してくれます。

7. 結果

出力した画像はこのようになりました。
なお、今回は輝度値のみで学習を行っているため、カラー画像には対応していません。

• 元画像
  

• 低解像度画像
  PSNR:30.83
  

• 生成画像
  PSNR:31.43
  

論文ほどの数値ではありませんが、補間法よりはPSNRの値は大きくなりました。
論文の数値に近づけるには、学習回数などのパラメータ設定・学習データ数の増加などを行えばよさそうです。

最後に元画像・低解像度画像・生成画像の一部を並べて表示してみます。

数値の差があまりないので目視ではあまり変化は見えないですね...

8. コードの全容

前述の通り、Githubに載せています。
pythonのファイルは主に3つあります。
各ファイルの役割は以下の通りです。

• data_create.py : データ生成に関するコード。
• model.py : 超解像のアルゴリズムに関するコード。
• main.py : 主に使用するコード。

9. まとめ

今回は、以前作っていたSRCNNのコードをもう一度見直し、記事も見直しました。

記事が長くなってしまいましたが、最後まで読んでくださりありがとうございました。

参考文献

・Image Super-Resolution Using Deep Convolutional Networks
　実装の参考にした論文。
・画素数の壁を打ち破る 複数画像からの超解像技術
　超解像の説明のために使用。
・DIV2K dataset
　今回使用したデータセット。