-
Notifications
You must be signed in to change notification settings - Fork 10
/
styleTransfer.py
96 lines (83 loc) · 4.99 KB
/
styleTransfer.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
import tensorflow as tf
import tensorflow.contrib as contrib
import numpy as np
import scipy.io as sio
import scipy.misc as misc
from PIL import Image
def conv(inputs, w, b):
w = tf.constant(w)
b = tf.constant(b)
return tf.nn.conv2d(inputs, w, [1, 1, 1, 1], "SAME") + b
def mapping(img):
return 255.0 * (img - np.min(img)) / (np.max(img) - np.min(img))
class StyleTransfer:
def __init__(self, H=256, W=256, C=3, alpha=1e-3, beta=1.0, iteration=500, content_path="./content//content.jpg", style_path="./style//style.jpg"):
self.content_img = tf.placeholder("float", [1, H, W, C])
self.style_img = tf.placeholder("float", [1, H, W, C])
self.target_img = tf.get_variable("target", shape=[1, H, W, C], initializer=tf.truncated_normal_initializer(stddev=0.02))
feature_bank_x = self.Network_vgg(self.target_img)
feature_bank_style = self.Network_vgg(self.style_img)
feature_bank_content = self.Network_vgg(self.content_img)
self.L_content = self.content_loss(feature_bank_x, feature_bank_content)
self.L_style = self.style_loss(feature_bank_x, feature_bank_style)
self.total_loss = alpha * self.L_content + beta * self.L_style
# self.Opt = tf.train.AdamOptimizer(0.0002).minimize(self.total_loss)
#L-BFGS
self.optimizer = tf.contrib.opt.ScipyOptimizerInterface(self.total_loss, method='L-BFGS-B',options={'maxiter': iteration, 'disp': 0})
self.sess = tf.Session()
self.sess.run(tf.global_variables_initializer())
self.train(H, W, C, content_path, style_path)
def train(self, H, W, C, content_path, style_path):
content_img = np.reshape(misc.imresize(np.array(Image.open(content_path)), [H, W], mode="RGB"), [1, H, W, C])
style_img = np.reshape(misc.imresize(np.array(Image.open(style_path)), [H, W], mode="RGB"), [1, H, W, C])
self.sess.run(tf.assign(self.target_img, content_img), feed_dict={self.content_img: content_img, self.style_img: style_img})
self.optimizer.minimize(self.sess, feed_dict={self.content_img: content_img, self.style_img: style_img})
L_content = self.sess.run(self.L_content, feed_dict={self.content_img: content_img, self.style_img: style_img})
L_style = self.sess.run(self.L_style, feed_dict={self.content_img: content_img, self.style_img: style_img})
L_total = self.sess.run(self.total_loss, feed_dict={self.content_img: content_img, self.style_img: style_img})
print("L_content: %g, L_style: %g, L_total: %g" % (L_content, L_style, L_total))
target_img = self.sess.run(self.target_img,feed_dict={self.content_img: content_img, self.style_img: style_img})
Image.fromarray(np.uint8(mapping(np.reshape(target_img, [H, W, C])))).save("./deepdream/target.jpg")
def content_loss(self, feature_bank_x, feature_bank_content):
#content loss
#squared-error
return tf.reduce_sum(tf.square(feature_bank_x["relu4_2"] - feature_bank_content["relu4_2"])) / 2.0
def style_loss(self, feature_bank_x, feature_bank_style):
#style loss
E = 0
for layer in feature_bank_style.keys():
if layer == "relu1_1" or layer=="relu2_1" or layer=="relu3_1" or layer=="relu4_1" or layer=="relu5_1":
w = 0.2
else:
w = 0
C = int(feature_bank_x[layer].shape[-1])
H = int(feature_bank_x[layer].shape[1])
W = int(feature_bank_x[layer].shape[2])
F = tf.reshape(tf.transpose(feature_bank_x[layer], [0, 3, 1, 2]), shape=[C, -1])
#Gram matrix of x
G_x = tf.matmul(F, tf.transpose(F))
C = int(feature_bank_style[layer].shape[-1])
F = tf.reshape(tf.transpose(feature_bank_style[layer], [0, 3, 1, 2]), shape=[C, -1])
#Gram matrix of style
G_s = tf.matmul(F, tf.transpose(F))
E += w * tf.reduce_sum(tf.square(G_x - G_s)) / (4 * C**2 * H**2 * W**2)
return E
def Network_vgg(self, inputs):
vgg_para = sio.loadmat("./vgg_para//vgg.mat")
layers = vgg_para["layers"]
feature_bank = {}
with tf.variable_scope("vgg"):
for i in range(37):
if layers[0, i][0, 0]["type"] == "conv":
w = layers[0, i][0, 0]["weights"][0, 0]
b = layers[0, i][0, 0]["weights"][0, 1]
with tf.variable_scope(str(i)):
inputs = conv(inputs, w, b)
elif layers[0, i][0, 0]["type"] == "relu":
inputs = tf.nn.relu(inputs)
feature_bank[layers[0, i][0, 0]["name"][0]] = inputs
else:
inputs = tf.nn.max_pool(inputs, [1, 2, 2, 1], [1, 2, 2, 1], "SAME")
return feature_bank
if __name__ == "__main__":
st = StyleTransfer(H=512, W=512, C=3, alpha=1e-5, beta=1.0, iteration=500, content_path="./content//content.jpg", style_path="./style//style.jpg")