tensorflow_hub实现风格迁移

Posted 2021-12-07 @微笑男孩

准备环境：

cuda是10的版本    cudart64_100.dll

下面是效果图

下面是代码，直接复制到编译器运行即可，可能会报错， 要看自己电脑有没有对应的库 和 图片的路径记得改成自己的  清晰度高的图片效果更加哦~~~

import matplotlib.pyplot as plt
import tensorflow_hub as hub
import tensorflow as tf
import numpy as np
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '/gpu:0'

# 两张图片  我们想要做自己的风格迁移 只需要将下面的图片的地址替换为自己的即可
# 风格图片
image1 = plt.imread(r"D:\\All-Python\\tensorflow_Hub\\风格迁移\\Image\\img_7.png")

# 原始图片
image2 = plt.imread(r"D:\\All-Python\\tensorflow_Hub\\风格迁移\\Image\\img_6.png")

# 加载模型  从tensorflow_hub 官网下载
hub_model = hub.load('https://hub.tensorflow.google.cn/google/magenta/arbitrary-image-stylization-v1-256/2')

# 改变维度  因为传入模型需要四维数据
image3 = image1[np.newaxis, :, :, :]
image4 = image2[np.newaxis, :, :, :]

# 规范数据  将格式转换为Tensor
image3 = tf.convert_to_tensor(image3, dtype=tf.float32)
image4 = tf.convert_to_tensor(image4, dtype=tf.float32)

# 第一个参数是内容图片 第二个参数是风格图片
outputs = hub_model(image4, image3)

# 定义一个画板 显示图片
plt.figure()
# 表示将整个画板分为 1行 3列 当前位置为 1
plt.subplot(1, 3, 1)
plt.imshow(image1)
# 表示将整个画板分为 1行 3列 当前位置为 2
plt.subplot(1, 3, 2)
plt.imshow(image2)

plt.subplot(1, 3, 3)
# 结果得到的是一个列表  相当于得到了五维的数据 
# 我们提取其中的RGB数据用来显示图片
plt.imshow(outputs[0][0])
plt.show()

 有很多图片的组合效果还是很不错的：