Test for the image style transfer by Ghiasi et al.

tf-image-stilizer.py

# Check it working on: https://www.kaggle.com/matiasmiche/image-stylization

import matplotlib.pyplot as plt
import tensorflow_hub as hub
import tensorflow as tf
import numpy as np
import cv2

content_filename = 'content_image.jpg'
style_filename = 'style_image.jpg'

model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')

def load_image(img_path):
	img = tf.io.read_file(img_path)
	img = tf.image.decode_image(img, channels = 3)
	img = tf.image.convert_image_dtype(img, tf.float32)
	img = img[tf.newaxis, :]
	return img

content_image = load_image(content_filename)
style_image = load_image(style_filename)

stylized_image = model(tf.constant(content_image), tf.constant(style_image))[0]

plt.imshow(np.squeeze(stylized_image))
plt.show()

cv2.imwrite('generated.jpg', cv2.cvtColor(np.squeeze(stylized_image)*255, cv2.COLOR_BGR2RGB))