examples/neural-style-transfer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "DSC160 Data Science and the Arts - Twomey - Spring 2020 - [dsc160.roberttwomey.com](http://dsc160.roberttwomey.com)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Neural Style Transfer with PyTorch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook walks through a replication of the Neural Style Transfer from [Image Style Transfer Using Convolutional Neural Networks](https://www.cv-foundation.org/openaccess/content_cvpr_2016/papers/Gatys_Image_Style_Transfer_CVPR_2016_paper.pdf) by Gatys et. al (2016) using pytorch.\n",
    "\n",
    "This implementation uses a pre-trained VGG19 model as the basis for the style and content features used in optimization. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision.models as models\n",
    "import cv2\n",
    "import os\n",
    "\n",
    "import copy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "make a directory to store our output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists('./debug/'):\n",
    "     os.makedirs('./debug/')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "do we have a GPU? if so, use it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Input Images\n",
    "\n",
    "Write helper function to load images, resize to desired image size for input to VGG19 (512x512 with GPU), and also scale image intensities from 0-255 (as loaded with PIL) to 0.0-1.0 with using torch's [`transform.ToTensor()`](https://pytorch.org/docs/stable/torchvision/transforms.html#torchvision.transforms.ToTensor)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imsize = 512 if torch.cuda.is_available() else 128  # use small size if no gpu\n",
    "\n",
    "# create an input image pre-processer\n",
    "loader = transforms.Compose([\n",
    "    transforms.Resize(imsize),  # scale imported image\n",
    "    transforms.ToTensor()])  # transform it into a torch tensor (changes range to 0 - 1.0)\n",
    "\n",
    "\n",
    "def image_loader(image_name):\n",
    "    image = Image.open(image_name)\n",
    "    # fake batch dimension required to fit network's input dimensions\n",
    "    image = loader(image).unsqueeze(0)\n",
    "    return image.to(device, torch.float)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Read in our style and content images (stored in the img subdirectory):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "style_img = image_loader(\"./img/style_0.png\")\n",
    "content_img = image_loader(\"./img/content_0.png\")\n",
    "\n",
    "assert style_img.size() == content_img.size(), \\\n",
    "    \"we need to import style and content images of the same size\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "make a helper function to show images: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "unloader = transforms.ToPILImage()  # reconvert into PIL image\n",
    "\n",
    "plt.ion() # interactive mode\n",
    "\n",
    "def imshow(tensor, title=None):\n",
    "    image = tensor.cpu().clone()  # we clone the tensor to not do changes on it\n",
    "    image = image.squeeze(0)      # remove the fake batch dimension\n",
    "    image = unloader(image)\n",
    "    plt.imshow(image)\n",
    "    if title is not None:\n",
    "        plt.title(title)\n",
    "    plt.pause(0.001) # pause a bit so that plots are updated\n",
    "\n",
    "\n",
    "plt.figure()\n",
    "imshow(style_img, title='Style Image')\n",
    "\n",
    "plt.figure()\n",
    "imshow(content_img, title='Content Image')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loss Functions\n",
    "\n",
    "Neural style transfer minimizes both a style loss (Gram matrices) and a content loss (conv feature map differences). Next we will define these losses.\n",
    "\n",
    "### Content Loss\n",
    "\n",
    "Mean Squared Error (MSE) between content features and current features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ContentLoss(nn.Module):\n",
    "\n",
    "    def __init__(self, target,):\n",
    "        super(ContentLoss, self).__init__()\n",
    "        # we 'detach' the target content from the tree used\n",
    "        # to dynamically compute the gradient: this is a stated value,\n",
    "        # not a variable. Otherwise the forward method of the criterion\n",
    "        # will throw an error.\n",
    "        self.target = target.detach()\n",
    "\n",
    "    def forward(self, input):\n",
    "        self.loss = F.mse_loss(input, self.target)\n",
    "        return input"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Style Loss\n",
    "\n",
    "Style loss is based off of the difference in Gram matrices beween style reference and current image. The gram matrix for a given layer of the network is made by taking each of the (b) feature maps at that layer, linearizing them (flattening each c * d image into a single N = c*d vector), putting those into a new matrix, and calculating it's product with its transpose."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gram_matrix(input):\n",
    "    a, b, c, d = input.size()  # a=batch size(=1)\n",
    "    # b=number of feature maps\n",
    "    # (c,d)=dimensions of a f. map (N=c*d)\n",
    "\n",
    "    features = input.view(a * b, c * d)  # resise F_XL into \\hat F_XL\n",
    "\n",
    "    G = torch.mm(features, features.t())  # compute the gram product\n",
    "\n",
    "    # we 'normalize' the values of the gram matrix\n",
    "    # by dividing by the number of element in each feature maps.\n",
    "    return G.div(a * b * c * d)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Mean Squared Error (MSE) between target gram matrix and current gram matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class StyleLoss(nn.Module):\n",
    "\n",
    "    def __init__(self, target_feature):\n",
    "        super(StyleLoss, self).__init__()\n",
    "        self.target = gram_matrix(target_feature).detach()\n",
    "\n",
    "    def forward(self, input):\n",
    "        G = gram_matrix(input)\n",
    "        self.loss = F.mse_loss(G, self.target)\n",
    "        return input"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model\n",
    "\n",
    "Load our pretrained VGG19 model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cnn = models.vgg19(pretrained=True).features.to(device).eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "VGG19 is trained on images that are normalized with mean and std as given below. We need to normalize our images to have the same mean and std."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cnn_normalization_mean = torch.tensor([0.485, 0.456, 0.406]).to(device)\n",
    "cnn_normalization_std = torch.tensor([0.229, 0.224, 0.225]).to(device)\n",
    "\n",
    "# create a module to normalize input image so we can easily put it in a\n",
    "# nn.Sequential\n",
    "class Normalization(nn.Module):\n",
    "    def __init__(self, mean, std):\n",
    "        super(Normalization, self).__init__()\n",
    "        # .view the mean and std to make them [C x 1 x 1] so that they can\n",
    "        # directly work with image Tensor of shape [B x C x H x W].\n",
    "        # B is batch size. C is number of channels. H is height and W is width.\n",
    "        self.mean = torch.tensor(mean).view(-1, 1, 1)\n",
    "        self.std = torch.tensor(std).view(-1, 1, 1)\n",
    "\n",
    "    def forward(self, img):\n",
    "        # normalize img\n",
    "        return (img - self.mean) / self.std"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get Style Model and Losses\n",
    "<a name='select-layers' /></a>\n",
    "\n",
    "Content loss is calculated based on a single convolutional layer below, style loss is calculated based on the gram matrices for features from 5 different convolutional layers as seen below. We need to specify which VGG19 layers we want to use:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "content_layers_default = ['conv_4']\n",
    "style_layers_default = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following section takes the given parameters above (cnn network, normalization parameters, input images, and lists of layers) and returns the style network and the list of loss objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_style_model_and_losses(cnn, normalization_mean, normalization_std,\n",
    "                               style_img, content_img,\n",
    "                               content_layers=content_layers_default,\n",
    "                               style_layers=style_layers_default):\n",
    "    cnn = copy.deepcopy(cnn)\n",
    "\n",
    "    # normalization module\n",
    "    normalization = Normalization(normalization_mean, normalization_std).to(device)\n",
    "\n",
    "    # just in order to have an iterable access to or list of content/syle\n",
    "    # losses\n",
    "    content_losses = []\n",
    "    style_losses = []\n",
    "\n",
    "    # assuming that cnn is a nn.Sequential, so we make a new nn.Sequential\n",
    "    # to put in modules that are supposed to be activated sequentially\n",
    "    model = nn.Sequential(normalization)\n",
    "\n",
    "    i = 0  # increment every time we see a conv\n",
    "    for layer in cnn.children():\n",
    "        if isinstance(layer, nn.Conv2d):\n",
    "            i += 1\n",
    "            name = 'conv_{}'.format(i)\n",
    "        elif isinstance(layer, nn.ReLU):\n",
    "            name = 'relu_{}'.format(i)\n",
    "            # The in-place version doesn't play very nicely with the ContentLoss\n",
    "            # and StyleLoss we insert below. So we replace with out-of-place\n",
    "            # ones here.\n",
    "            layer = nn.ReLU(inplace=False)\n",
    "        elif isinstance(layer, nn.MaxPool2d):\n",
    "            name = 'pool_{}'.format(i)\n",
    "        elif isinstance(layer, nn.BatchNorm2d):\n",
    "            name = 'bn_{}'.format(i)\n",
    "        else:\n",
    "            raise RuntimeError('Unrecognized layer: {}'.format(layer.__class__.__name__))\n",
    "\n",
    "        model.add_module(name, layer)\n",
    "\n",
    "        if name in content_layers:\n",
    "            # add content loss:\n",
    "            target = model(content_img).detach()\n",
    "            content_loss = ContentLoss(target)\n",
    "            model.add_module(\"content_loss_{}\".format(i), content_loss)\n",
    "            content_losses.append(content_loss)\n",
    "\n",
    "        if name in style_layers:\n",
    "            # add style loss:\n",
    "            target_feature = model(style_img).detach()\n",
    "            style_loss = StyleLoss(target_feature)\n",
    "            model.add_module(\"style_loss_{}\".format(i), style_loss)\n",
    "            style_losses.append(style_loss)\n",
    "\n",
    "    # now we trim off the layers after the last content and style losses\n",
    "    for i in range(len(model) - 1, -1, -1):\n",
    "        if isinstance(model[i], ContentLoss) or isinstance(model[i], StyleLoss):\n",
    "            break\n",
    "\n",
    "    model = model[:(i + 1)]\n",
    "\n",
    "    return model, style_losses, content_losses"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Starting Image\n",
    "\n",
    "We can start with white noise or the input image. Here we will start with the input image as the target of our optimization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_img = content_img.clone()\n",
    "# if you want to use white noise instead uncomment the below line:\n",
    "# input_img = torch.randn(content_img.data.size(), device=device)\n",
    "\n",
    "# add the original input image to the figure:\n",
    "plt.figure()\n",
    "imshow(input_img, title='Input Image')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Neural Style Transfer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "create an optimizer for gradient descent:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_input_optimizer(input_img):\n",
    "    # this line to show that input is a parameter that requires a gradient\n",
    "    optimizer = optim.LBFGS([input_img.requires_grad_()])\n",
    "    return optimizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The body of the style transfer methods:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_style_transfer(cnn, normalization_mean, normalization_std,\n",
    "                       content_img, style_img, input_img, num_steps=300,\n",
    "                       style_weight=1000000, content_weight=1):\n",
    "    \"\"\"Run the style transfer.\"\"\"\n",
    "    print('Building the style transfer model..')\n",
    "    model, style_losses, content_losses = get_style_model_and_losses(cnn,\n",
    "        normalization_mean, normalization_std, style_img, content_img)\n",
    "    optimizer = get_input_optimizer(input_img)\n",
    "\n",
    "    print('Optimizing..')\n",
    "    run = [0]\n",
    "    while run[0] <= num_steps:\n",
    "\n",
    "        def closure():\n",
    "            # correct the values of updated input image\n",
    "            input_img.data.clamp_(0, 1)\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            model(input_img)\n",
    "            style_score = 0\n",
    "            content_score = 0\n",
    "\n",
    "            for sl in style_losses:\n",
    "                style_score += sl.loss\n",
    "            for cl in content_losses:\n",
    "                content_score += cl.loss\n",
    "\n",
    "            style_score *= style_weight\n",
    "            content_score *= content_weight\n",
    "\n",
    "            loss = style_score + content_score\n",
    "            loss.backward()\n",
    "\n",
    "            run[0] += 1\n",
    "            if run[0] % 50 == 0:\n",
    "                print(\"run {}:\".format(run))\n",
    "                print('Style Loss : {:4f} Content Loss: {:4f}'.format(\n",
    "                    style_score.item(), content_score.item()))\n",
    "                out_img = input_img.data.cpu().squeeze().permute(1,2,0).numpy();\n",
    "                cv2.imwrite(os.path.join('debug', str(run[0]) + '.png'), out_img*255);\n",
    "                print()\n",
    "\n",
    "            return style_score + content_score\n",
    "\n",
    "        optimizer.step(closure)\n",
    "\n",
    "    # a last correction...\n",
    "    input_img.data.clamp_(0, 1)\n",
    "\n",
    "    return input_img"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the algorithm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = run_style_transfer(cnn, cnn_normalization_mean, cnn_normalization_std,\n",
    "                            content_img, style_img, input_img)\n",
    "\n",
    "plt.figure()\n",
    "imshow(output, title='Output Image')\n",
    "\n",
    "# sphinx_gallery_thumbnail_number = 4\n",
    "plt.ioff()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extensions\n",
    "\n",
    "- Repeat with your own style and content image. \n",
    "  - What would be an interesting combination of style and content, and why?\n",
    "- Extend the training steps? how does this change the result?\n",
    "- Change the Content and Style weights above. How does this change the results?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "- A Neural Algorithm of Artistic Style by Gatys et. al - https://arxiv.org/abs/1508.06576\n",
    "- More on VGG16 from pytorch: https://pytorch.org/hub/pytorch_vision_vgg/\n",
    "- More detail on the pytorch demo page: https://pytorch.org/tutorials/advanced/neural_style_tutorial.html"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}