Neural Style Transfer in Python

Neural Style Transfer (NST) is a deep learning technique that allows us to combine the artistic style of one image with the content of another, creating unique, stylized images from ordinary photos. NST opens up a world of creative possibilities!

What is NST?

NST is a process that uses Convolutional Neural Networks (CNNs) to blend the content and style of two separate images. Content refers to an image's structural details, such as the shapes, objects, or people it contains. Style, conversely, encompasses the textures, colors, brush strokes, and patterns of a specific artwork or artist.

For example, you could use NST to create an interesting image by blending the content of Leonardo da Vinci's Mona Lisa with the style of Vincent van Gogh's The Starry Night.

Mona Lisa by Leonardo da Vinci (Content)

Starry Night by Vincent van Gogh (Style)

Target Image

Code

Let’s code a Neural Style Transfer in Python using PyTorch! This project will be coded in a Google Collaboratory Notebook.

First, we need to import the necessary libraries:

# Import necessary libraries

import torch
from torchvision import transforms, models
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np

# Use cuda if it is available
device = ("cuda" if torch.cuda.is_available() else "cpu")

For our model, we are going to use VGG-19. VGG-19 is a convolutional neural network that is 19 layers deep. It has been trained on over a million images from the ImageNet database. We want to use the pre-trained model, so we freeze the model’s parameters, preventing them from being updated during training:

# Use VGG-19
# VGG-19 is a CNN that is 19 layers deep that has been pre-trained on more than a million images from the ImageNet database

model = models.vgg19(pretrained=True).features
for p in model.parameters():
  p.requires_grad = False
model.to(device)

# Print the layers
print(model)

Once we have the model, we can print the layers:

