TensorFlow实现图像风格转换 V1算法

东方耀

TensorFlow实现图像风格转换 V1算法实现

1. # -*- coding: utf-8 -*-
   
2. __author__ = 'dongfangyao'
   
3. __date__ = '2019/2/4 下午3:33'
   
4. __product__ = 'PyCharm'
   
5. __filename__ = '10_image_style_conversion'
   
6. """
   
7.    TensorFlow实现图像风格转换 V1算法
   
8.    v1算法的缺点：每次需要随机初始化图像的变量 使用GD下降 运行多次 效率低
   
9.    v1算法训练的是 图像本身
   
10.    v2算法训练的是 网络（Image Transform Net）
    
11.    v3算法：重新定义了风格损失的计算方法 放弃了Gram矩阵计算相似度 用了最match小块分割方法（patch为单位）
    
12. """
    
13. import os
    
14. import numpy as np
    
15. import tensorflow as tf
    
16. import time
    
17. from PIL import Image
    
18. import matplotlib.pyplot as plt
    
19. 
    
20. VGG_MEAN = [103.939, 116.779, 123.68]
    
21. 
    
22. 
    
23. class VGGNet:
    
24.     """
    
25.     构建VGG16的网络结构 并从预处理模型中加载训练好的参数
    
26.     """
    
27.     def __init__(self, data_dict):
    
28.         self.data_dict = data_dict
    
29. 
    
30.     def get_conv_kernel(self, name):
    
31.         # 卷积核的参数 w:0 bias：1
    
32.         return tf.constant(self.data_dict[name][0], name='conv')
    
33. 
    
34.     def get_fc_weight(self, name):
    
35.         return tf.constant(self.data_dict[name][0], name='fc')
    
36. 
    
37.     def get_bias(self, name):
    
38.         return tf.constant(self.data_dict[name][1], name='bias')
    
39. 
    
40.     def conv_layer(self, inputs, name):
    
41.         """构建一个卷积计算层"""
    
42.         # 多使用name_scope的好处：1、防止参数命名冲突 2、可视化的显示规整
    
43.         with tf.name_scope(name):
    
44.             conv_w = self.get_conv_kernel(name)
    
45.             conv_b = self.get_bias(name)
    
46.             # tf.layers.conv2d() 里面没有参数的 不能用了
    
47.             result = tf.nn.conv2d(inputs, conv_w, [1, 1, 1, 1], padding='SAME')
    
48.             result = tf.nn.bias_add(result, conv_b)
    
49.             result = tf.nn.relu(result)
    
50.             return result
    
51. 
    
52.     def pooling_layer(self, inputs, name):
    
53.         """构建一个池化层 tf.layers.max_pooling2d()"""
    
54.         result = tf.nn.max_pool(inputs, [1, 2, 2, 1], [1, 2, 2, 1], padding='SAME', name=name)
    
55.         return result
    
56. 
    
57.     def fc_layer(self, inputs, name, activation=tf.nn.relu):
    
58.         """构建全连接层的计算"""
    
59.         with tf.name_scope(name):
    
60.             fc_w = self.get_fc_weight(name)
    
61.             fc_b = self.get_bias(name)
    
62.             result = tf.nn.bias_add(tf.matmul(inputs, fc_w), fc_b)
    
63.             if activation is None:
    
64.                 return result
    
65.             else:
    
66.                 return activation(result)
    
67. 
    
68.     def flatten_op(self, inputs, name):
    
69.         """展平操作 tf.layers.flatten()"""
    
70.         with tf.name_scope(name):
    
71.             # [N H W C]---> [N H*W*C]
    
72.             x_shape = inputs.get_shape().as_list()
    
73.             dim = 1
    
74.             for d in x_shape[1:]:
    
75.                 dim *= d
    
76.             inputs = tf.reshape(inputs, shape=[-1, dim])
    
77.             return inputs
    
78. 
    
79.     def build(self, input_rgb):
    
80.         """构建vgg16网络 提取特征 FP过程
    
81.            参数：
    
82.            --input_rgb: [1, 224, 224, 3]
    
83.         """
    
84.         start_time = time.time()
    
