はじめに

「waifu2x」というWebサービスをご存じだろうか。waifu2xは深層学習を利用した画像処理モデルを用いて、アップロードされた画像に対して画像拡大およびノイズ除去を提供するWebサービスである。

waifu2xのソースコードは元々githubに公開されており、その派生バージョンが当時数々存在した。その中の一つにPython版でwaifu2xの処理を再現した有志が存在した。そちらを紹介させてもらおうと思う。
当時のソースコードであるため、現在では非推奨となっているライブラリも存在するが、それも含めて解説していく。

■ソースコードのURL

marcan.st

https://marcan.st/transf/waifu2x.py

Contents

コードの詳細な解説
注意点

コードの詳細な解説

インポート

import json, sys, numpy as np
from scipy import misc, signal
from PIL import Image

json: モデルデータを読み込むために使用
sys: コマンドライン引数や標準エラー出力を扱うために使用
numpy: 数値計算を効率的に行うためのライブラリ
scipy.misc: 古いSciPyの画像処理関連機能（現在非推奨）
scipy.signal: 畳み込み処理に使用
Pillow (PIL): 画像の読み込み、変換、保存を行うためのライブラリ

コマンドライン引数の取得

infile, outfile, modelpath = sys.argv[1:]

infile: 入力画像のパス
outfile: 出力画像のパス
modelpath: 学習済みモデルのJSONファイルのパス

モデルデータの読み込み

model = json.load(open(modelpath))

model: JSON形式で保存された深層学習モデルのパラメータを格納。

入力画像の読み込みと前処理

im = Image.open(infile).convert("YCbCr")
im = misc.fromimage(im.resize((2*im.size[0], 2*im.size[1]), resample=Image.NEAREST)).astype("float32")
planes = [np.pad(im[:,:,0], len(model), "edge") / 255.0]

入力画像をYCbCr色空間に変換
Y: 輝度（明るさ）
Cb, Cr: 色差成分
最近傍補間を用いて画像を2倍に拡大。
画像のYチャンネル（輝度成分）のみを抽出し、モデルサイズに応じたパディングを適用。

モデルの畳み込み処理

count = sum(step["nInputPlane"] * step["nOutputPlane"] for step in model)
progress = 0
for step in model:
assert step["nInputPlane"] == len(planes)
assert step["nOutputPlane"] == len(step["weight"]) == len(step["bias"])
o_planes = []
for bias, weights in zip(step["bias"], step["weight"]):
partial = None
for ip, kernel in zip(planes, weights):
p = signal.convolve2d(ip, np.float32(kernel), "valid")
if partial is None:
partial = p
else:
partial += p
progress += 1
sys.stderr.write("\r%.1f%%..." % (100 * progress / float(count)))
partial += np.float32(bias)
o_planes.append(partial)
planes = [np.maximum(p, 0) + 0.1 * np.minimum(p, 0) for p in o_planes]

モデルの各ステップ（レイヤー）を順次処理。
nInputPlane: 入力チャンネル数。
nOutputPlane: 出力チャンネル数。
weights: 畳み込みカーネル。
bias: バイアス項。
畳み込み演算を適用し、ReLUベースの活性化関数を適用。
正の値はそのまま、負の値は0.1倍。

出力画像の生成

assert len(planes) == 1
im[:,:,0] = np.clip(planes[0], 0, 1) * 255
misc.toimage(im, mode="YCbCr").convert("RGB").save(outfile)
sys.stderr.write("Done\n")

処理済みの輝度チャンネルを画像データに戻す。
輝度値をクリップし、YCbCrからRGB形式に変換。
変換後の画像を指定されたファイルに保存。

注意点

非推奨モジュールの使用

– scipy.miscのfromimageとtoimageは非推奨。代替としてnumpyやPillowを利用すべき。
2. リソース消費
– 高解像度画像や複雑なモデルは大量のメモリを消費する場合がある。

まとめ

waifu2xの派生バージョンの一つであるPython版waifu2xのソースコードを紹介した。
これを参考に、TensorFlowやPyTorchなどのモダンなライブラリを用いて改修すれば、waifu2xの処理を再現することも可能となるかもしれない。

Python版waifu2xのソースコードを解読してみる