まひろ量子のハックログ

プログラミングや機械学習などの知識を記録・共有します

深層学習の「超解像」でモザイクは除去できるのか (Tensorflow + Keras)

f:id:twx:20180909143610p:plain

画像出典: http://mmlab.ie.cuhk.edu.hk/projects/SRCNN.html

結論から言うと

うまくいかなかった。

あくまで実験記録として記事を書くが、ここにあるコードを真似してもうまく高解像度化できないので注意していただきたい。

超解像をやってみた。

モザイク画像から元画像を推定するAutoEncoderを作ってみた。深層学習フレームワークはTensorflow+Kerasを使用した。 参考にした論文はこれだ。

Learning a Deep Convolutional Network for Image Super-Resolution

SRCNNという、3層で構成されたシンプルなネットワークだ。任意のshapeの画像を受け取り、画像サイズを変えないようにpaddingを付与してCNNを3回繰り返すという処理を行う。活性化関数はReLuとした。

# ネットワークの定義
model = Sequential()
model.add(Conv2D(
    filters=64,
    kernel_size=9,
    padding='same',
    activation='relu',
    input_shape=(None, None, 3)
))
model.add(Conv2D(
    filters=32,
    kernel_size=1,
    padding='same',
    activation='relu'
))
model.add(Conv2D(
    filters=3,
    kernel_size=5,
    padding='same'
))

低解像度画像を作る

元画像を10%ほど粗くしたモザイク画像を作り、モザイク画像から元画像を推定させる。画像をモザイク化する関数は以下の記事を参考にしてOpenCVを用いて実装した。

note.nkmk.me

# 受け取ったnp arrayにモザイクをかけnp arrayとしてreturnする
def drop_resolution(arr, ratio=0.1):    
    tmp = cv2.resize(arr, None, fx=ratio, fy=ratio, interpolation=cv2.INTER_NEAREST)
    new_img = cv2.resize(tmp, arr.shape[:2][::-1], interpolation=cv2.INTER_NEAREST)
    new_arr = img_to_array(new_img)
    
    return new_arr

試しにモザイク化してみた。こんな感じになった。

f:id:twx:20180909144802p:plain

学習しよう

コード全文を以下に掲載する。Google ColaboratoryでGPUを使って回すと、だいたい30分〜40分くらいで終わる。

Google Colaboratoryで学習するには、画像データをGoogle Driveにアップロードしておき、ColaboratoryでGoogle Driveにアクセスできるように設定しておく必要がある。詳しくは以下の記事を参照していただきたい。

www.mahirokazuko.com

# 各種モジュールのインポート
import os
import glob
import math
import random
import cv2
import numpy as np
from tensorflow.python import keras
from tensorflow.python.keras import backend as K
from tensorflow.python.keras.models import Model, Sequential
from tensorflow.python.keras.preprocessing.image import ImageDataGenerator, load_img, img_to_array, array_to_img
from tensorflow.python.keras.layers import Add, Input, Conv2D, Conv2DTranspose, Dense, Input, MaxPooling2D, UpSampling2D, Lambda

# 受け取ったnp arrayにモザイクをかけnp arrayとしてreturnする
def mosaic(arr, ratio=0.1):    
    tmp = cv2.resize(arr, None, fx=ratio, fy=ratio, interpolation=cv2.INTER_NEAREST)
    new_img = cv2.resize(tmp, arr.shape[:2][::-1], interpolation=cv2.INTER_NEAREST)
    new_arr = img_to_array(new_img)
    return new_arr

# yieldをもちいてミニバッチをreturnするジェネレータ
def data_generator(data_dir, mode, mosaic_ratio=0.1, target_size=(128, 128), batch_size=32, shuffle=True):
    for imgs in ImageDataGenerator().flow_from_directory(
        directory=data_dir,
        classes=[mode],
        class_mode=None,
        color_mode='rgb',
        target_size=target_size,
        batch_size=batch_size,
        shuffle=shuffle
    ):
        x = np.array([mosaic(img_to_array(img), mosaic_ratio) for img in imgs])
        yield x/255., imgs/255.

        
DATA_DIR = 'data/'
BATCH_SIZE = 100

# 上で定義したジェネレータをもちいて学習データとテストデータをロード
train_data_generator = data_generator(DATA_DIR, 'train', batch_size=BATCH_SIZE, shuffle=True)
test_data_generator = data_generator(DATA_DIR, 'test', batch_size=BATCH_SIZE, shuffle=False)


# ネットワークの定義
model = Sequential()
model.add(Conv2D(
    filters=64,
    kernel_size=9,
    padding='same',
    activation='relu',
    input_shape=(None, None, 3)
))
model.add(Conv2D(
    filters=32,
    kernel_size=1,
    padding='same',
    activation='relu'
))
model.add(Conv2D(
    filters=3,
    kernel_size=5,
    padding='same'
))

# ピーク信号対雑音比(PSNR)を「解像度の高さ」の指標とする。この値が20以上くらいになると比較的、解像度が良いとされている。
def psnr(y_true, y_pred):
    return -10*K.log(
        K.mean(K.flatten((y_true - y_pred))**2)
    )/np.log(10)

# モデルをコンパイルする。損失関数は二乗誤差。最適化アルゴリズムはAdam。評価指標は上で定義したpsnr。
model.compile(
    loss='mean_squared_error', 
    optimizer='adam', 
    metrics=[psnr]
)

# 学習させる
model.fit_generator(
    train_data_generator,
    validation_data=test_data_generator,
    validation_steps=1,
    steps_per_epoch=100,
    epochs=50
)

結果を見てみる

# 未知のデータに対してモザイク除去を試してみる

from IPython.display import display_png

unknown_img_t = img_to_array( load_img('pingpong.jpg') )
unknown_img_x = mosaic(unknown_img_t)
unknown_img_y = model.predict(unknown_img_x.reshape(1,128,128,3))[0]

display_png( array_to_img( unknown_img_t ) )
display_png( array_to_img( unknown_img_x ) )
display_png( array_to_img( unknown_img_y ) )

以下のようになった。上から順に、オリジナル画像、モザイク画像、モザイク除去画像だ。

f:id:twx:20180909163941p:plain

確かに、少しはマシになっているのかもしれないが、モザイクが除去できているとは言えない結果となった。

冒頭の元論文では、「少しぼやけた画像」に対して高解像度化の効果が得られたと書かれていた。さすがに、今回のモザイクのようにかなり粗い画像に対しては期待できる効果は得られないようだ。

GAN等を使ったほうが、良い結果になるかもしれない。

以上、今回は超解像を試してみました。 良い記事だと思っていただけた方は下の「★+」ボタンのクリック、SNSでのシェア、「読者になる」ボタンのクリックをお願いいたします。 それではまた〜

Kozuko Mahiro's Hacklog ―― Copyright © 2018 Mahiro Kazuko