はじめに
深層学習を使用した単一画像超解像について記事を書くが、すでに多くの記事があるので、個人的に気づきがあった点を重点的に記すことにする。一通りの処理を実装したGoogle Colabのノートブックも公開する。
単一画像超解像とは
Single Image Super Resolution(SISR)は低画質画像を入力として、高画質画像を出力とする処理。
2018年ぐらいまでの主要な手法については下記の記事に詳しい。
トップ学会採択論文にみる、超解像ディープラーニング技術のまとめ
下記論文には各種手法の性能比較等がある。
A Deep Journey into Super-resolution: A Survey
以下のページは最新を含む多くの関連論文へのリンクがある。
Awesome-Super-Resolution
この記事では、実装しやすい割りに高性能なEDSRを実装する。
Enhanced Deep Residual Networks for Single Image Super-Resolution
Code
実装上の注意点
以下、実装して気づいた点など
Datasetの生成
これはTensorflow(というかTPU)固有の問題かもしれないが、意外と工数がかかったので最初に記す。
SISRでは標準的にDIV2Kと言うデータが学習用に使用される。本来はここからランダムに領域を抽出して学習データにするのだが、高画質画像のPNGをまとめたzipファイルでも3GBあり、学習データはオンメモリにするのではなくファイルから逐次読み出したいのだが、TPUではファイルアクセスさせるようなDatasetはエラーが出る(Google Storage Serviceが必要)。
仕方がないので事前に一部の領域をパッチとして用意して、tf.data.Dataset.from_tensor_slicesを使ってDatasetを生成する事にした。from_tensor_slices自体にもデータ量の制限(多分２G)があるようで、それを超えるサイズを作りたい場合は個別に生成してからconcatenateで結合させる必要があった。
今回の実装は、パッチの数を減らすために事前の前処理でエッジが多い領域を優先して切り抜くことで、約半分の領域だけ使って学習させることにした。
PSNRの計算方法
超解像によって画像をどの程度復元できたか、はPSNR(Peak Signal-to-Noise Ratio)などを使って定量的に評価する。超解像ではPSNRはRGBで計算するのではなく、以下のような方法で行うのが一般的なようだ。
- YCbCRのY成分のみで、Yの範囲は16〜235
- 画像端の数ピクセルは計算対象外とする
正則化手法は使わない
CNNを使った画像認識では各種正則化で性能向上させることが一般的だが、超解像ではCNNを使う場合でも正則化は効果が無いか寧ろ逆効果な場合が多いようだ。
したがってEDSRではBatchNormalizationやWeightDecay等は使わない。
また、画像認識ではOptimizerにSGDを使った方がAdam等より性能が良い場合がほとんどだが、超解像ではAdamがよく使われ、実際にEDSRではAdamの方が好成績だった。SGD自体に正則化の効果があるようなので関連があるのではないかと思われる。
「画像認識は抽象化の学習」で「超解像では具体化の学習」だから、と言う説明でそれなりに納得はするが、なんだか不思議な感じがする。
実装
事前のパッチ化
先に述べたDatasetの作成処理で、大体5分ぐらいかかるが、この時間はTPUを使うことの高速化で挽回できる。時間がかかるので進捗表示にtqdmを始めて使ってみたが、これは便利。
事前にDIV2KのZipファイルをダウンロードしておくことが必要。
import zipfile
import os
import io
import random
from PIL import Image
from PIL import ImageFilter
import numpy as np
from tqdm.notebook import tqdm
from tensorflow.keras import backend
import tensorflow as tf
def make_dataset_from_div2k( filenames, num_patches, scale=2, patch_size = 96):
print("make_dataset_from_zip", filenames)
lr_patch_size = patch_size//scale
hr_filename = filenames[0]
lr_filename = filenames[1]
with zipfile.ZipFile(hr_filename) as hr_zip:
img_files = 0
for img_filename in hr_zip.namelist():
if not img_filename.endswith(".png"):
continue
img_files += 1
pbar = tqdm(total=img_files)
img_count = 0
average_patches_in_file = num_patches/img_files
imgbuf = None
patches_list = []
dataset = None
file_count = 0
for img_filename in hr_zip.namelist():
if not img_filename.endswith(".png"):
continue
lr_img_filename = f'{img_filename[0:11]}_LR_bicubic/X{scale}/{img_filename[-8:-4]}x{scale}.png'
with hr_zip.open(img_filename) as img_file:
img_bin = io.BytesIO(img_file.read())
img = Image.open(img_bin)
num_h = img.height//patch_size
num_w = img.width//patch_size
patches = num_h*num_w
patches_list.append(patches)
patch_list = []
for h in range(num_h):
for w in range(num_w):
x = w * patch_size
y = h * patch_size
cropped = img.crop((x, y, x+patch_size, y+patch_size))
img_find_edges = cropped.filter(ImageFilter.FIND_EDGES)
patch_list.append( {"score":np.mean(img_find_edges),
"img":cropped,
"pos":(x,y) })
patch_list.sort(reverse=True,key=lambda x:x['score'])
patches_in_file = 0
with zipfile.ZipFile(lr_filename) as lr_zip:
with lr_zip.open(lr_img_filename) as img_file:
img_bin = io.BytesIO(img_file.read())
img = Image.open(img_bin)
patches_in_file = int(average_patches_in_file*(file_count+1)-img_count)
patches_in_file = min(patches, patches_in_file)
for i in range(int(patches_in_file)):
if imgbuf is None:
remain = num_patches-img_count
if remain==0:
break
img_count_in_ds = 0
num_patches_in_ds = min(1024*32, remain)
imgbuf=np.zeros((num_patches_in_ds, patch_size, patch_size, 3), dtype=np.uint8)
lr_imgbuf=np.zeros((num_patches_in_ds, lr_patch_size, lr_patch_size, 3), dtype=np.uint8)
imgbuf[img_count_in_ds] = np.array( patch_list[i]["img"] )
x,y = patch_list[i]["pos"]
lr_img = img.crop((x//scale, y//scale, (x+patch_size)//scale, (y+patch_size)//scale))
lr_imgbuf[img_count_in_ds] = np.array( lr_img )
img_count+=1
img_count_in_ds+=1
if img_count_in_ds==num_patches_in_ds:
if dataset is None:
dataset = tf.data.Dataset.from_tensor_slices((lr_imgbuf,imgbuf))
else:
dataset = dataset.concatenate( tf.data.Dataset.from_tensor_slices((lr_imgbuf,imgbuf)) )
imgbuf = None
lr_imgbuf = None
img.close()
file_count+=1
# print(file_count)
pbar.update(1)
if imgbuf is not None:
print("!!")
# print(min(patches_list))
pbar.close()
return dataset
EDSR
以下、EDSRのコード。大変シンプルだが、scaling_factorで調整しないとモデルが崩壊する場合がある。この辺が正則化を入れない弊害か。
モデルとしてはInputでサイズを指定していないので、メモリが許す限りどんなサイズの画像でも入力可能。
論文では48x48を低画質のパッチとし、高画質画像は倍率(2/3/4倍)のサイズをかけた大きさの画像として訓練しているが、この実装では96x96を高画質のパッチとして、2/3/4で割ったサイズの画像を低画質の入力サイズとしているので、２倍以外は論文の学習よりも計算量が減っている。
import tensorflow as tf
import tensorflow.keras.layers as layers
import numpy as np
DIV2K_RGB_MEAN = np.array([0.4488, 0.4371, 0.4040])* 255
def normalize(x, rgb_mean=DIV2K_RGB_MEAN):
return (x - rgb_mean) / 127.5
def denormalize(x, rgb_mean=DIV2K_RGB_MEAN):
return (x * 127.5) + rgb_mean
def upsample(x, scale, num_filters):
if 2 <= scale <= 3 :
x = layers.Conv2D(num_filters * (scale ** 2), 3, padding='same')(x)
x = layers.Lambda(lambda y: tf.nn.depth_to_space(y, scale))(x)
elif scale == 4:
x = layers.Conv2D(num_filters * (2 ** 2), 3, padding='same')(x)
x = layers.Lambda(lambda y: tf.nn.depth_to_space(y, 2))(x)
x = layers.Conv2D(num_filters * (2 ** 2), 3, padding='same')(x)
x = layers.Lambda(lambda y: tf.nn.depth_to_space(y, 2))(x)
return x
def create_model(num_filters=64, num_res_blocks=8, scale=2, scaling_factor=None, activation='relu'):
inputs = layers.Input( (None, None, 3) )
CONV_KWARGS ={
'padding':'same',
'kernel_initializer':'he_uniform',
'bias_initializer':'he_uniform',
}
lowres = layers.Lambda(normalize)(inputs)
# EDSR
def res_block(x_in, filters, scaling=None):
x = layers.Conv2D(filters, 3, activation=activation, **CONV_KWARGS)(x_in)
x = layers.Conv2D(filters, 3, **CONV_KWARGS)(x)
if scaling:
x = layers.Lambda(lambda t: t * scaling)(x)
x = layers.Add()([x_in, x])
return x
x = lowres = layers.Conv2D(num_filters, 3, **CONV_KWARGS)(lowres)
for i in range(num_res_blocks):
x = res_block(x, num_filters, scaling=scaling_factor)
x = layers.Conv2D(num_filters, 3, **CONV_KWARGS)(x)
lowres = layers.Add()([x, lowres])
# Upsampling
highres = upsample(lowres, scale, num_filters)
highres = layers.Conv2D(3, 3, **CONV_KWARGS)(highres)
highres = layers.Lambda(denormalize)(highres)
return tf.keras.models.Model(inputs, highres)
実験結果
論文に近い形で実装したつもりだが、学習率に関しては少し調整してある。
Set5/Set14/BSD100/Urban100は超解像の評価によく用いられるデータセットで下記のgithubからダンロードしている。
https://github.com/jbhuang0604/SelfExSR.git
DIV2KについてはGoogle Drive内にzipがある前提の実装なので、実際に試す場合は最初のセルを適宜修正する必要がある。
訓練時間短縮のため、DIV2Kは訓練データのみ使用し、５エポックに１回Set14のPSNRを計測するように実装してある。
以下4xでの実験結果で、数値はPSNR。最後の行以外は論文内の値。
| Method | Set5 | Set14 | BSD100 | Urban100 |
|---|---|---|---|---|
| Bicubic | 28.42 | 26.00 | 25.96 | 23.14 |
| EDSR(Paper) | 32.46 | 28.80 | 27.71 | 26.64 |
| EDSR | 32.40 | 28.72 | 27.56 | 26.29 |
実験では、論文の値に若干届かないものの、かなり近いところまで再現できている。
EDSR(パラメータ数40M)だと3時間程度かかるが、論文にあるbaseline設定(パラメータ数1.5M)なら30分ぐらいで終わる。