85.         print('building start...')
    
86. 
    
87.         # r, g, b = tf.split(value=input_rgb, num_or_size_splits=[1, 1, 1], axis=3)
    
88.         r, g, b = tf.split(value=input_rgb, num_or_size_splits=3, axis=3)
    
89.         # 输入vgg的图像是bgr的顺序（跟opencv一样 倒序的）而不是rgb
    
90.         x_bgr = tf.concat(values=[
    
91.             b - VGG_MEAN[0],
    
92.             g - VGG_MEAN[1],
    
93.             r - VGG_MEAN[2]
    
94.         ], axis=3)
    
95. 
    
96.         assert x_bgr.get_shape().as_list()[1:] == [224, 224, 3]
    
97.         # stage 1
    
98.         self.conv1_1 = self.conv_layer(x_bgr, 'conv1_1')
    
99.         self.conv1_2 = self.conv_layer(self.conv1_1, 'conv1_2')
    
100.         self.pool1 = self.pooling_layer(self.conv1_2, 'pool1')
     
101.         # stage 2
     
102.         self.conv2_1 = self.conv_layer(self.pool1, 'conv2_1')
     
103.         self.conv2_2 = self.conv_layer(self.conv2_1, 'conv2_2')
     
104.         self.pool2 = self.pooling_layer(self.conv2_2, 'pool2')
     
105. 
     
106.         # stage 3
     
107.         self.conv3_1 = self.conv_layer(self.pool2, 'conv3_1')
     
108.         self.conv3_2 = self.conv_layer(self.conv3_1, 'conv3_2')
     
109.         self.conv3_3 = self.conv_layer(self.conv3_2, 'conv3_3')
     
110.         self.pool3 = self.pooling_layer(self.conv3_3, 'pool3')
     
111.         # stage 4
     
112.         self.conv4_1 = self.conv_layer(self.pool3, 'conv4_1')
     
113.         self.conv4_2 = self.conv_layer(self.conv4_1, 'conv4_2')
     
114.         self.conv4_3 = self.conv_layer(self.conv4_2, 'conv4_3')
     
115.         self.pool4 = self.pooling_layer(self.conv4_3, 'pool4')
     
116.         # stage 5
     
117.         self.conv5_1 = self.conv_layer(self.pool4, 'conv5_1')
     
118.         self.conv5_2 = self.conv_layer(self.conv5_1, 'conv5_2')
     
119.         self.conv5_3 = self.conv_layer(self.conv5_2, 'conv5_3')
     
120.         self.pool5 = self.pooling_layer(self.conv5_3, 'pool5')
     
121.         # stage 6
     
122.         # self.flatten5 = self.flatten_op(self.pool5, 'flatten5')
     
123.         # self.fc6 = self.fc_layer(self.flatten5, 'fc6')
     
124.         # self.fc7 = self.fc_layer(self.fc6, 'fc7')
     
125.         # self.fc8 = self.fc_layer(self.fc7, 'fc8', activation=None)
     
126.         # self.prob = tf.nn.softmax(self.fc8, name='prob')
     
127. 
     
128.         # print(self.prob.shape)
     
129. 
     
130.         print('buliding finished 耗时：%4ds' % (time.time() - start_time))
     
131.         pass
     
132. 
     
133. vgg16_npy_path = './vgg16.npy'
     
134. # vgg16_data = np.load(vgg16_npy_path, encoding='latin1')
     
135. # data_dict = vgg16_data.item()
     
136. #
     
137. # print(data_dict.keys())
     
138. #
     
139. # vgg16_for_result = VGGNet(data_dict)
     
140. #
     
141. # image_rgb = tf.placeholder(dtype=tf.float32, shape=[1, 224, 224, 3], name='image_rgb')
     
142. #
     
143. # vgg16_for_result.build(image_rgb)
     
144. 
     
145. # 224 * 224
     
146. content_img_path = './img/content.jpeg'
     
147. style_img_path = './img/style.jpeg'
     
148. 
     
149. num_steps = 100
     
150. learning_rate = 10
     
151. 
     
152. lambda_c = 0.1
     
153. lambda_s = 500
     
154. 
     