Sequential(
  (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (1): ReLU(inplace=True)
  (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (3): ReLU(inplace=True)
  (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (6): ReLU(inplace=True)
  (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (8): ReLU(inplace=True)
  (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (11): ReLU(inplace=True)
  (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (13): ReLU(inplace=True)
  (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (15): ReLU(inplace=True)
  (16): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (17): ReLU(inplace=True)
  (18): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (19): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (20): ReLU(inplace=True)
  (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (22): ReLU(inplace=True)
  (23): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (24): ReLU(inplace=True)
  (25): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (26): ReLU(inplace=True)
  (27): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (29): ReLU(inplace=True)
  (30): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (31): ReLU(inplace=True)
  (32): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (33): ReLU(inplace=True)
  (34): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (35): ReLU(inplace=True)
  (36): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
)

Next, we want to create a function that will return the activation features of the model given an input tensor. We will use this to get the activation features of the content and style images. We will mainly use the layers to capture style and content loss. This function iterates through the model's layers, extracting activations at specific layers and storing them in the features dictionary.

def model_activations(input,model):
  # 6 Activation layers
  # Layers where the dimensions are changed
  layers = {
    '0' : 'conv1_1',
    '5' : 'conv2_1',
    '10': 'conv3_1',
    '19': 'conv4_1',
    '21': 'conv4_2',
    '28': 'conv5_1'
  }
  features = {}
  x = input
  x = x.unsqueeze(0)
  for name,layer in model._modules.items():
      x = layer(x)
      if name in layers:
          features[layers[name]] = x

  return features

We also want to ensure that the images are a manageable resolution, so we will use the transforms in PyTorch to Resize images to 400x400 pixels, transform them into tensors, and normalize the data to make it easier for the model to understand:

transform = transforms.Compose([transforms.Resize(400),
                              transforms.ToTensor(),
                              transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))])

Now, we have to import images into the project. Because we are using a Google Colab Notebook, we will do this using Google Drive:

from google.colab import drive
drive.mount('/content/drive')

# Change the Image Folder Path to wherever the images are in your Google Drive
path = "drive/My Drive/[Image Folder Path]"

content = Image.open(path + "content.jpg").convert("RGB")
content = transform(content).to(device)
print("Content shape: ", content.shape)
style = Image.open(path + "style.jpg").convert("RGB")
style = transform(style).to(device)
print("Style shape: ", style.shape)

From the output, we can see that both images have three color channels (Red, Green, and Blue) and two dimensions.

Content shape:  torch.Size([3, 595, 400])
Style shape:  torch.Size([3, 400, 490])

We also want to create a function that can convert image tensors to a format that MatPlotLib can display:

def image_convert(image):
    x = image.to("cpu").clone().detach().numpy().squeeze()
    x = x.transpose(1,2,0)
    x = x*np.array((0.5,0.5,0.5)) + np.array((0.5,0.5,0.5))
    return np.clip(x,0,1)

We can use this function to display the content and style images:

fig, (ax1,ax2) = plt.subplots(1,2)

ax1.imshow(image_convert(content),label = "Content")
ax2.imshow(image_convert(style),label = "Style")
plt.show()

Content and Style Images displayed by MatPlotLib

Define a function that performs matrix multiplication of a feature with its transpose to create the gram matrix. The gram matrix captures the relationships between different feature maps within the network, which helps represent style elements like color and texture:

def gram_matrix(imgfeature):
    _,d,h,w = imgfeature.size()
    imgfeature = imgfeature.view(d,h*w)
    gram_mat = torch.mm(imgfeature,imgfeature.t())

    return gram_mat

We want to create a target image with the same dimensions as the content image and extract features from the style and content images using the VGG-19 model:

target = content.clone().requires_grad_(True).to(device)

# Get the features of the content and style images
style_features = model_activations(style, model)
content_features = model_activations(content, model)

We also want to define the style weights for the model. Experiment with the weights, epochs, and learning rate for the best results. Content is captured by the higher layers, while the lower layers capture style:

style_weight_measurements = {"conv1_1" : 1.0, "conv2_1" : 0.8, "conv3_1" : 0.4, "conv4_1" : 0.2, "conv5_1" : 0.1}

style_grams = {layer:gram_matrix(style_features[layer]) for layer in style_features}

# To accurately capture style, the style weight has to be a lot higher than the content weight
content_weight = 1000
style_weight = 1e8

print_after = 200
epochs = 4000
learning_rate = 0.1

# Use the Adam Optimizer
optimizer = torch.optim.Adam([target],lr=learning_rate)

content_weight: Controls the influence of the content image on the final output.

Get Rishab Alagharu’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

style_weight: Controls the influence of the style image on the final output.

epochs: The total number of training iterations.

learning_rate: Determines the step size during optimization.

Finally, we can generate a target image from the content and style images. We want to optimize for the number of epochs. For every iteration/epoch, we want to:

• Extract features from the current target image.
• Calculate the content loss by comparing the features of the target image to those of the content image.
• The content loss is calculated by:

Content Loss Formula

• Calculate the style loss by comparing the gram matrices of the target image to the gram matrices of the style image.
• Find the total loss by combining the content and style loss using the specified weights.
• For every 10th epoch, we want to print the loss.
• Calculate the gradients of the loss of the target image’s pixels and update the target image based on the calculated gradients.
• Every 200 epochs, we want to display the new image. Change the value of print_after to see more target images.

# Iterate for the number of epochs (Optimization)
for i in range(1, epochs+1):
    target_features = model_activations(target, model)

    # Get the average loss which is calculated by the (content - target)^2
    content_loss = torch.mean((content_features['conv4_2']-target_features['conv4_2'])**2)

    style_loss = 0

    # Iterate through each of the style weights
    for layer in style_weight_measurements:
        style_gram = style_grams[layer]
        target_gram = target_features[layer]
        _,d,w,h = target_gram.shape
        target_gram = gram_matrix(target_gram)

        style_loss += (style_weight_measurements[layer]*torch.mean((target_gram-style_gram)**2))/d*w*h # Normalize with depth, width, and height

    total_loss = (content_weight * content_loss) + (style_weight * style_loss)

    # Print every 10th epoch
    if i % 10 == 0:
        print("epoch ", i ," ", total_loss)

    optimizer.zero_grad()
    total_loss.backward()
    optimizer.step()

    # Display the image after print_after epochs
    if i % print_after == 0:
        plt.imshow(image_convert(target),label="Epoch "+str(i))
        plt.show()
        plt.imsave(str(i)+'.png',image_convert(target),format='png')

If you want to see the complete code, this is the Google Colab Notebook: https://colab.research.google.com/drive/1FLn333yCPvPvFRVI6ceC4SbPjhbIc6ay?usp=sharing

Limitations and Challenges of NST

Despite its many uses, NST also has some challenges and limitations:

• Computational Intensity: The optimization process in NST is resource-intensive, making it impractical for real-time applications without specialized hardware. Also, in our project, the image's resolution is limited due to computational constraints.

• Content Complexity: The algorithm may struggle with more complex or less defined images, where content and style differentiation becomes challenging.

• Control Over Output: Fine-tuning the final output is often tricky, as adjusting style and content weights can produce unpredictable results.

Conclusion

Neural Style Transfer represents a powerful fusion of artificial intelligence and art, allowing for creative expression through machine learning. The ability to create new images that blend content and artistic style has expanded possibilities in digital art and media and contributed to our understanding of image processing and neural networks.

Thank you for reading, and good luck!