loss_functions.py

def gram_matrix(x):
  x = tf.transpose(x, (2, 0, 1))
  features = tf.reshape(x, (tf.shape(x)[0], -1))
  gram = tf.matmul(features, tf.transpose(features))
  return gram
  
def style_loss(style, combination):
  S = gram_matrix(style)
  C = gram_matrix(combination)
  channels = 3
  size = img_rows * img_cols
  return tf.reduce_sum(tf.square(S - C)) / (channels * (3 ** 2) * (size ** 2))

def content_loss(base, combination):
  return tf.reduce_sum(tf.square(combination - base))

def total_variation_loss(x):
  a = tf.square(x[:, : img_rows - 1, : img_cols - 1, :] - x[:, 1:, : img_cols - 1, :])
  b = tf.square(x[:, : img_rows - 1, : img_cols - 1, :] - x[:, : img_rows - 1, 1:, :])
  return tf.reduce_sum(tf.pow(a + b, 1.25))