155. output_dir = './img/output_img'
     
156. 
     
157. if not os.path.exists(output_dir):
     
158.     os.mkdir(output_dir)
     
159. 
     
160. 
     
161. def initial_image(shape, mean, stddev):
     
162.     initial_image = tf.truncated_normal(shape=shape, mean=mean, stddev=stddev, dtype=tf.float32)
     
163.     return tf.Variable(initial_image)
     
164. 
     
165. 
     
166. initial_image_result = initial_image([1, 224, 224, 3], mean=127.5, stddev=20)
     
167. 
     
168. 
     
169. def read_img(image_name):
     
170.     img = Image.open(image_name)
     
171.     # [224 224 3]
     
172.     np_img = np.array(img)
     
173.     # np_img = tf.reshape(np_img, shape=[1, 224, 224, 3])
     
174.     # [1 224 224 3]
     
175.     # np_img = np.reshape(np_img, newshape=[1, 224, 224, 3])
     
176.     np_img = np.asarray([np_img], dtype=np.float32)
     
177.     print(np_img.shape)
     
178.     return np_img
     
179. 
     
180. # plt.imshow(read_img(style_img_path)[0])
     
181. # plt.show()
     
182. 
     
183. 
     
184. content_img_arr_val = read_img(content_img_path)
     
185. style_img_arr_val = read_img(style_img_path)
     
186. 
     
187. content_img = tf.placeholder(dtype=tf.float32, shape=[1, 224, 224, 3], name='content_img')
     
188. style_img = tf.placeholder(dtype=tf.float32, shape=[1, 224, 224, 3], name='style_img')
     
189. 
     
190. # initial_image_result content_img style_img 三张图片进入vgg网络 提取特征
     
191. vgg16_data = np.load(vgg16_npy_path, encoding='latin1')
     
192. data_dict = vgg16_data.item()
     
193. 
     
194. print(data_dict.keys())
     
195. 
     
196. vgg16_for_initial_result = VGGNet(data_dict)
     
197. vgg16_for_content_img = VGGNet(data_dict)
     
198. vgg16_for_style_img = VGGNet(data_dict)
     
199. 
     
200. vgg16_for_content_img.build(content_img)
     
201. vgg16_for_style_img.build(style_img)
     
202. vgg16_for_initial_result.build(initial_image_result)
     
203. 
     
204. # 定义提取哪些层的特征 cnn的
     
205. 
     
206. # 内容特征 越低层越精细
     
207. content_features = [
     
208.     vgg16_for_content_img.conv1_2,
     
209.     vgg16_for_content_img.conv2_2,
     
210.     # vgg16_for_content_img.conv3_3,
     
211.     # vgg16_for_content_img.conv4_3,
     
212.     # vgg16_for_content_img.conv5_3
     
213. ]
     
214. 
     
215. # 结果的内容特征必须与内容特征一致
     
216. result_content_features = [
     
217.     vgg16_for_initial_result.conv1_2,
     
218.     vgg16_for_initial_result.conv2_2,
     
219.     # vgg16_for_initial_result.conv3_3,
     
220.     # vgg16_for_initial_result.conv4_3,
     
221.     # vgg16_for_initial_result.conv5_3
     
222. ]
     
223. 
     
224. # 风格特征 越高层越抽象
     
225. style_features = [
     
226.     # vgg16_for_style_img.conv1_2,
     
227.     # vgg16_for_style_img.conv2_2,
     
228.     # vgg16_for_style_img.conv3_3,
     
229.     vgg16_for_style_img.conv4_3,
     
230.     vgg16_for_style_img.conv5_3
     
231. ]
     
232. 
     
233. # 结果的风格特征必须与风格特征一致
     
234. result_style_features = [
     
235.     # vgg16_for_initial_result.conv1_2,
     
236.     # vgg16_for_initial_result.conv2_2,
     
237.     # vgg16_for_initial_result.conv3_3,
     
238.     vgg16_for_initial_result.conv4_3,
     
239.     vgg16_for_initial_result.conv5_3
     
240. ]
     
241. 
     
242. # 开始计算损失
     
