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
97
|
import tensorflow as tf
import numpy as np
from tensorflow.keras.layers import \
Conv2D, AveragePooling2D
from skimage import transform
import hyperparameters as hp
def get_gram(style_output):
style_shape = tf.shape(style_output)
output = tf.linalg.einsum('bijc,bijd->bcd', style_output, style_output)
dimensions = style_shape[1] * style_shape[2]
dimensions = tf.cast(dimensions, tf.float32)
return output / dimensions
class YourModel(tf.keras.Model):
""" Your own neural network model. """
def __init__(self, content_image, style_image): #normalize these images to float values
super(YourModel, self).__init__()
self.content_image = transform.resize(content_image, tf.shape(style_image), anti_aliasing=True, preserve_range=True)
self.content_image = tf.image.convert_image_dtype(self.content_image, tf.float32)
self.style_image = transform.resize(style_image, tf.shape(style_image), anti_aliasing=True, preserve_range=True)
self.style_image = tf.image.convert_image_dtype(self.style_image, tf.float32)
image = tf.image.convert_image_dtype(content_image, tf.float32)
self.x = tf.Variable([image])
self.content_weight = hp.alpha
self.style_weight = hp.beta
self.photo_layers = ['block5_conv2']
self.style_layers = ['block1_conv1', 'block2_conv1', 'block3_conv1', 'block4_conv1', 'block5_conv1']
self.num_photo_layers = len(self.photo_layers)
self.num_style_layers = len(self.style_layers)
self.optimizer = tf.keras.optimizers.Adam(learning_rate=hp.learning_rate, beta_1=hp.beta_1, epsilon=hp.epsilon)
self.vgg16 = tf.keras.applications.VGG16(include_top=False, weights='vgg16_imagenet.h5')
self.vgg16.trainable = False
# creating the Gram Matrix
p_output = self.vgg16.get_layer(self.photo_layers[0]).output
style_output = []
for layer in self.style_layers:
style_output.append(self.vgg16.get_layer(layer).output)
G = [get_gram(x) for x in style_output]
self.vgg16 = tf.keras.Model([self.vgg16.input], [p_output, G])
# figure this out Michael
img_to_np = lambda img: np.array([img * 255])
self.content_target = self.vgg16(img_to_np(content_image))[0]
self.style_target = self.vgg16(img_to_np(style_image))[1]
# create a map of the layers to their corresponding number of filters if it is a convolutional layer
def call(self, x):
# call onto our pretrained network, since we don't have a classifcation head to follow
x = self.vgg16(x * 255)
return x
def loss_fn(self, x):
x = self.call(x)
content_l = self.content_loss(x[0], self.content_target)
style_l = self.style_loss(x[1], self.style_target)
return (self.content_weight * content_l) + (self.style_weight * style_l)
def content_loss(self, photo_layers, input_layers):
return tf.reduce_mean(tf.square(photo_layers - input_layers))
def style_loss(self, art_layers, input_layers):
layer_losses = []
for created, target in zip(art_layers, input_layers):
reduced = tf.reduce_mean(tf.square(created - target))
layer_losses.append(reduced)
return tf.add_n(layer_losses)
def train_step(self, epoch):
with tf.GradientTape(watch_accessed_variables=False) as tape:
tape.watch(self.x)
# loss = self.loss_fn(self.content_image, self.style_image, self.x)
loss = self.loss_fn(self.x)
print('\rEpoch {}: Loss: {:.4f}'.format(epoch, loss), end='')
gradients = tape.gradient(loss, self.x)
self.optimizer.apply_gradients([(gradients, self.x)])
self.x.assign(tf.clip_by_value(self.x, clip_value_min=0.0, clip_value_max=1.0))
|
