学習済み画風変換モデルを使って画像を好みの画風に変換する

pip install tensorflow
pip install tensorflow_hub
pip install scikit-image

# ライブラリを読み込む
import tensorflow as tf
import tensorflow_hub as hub

import matplotlib.pyplot as plt
from skimage import io

# 画風を変換する画像を指定する
content_path = r'C:\tmp\image\sample.png'

# 画風を指定する
style_path = r'C:\tmp\style\wave.jpg'

# 画像の縮尺をmax_dimで調整する
def img_scaler(image, max_dim=512):
    # 読み込んだ画像の　テンソルをfloat32でセットする
    original_shape = tf.cast(tf.shape(image)[:-1], tf.float32)

    # 画像の縮尺を設定する
    scale_ratio = max_dim / max(original_shape)

    # 設定した縮尺とint32で、new_shapeにテンソルをセットする
    new_shape = tf.cast(original_shape * scale_ratio, tf.int32)

    # 画像の縮尺を変換する
    return tf.image.resize(image, new_shape)


# モデルが画像を読み込めるように変換する
def load_img(path_to_img):
    # 画像をデータとして読み込む
    img = tf.io.read_file(path_to_img)

    # 画像データをRGBのテンソルに変換する
    img = tf.image.decode_image(img, channels=3)

    # 変換した配列のshapeを確認する
    # [縦 横 チャンネル]の配列になる
    print('変換する画像', path_to_img)
    print('変換前のshape', img.get_shape()[:])

    # RGBのテンソルのデータ型をfloat32に変換する
    img = tf.image.convert_image_dtype(img, tf.float32)

    # 画像の縮尺を変換する
    img = img_scaler(img)

    # 変換した配列のshapeを確認する
    # [縦 横 チャンネル]の配列になる
    print('返還後のshape', img.get_shape()[:])

    return img[tf.newaxis, :]


# 画像を画風変換モデルで扱えるようにテンソルへ変換する
content_image = load_img(content_path)
style_image = load_img(style_path)


# tensorflow_hubからmagentaの任意の画風に変換するためのモジュールを読み込む
hub_module = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')

# 画像の画風を変換する
stylized_image = hub_module(tf.constant(content_image), tf.constant(style_image))[0]

# 画像を描画するための準備
fig = plt.figure()
plt.imshow(stylized_image[0])

# 軸と目盛を表示しない設定
plt.axis('off')
# 画像を保存するときに余白を無くすための設定
plt.subplots_adjust(left=0, right=1, bottom=0, top=1)

# 画風を変換した画像を表示する
plt.show()

# 画風を変換した画像を保存する
io.imsave(r'C:\tmp\result\sample_wave.png', stylized_image[0])