243. content_loss = tf.zeros(shape=1, dtype=tf.float32)
     
244. # zip([1, 2], [3, 4]) ---> [(1, 3), (2, 4)] 两个数组变成 一个数组
     
245. # c与c_的shape（卷积激励之后）：[1 height width channel] 在通道axis=[1, 2, 3]求平均
     
246. # loss = mse 平方差损失函数
     
247. for c, c_ in zip(content_features, result_content_features):
     
248.     content_loss += tf.reduce_mean(tf.square(c - c_), axis=[1, 2, 3])
     
249.     pass
     
250. 
     
251. # 计算Gram矩阵
     
252. def gram_matrix(x):
     
253.     """Gram矩阵计算 k个feature_map 两两之间的关联性 相似度 k*k的矩阵
     
254.     Args：
     
255.     ---x feature_map from Conv层 shape：[1 height width channels]    """
     
256.     b, h ,w, ch = x.get_shape().as_list()
     
257.     features = tf.reshape(x, shape=[b, h*w, ch])
     
258.     # [ch ch] = [ch h*w] 矩阵相乘 [h*w ch]
     
259.     # 除以维度相乘是为了防止 值过大 除以统一的数
     
260.     # features是三维的 导致gram矩阵也是三维的
     
261.     gram = tf.matmul(features, features, adjoint_a=True) / tf.constant(h*w*ch, tf.float32)
     
262.     # features是2维的 导致gram矩阵也是2维的
     
263.     # gram = tf.matmul(tf.matrix_transpose(features[0]), features[0]) / tf.constant(h*w*ch, tf.float32)
     
264.     return gram
     
265.     pass
     
266. 
     
267. 
     
268. style_gram_matrix = [gram_matrix(feature) for feature in style_features]
     
269. result_style_gram_matrix = [gram_matrix(feature) for feature in result_style_features]
     
270. 
     
271. style_loss = tf.zeros(shape=1, dtype=tf.float32)
     
272. for s, s_ in zip(style_gram_matrix, result_style_gram_matrix):
     
273.     # loss = mse 平方差损失函数
     
274.     # axis=[0, 1, 2, 3] 这里会报错  gram矩阵输出是三维的
     
275.     style_loss += tf.reduce_mean(tf.square(s - s_), axis=[1, 2])
     
276.     # gram矩阵输出是二维的
     
277.     # style_loss += tf.reduce_mean(tf.square(s - s_), axis=[0, 1])
     
278.     pass
     
279. 
     
280. loss = content_loss * lambda_c + style_loss * lambda_s
     
281. 
     
282. with tf.name_scope('train'):
     
283.     train_op = tf.train.AdamOptimizer(learning_rate).minimize(loss)
     
284. 
     
285. 
     
286. init_op = tf.global_variables_initializer()
     
287. 
     
288. with tf.Session() as sess:
     
289.     sess.run(init_op)
     
290.     for step in range(num_steps):
     
291.         loss_value, content_loss_value, style_loss_value, _ = \
     
292.         sess.run(fetches=[loss, content_loss, style_loss, train_op], feed_dict={
     
293.             content_img: content_img_arr_val,
     
294.             style_img: style_img_arr_val
     
295.         })
     
296.         print('step:%d loss_value:%8.4f content_loss_value:%8.4f style_loss_value:%8.4f'
     
297.               % (step+1, loss_value[0], content_loss_value[0], style_loss_value[0]))
     
298.         # result_image：shape [224, 224, 3]
     
299.         result_image = initial_image_result.eval(sess)[0]
     
300.         # np.clip值裁剪 小于0的变为0 大于255的变为255
     
301.         result_image = np.clip(result_image, 0, 255)
     
302.         result_image = np.asarray(result_image, dtype=np.uint8)
     
303.         # np_img = np.asarray([np_img], dtype=np.float32)
     
304.         img = Image.fromarray(result_image)
     
305.         result_image_path = os.path.join(output_dir, 'result-%05d.jpg' % (step + 1))
     
306.         img.save(result_image_path)
     
307. 
     
308.     pass
     
309. 
     
310. 
     
311. 
     

复制代码

xsoft

谢谢老师提供的资料。