neural_style_transfer.py

# -*- coding: utf-8 -*-
"""Neural Style Transfer.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1wg2J2yPaUFq1kEN0kiHvwxSpH89THDek
"""

# Commented out IPython magic to ensure Python compatibility.
import numpy as np
import tensorflow as tf
# import matplotlib.pyplot as plt
import cv2

import os
import tensorflow.keras.backend as K
import sys
import time
import PIL
# %matplotlib inline

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

MODEL_PATH = os.path.abspath('model/style_transfer.h5')

n_C = 3
IMAGE_WIDTH, IMAGE_HEIGHT = 600,600


style_targets = None
content_targets = None

content_layers = ['block4_conv2','block5_conv2']
style_layers=['block1_conv1','block2_conv1','block3_conv1','block4_conv1']
style_weights = [1,1,1,0.7]
num_content_layers = len(content_layers)
num_style_layers = len(style_layers)
alpha = 1e3
beta = 1e8

total_variation_weight=30
opt = tf.optimizers.Adam(learning_rate=0.02, beta_1=0.99, epsilon=1e-1)


model = None

def load_model():
    global model
    if model == None:
        model = tf.keras.models.load_model(MODEL_PATH)
        print('Model loaded')
    return model


def content_loss(c,g):
  return K.sum(K.square(c-g))

def gram_matrix(v):
  v = tf.reshape(v,[-1,n_C])
  
  return K.dot(K.transpose(v),v)

def Gram_Matrix(x):
  features = K.batch_flatten(K.permute_dimensions(x, (2, 0, 1)))
  gram = K.dot(features, K.transpose(features))
  return gram

def style_loss(s,g,w):
  S = Gram_Matrix(s[0])
  C = Gram_Matrix(g[0])
  size = IMAGE_WIDTH * IMAGE_HEIGHT
  return w * K.sum(K.square(S - C)) / (4.0 * (n_C ) * (size))

def high_pass_x_y(image):
  x_var = image[:,:,1:,:] - image[:,:,:-1,:]
  y_var = image[:,1:,:,:] - image[:,:-1,:,:]
  
  return x_var, y_var

def total_variation_loss(image):
  x_deltas, y_deltas = high_pass_x_y(image)
  return tf.reduce_sum(tf.abs(x_deltas)) + tf.reduce_sum(tf.abs(y_deltas))



def clip_0_1(image):
  return tf.clip_by_value(image, clip_value_min=0.0, clip_value_max=1.0)



def total_loss(outputs):
  style_outputs = outputs[:num_style_layers]
  content_outputs = outputs[num_style_layers:]
  c_loss = K.sum([content_loss(content,content_targets[i]) for i,content in enumerate(content_outputs)])
  s_loss = K.sum([style_loss(style,style_targets[i],style_weights[i]) for i,style in enumerate(style_outputs)])
  loss = alpha*c_loss + beta* s_loss

  return loss

def compute_content_loss_only(outputs):
  content_outputs =  content_outputs = outputs[num_style_layers:]
  c_loss = K.sum([content_loss(content,content_targets[i]) for i,content in enumerate(content_outputs)])
  loss = alpha * c_loss

  return loss

def compute_style_loss_only(outputs):
  style_outputs = outputs[:num_style_layers]
  s_loss = K.sum([style_loss(style,style_targets[i],style_weights[i]) for i,style in enumerate(style_outputs)])
  loss = s_loss

  return loss

def load_and_preprocess(path):
  img = cv2.imread(path)
  print(img.shape)
  img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
  img = cv2.resize(img,(IMAGE_WIDTH,IMAGE_HEIGHT),cv2.INTER_CUBIC)
  origi = img.copy()
  img = img.astype(np.float32)
  img = img/255.0
  # img = preprocess_input(img)
  img = img.reshape(-1,IMAGE_WIDTH,IMAGE_HEIGHT,n_C)
  img = tf.Variable(img,dtype=np.float32)
  return img,origi



# @tf.function()
def train_step(img):
	# keep track of our gradients
  with tf.GradientTape() as tape: 
    
    pred = model(img)
    loss = total_loss(pred)
    # loss += total_variation_weight*tf.image.total_variation(img)

    
  #Calculating gradinents of all trainable parameters
  grads = tape.gradient(loss, img)
  #Appling gradient optimization step
  opt.apply_gradients([(grads, img)])
  img.assign(clip_0_1(img))
  loss = np.mean(loss)
  return loss



def generate_image():
  img = np.random.randn(IMAGE_WIDTH,IMAGE_HEIGHT,n_C)

  img = img.astype(np.float32)
  img = img.reshape(-1,IMAGE_WIDTH,IMAGE_HEIGHT,n_C)
  img = tf.Variable(img,dtype=np.float32)
  return img


def tensor_to_image(tensor):
  tensor = tensor*255
  tensor = np.array(tensor, dtype=np.uint8)
  if np.ndim(tensor)>3:
    assert tensor.shape[0] == 1
    tensor = tensor[0]
  return PIL.Image.fromarray(tensor)




def perform_transformation(content_image,style_image,output_name):
  print("starting task")
  
  global style_targets, content_targets
  model=load_model()
  img,origi = load_and_preprocess(content_image)
  style_img,origi_style = load_and_preprocess(style_image)

  style_targets = model(style_img)[:num_style_layers]
  content_targets = model(img)[num_style_layers:]
  
  print("targets generated")
  
  g = generate_image()
  g_targets = model(g)
  print('starting generation')

  epochs = 1
  steps_per_epoch = 1

  step = 0
  print("Starting")

  for n in range(epochs):
    print(f"EPOCH:{n+1}")
    for m in range(steps_per_epoch):
      step += 1
      loss = train_step(img)

  print("Complete")
  art = img.numpy().squeeze()
  art = art * 255.0

  art = art.astype(np.uint8)
  art = cv2.cvtColor(art,cv2.COLOR_RGB2BGR)

  cv2.imwrite(output_name,art)
  return content_image, style_image,output_name