Fast Style Transfer with FastEstimator¶
[Paper] [Notebook] [TF Implementation] [Torch Implementation]
In this notebook we will demonstrate how to do a neural image style transfer with perceptual loss as described in Perceptual Losses for Real-Time Style Transfer and Super-Resolution. Typical neural style transfer involves two images: an image containing semantics that you want to preserve, and another image serving as a reference style. The first image is often referred as the content image and the other image as the style image. In this paper training images from the COCO2014 dataset are used to learn style transfer from any content image.
import tempfile
import cv2
import tensorflow as tf
import numpy as np
import fastestimator as fe
from fastestimator.backend import reduce_mean
from fastestimator.op.numpyop import LambdaOp
from fastestimator.op.numpyop.multivariate import Resize
from fastestimator.op.numpyop.univariate import Normalize, ReadImage
from fastestimator.trace.io import ModelSaver
from fastestimator.util import ImageDisplay, GridDisplay
from matplotlib import pyplot as plt
#Parameters
batch_size = 4
epochs = 2
train_steps_per_epoch = None
log_steps = 2000
style_weight=5.0
content_weight=1.0
tv_weight=1e-4
save_dir = tempfile.mkdtemp()
style_img_path = 'Vassily_Kandinsky,_1913_-_Composition_7.jpg'
test_img_path = 'panda.jpeg'
data_dir = None
In this notebook we will use Vassily Kandinsky's Composition 7 as a style image. We will also resize the style image to $256 \times 256$ to make the dimension consistent with that of COCO images.
style_img = cv2.imread(style_img_path)
assert style_img is not None, "cannot load the style image, please go to the folder with style image"
style_img = cv2.resize(style_img, (256, 256))
style_img = (style_img.astype(np.float32) - 127.5) / 127.5
style_img_disp = cv2.cvtColor((style_img + 1) * 0.5, cv2.COLOR_BGR2RGB)
ImageDisplay(image=style_img_disp, title='Vassily Kandinsky\'s Composition 7').show()