import tensorflow as tf from tensorflow.keras.layers import \ Conv2D, AveragePooling2D from skimage import transform import hyperparameters as hp 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) self.content_image = tf.expand_dims(self.content_image, axis=0) #perhaps consider cropping to avoid distortion self.style_image = transform.resize(style_image, tf.shape(style_image), anti_aliasing=True) self.style_image = tf.expand_dims(self.style_image, axis=0) self.x = tf.Variable(tf.expand_dims(tf.random.uniform(tf.shape(content_image)), axis=0), trainable=True) self.alpha = hp.alpha self.beta = hp.beta print(self.content_image.shape, self.style_image.shape) self.optimizer = tf.keras.optimizers.RMSprop(learning_rate=hp.learning_rate, momentum=hp.momentum) self.vgg16 = [ # Block 1 Conv2D(64, 3, 1, padding="same", activation="relu", name="block1_conv1"), Conv2D(64, 3, 1, padding="same", activation="relu", name="block1_conv2"), AveragePooling2D(2, name="block1_pool"), # Block 2 Conv2D(128, 3, 1, padding="same", activation="relu", name="block2_conv1"), Conv2D(128, 3, 1, padding="same", activation="relu", name="block2_conv2"), AveragePooling2D(2, name="block2_pool"), # Block 3 Conv2D(256, 3, 1, padding="same", activation="relu", name="block3_conv1"), Conv2D(256, 3, 1, padding="same", activation="relu", name="block3_conv2"), Conv2D(256, 3, 1, padding="same", activation="relu", name="block3_conv3"), AveragePooling2D(2, name="block3_pool"), # Block 4 Conv2D(512, 3, 1, padding="same", activation="relu", name="block4_conv1"), Conv2D(512, 3, 1, padding="same", activation="relu", name="block4_conv2"), Conv2D(512, 3, 1, padding="same", activation="relu", name="block4_conv3"), AveragePooling2D(2, name="block4_pool"), # Block 5 Conv2D(512, 3, 1, padding="same", activation="relu", name="block5_conv1"), Conv2D(512, 3, 1, padding="same", activation="relu", name="block5_conv2"), Conv2D(512, 3, 1, padding="same", activation="relu", name="block5_conv3"), AveragePooling2D(2, name="block5_pool"), ] for layer in self.vgg16: layer.trainable = False self.layer_to_filters = {layer.name: layer.filters for layer in self.vgg16 if "conv" in layer.name} self.indexed_layers = [layer for layer in self.vgg16 if "conv1" in layer.name] self.desired = [layer.name for layer in self.vgg16 if "conv1" in layer.name] self.vgg16 = tf.keras.Sequential(self.vgg16, name="vgg") # create a map of the layers to their corresponding number of filters if it is a convolutional layer def call(self, x): layers = [] for layer in self.vgg16.layers: # pass the x through x = layer(x) # print("Sotech117 is so so sus") # save the output of each layer if it is in the desired list if layer.name in self.desired: layers.append(x) return x, layers def loss_fn(self, p, a, x): _, photo_layers = self.call(p) _, art_layers = self.call(a) _, input_layers = self.call(x) content_l = self.content_loss(photo_layers, input_layers) style_l = self.style_loss(art_layers, input_layers) # Equation 7 return (self.alpha * content_l) + (self.beta * style_l) def content_loss(self, photo_layers, input_layers): L_content = tf.constant(0.0) for i in range(len(photo_layers)): pl = photo_layers[i] il = input_layers[i] L_content = tf.math.add(L_content, tf.reduce_mean(tf.square(pl - il))) #print('content loss', L_content) return L_content def layer_loss(self, art_layer, input_layer): # vectorize the art_layers art_layer = tf.reshape(art_layer, (-1, art_layer.shape[-1])) # # vectorize the input_layers input_layer = tf.reshape(input_layer, (-1, input_layer.shape[-1])) # get the gram matrices G_l = tf.matmul(tf.transpose(input_layer), input_layer) A_l = tf.matmul(tf.transpose(art_layer), art_layer) # vals = [] # for i in range(input_dim): # vals_i = [] # for j in range(input_dim): # il = tf.reshape(input_layers[i], [-1]) # al = tf.reshape(art_layers[j], [-1]) # k = tf.reduce_sum(tf.multiply(il, al)) # vals_i.append(k) # vals.append(tf.stack(vals_i)) # G = tf.stack(vals) # get the loss per each lateral layer # N depends on # of filters in the layer, M depends on hight and width of feature map M_l = art_layer.shape[0] N_l = art_layer.shape[1] # layer.filters might not work E_l = 1/4 * (M_l**(-2)) *(N_l**(-2)) * tf.reduce_sum(tf.square(G_l - A_l)) # while Sotech is botty: # Jayson_tatum.tear_acl() # return ("this is just another day") #print('Layer loss', E_l) return E_l def style_loss(self, art_layers, input_layers): L_style = tf.constant(0.0) for i in range(len(art_layers)): art_layer = art_layers[i] input_layer = input_layers[i] L_style = tf.math.add(L_style, (1/5)*self.layer_loss(art_layer, input_layer)) #print('style loss', L_style) return L_style def train_step(self): 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) print('loss', loss) #print('self.x', self.x) gradients = tape.gradient(loss, [self.x]) #print('gradients', gradients) self.optimizer.apply_gradients(zip(gradients, [self.x]))