MXNet-Gluon-Style-Transferを使う（複数の変換を一度に実行）

はじめに

以下のMXNet-Gluon-Style-Transferを使ってみたいと思います。
github.com
以前にも試したことがあります。
touch-sp.hatenablog.com
今回は複数の変換を一度に実行するためのスクリプトを書きました。
また、GPUを使う場合と使わない場合の両方で実行可能なスクリプトにしました。

結果

サンプル画像（元画像）はこちらからダウンロードさせて頂きました。
（chihiro sakaiによるPixabayからの画像）

f:id:touch-sp:20201212172949j:plain:w300 — 元画像

f:id:touch-sp:20201212174409p:plain — 変換後画像

1回の実行で21個の変換が行われます。

環境構築

Windows 10
Python 3.7.9

（GTX 1080  CUDA Toolkit 10.1）

PythonにインストールするのはMXNetとPillowだけです。

GPUがある人

pip install mxnet-cu101==1.7.0 -f https://dist.mxnet.io/python/cu101
pip install pillow

「cu101」の部分はCUDA Toolkitのバージョンに合わせて変更して下さい。
最終的にはこうなりました。（pip freeze）

certifi==2020.12.5
chardet==3.0.4
graphviz==0.8.4
idna==2.6
mxnet-cu101==1.7.0
numpy==1.16.6
Pillow==8.0.1
requests==2.18.4
urllib3==1.22

GPUがない人（CPUのみ）

pip install mxnet
pip install pillow

最終的にはこうなりました。（pip freeze）

certifi==2020.12.5
chardet==3.0.4
graphviz==0.8.4
idna==2.6
mxnet==1.7.0.post1
numpy==1.16.6
Pillow==8.0.1
requests==2.18.4
urllib3==1.22

方法

二つのPython スクリプトファイルを用意します。
「download.py」：モデルやスタイル画像をダウンロード、解凍するスクリプト
「transfer.py」：変換を実行するスクリプト

download.py

import os
import zipfile

from mxnet import gluon

url = 'https://github.com/StacyYang/MXNet-Gluon-Style-Transfer/archive/master.zip'
filename = gluon.utils.download(url)

with zipfile.ZipFile(filename) as f:
        f.extractall('.')

os.rename('MXNet-Gluon-Style-Transfer-master', 'master')

url2 = 'https://apache-mxnet.s3-accelerate.amazonaws.com/gluon/models/msgnet_21styles-2cb88353.zip'
filename2 = gluon.utils.download(url2)

with zipfile.ZipFile(filename2) as f:
        f.extractall('./master/models')

os.remove(filename)
os.remove(filename2)

os.chdir('./master/models')
os.rename('msgnet_21styles-2cb88353.params', '21styles.params')

transfer.py

import sys
import os
import glob

import mxnet as mx

from master import net
from master import utils

os.makedirs('result', exist_ok=True)

if mx.context.num_gpus() >0:
    ctx = mx.gpu()
    content_size = 1024
    set_cuda = 1
else:
    ctx = mx.cpu()
    content_size = 512
    set_cuda = 0

input_file = sys.argv[1]
style_files = glob.glob('./master/images/styles/*.jpg')

content_image = utils.tensor_load_rgbimage(input_file, ctx, size=content_size, keep_asp=True)

style_model = net.Net(ngf=128)
style_model.load_parameters('./master/models/21styles.params', ctx=ctx)

for style_file in style_files:

    output_filename = os.path.basename(style_file)

    style_image = utils.tensor_load_rgbimage(style_file, ctx, size=512)
    style_image = utils.preprocess_batch(style_image)

    style_model.set_target(style_image)
    output = style_model(content_image)
    utils.tensor_save_bgrimage(output[0], os.path.join('result', output_filename), set_cuda)