{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Style Transfer with Deep Neural Networks\n\n\nIn this notebook, we’ll *recreate* a style transfer method that is outlined in the paper, [Image Style Transfer Using Convolutional Neural Networks, by Gatys](https://www.cv-foundation.org/openaccess/content_cvpr_2016/papers/Gatys_Image_Style_Transfer_CVPR_2016_paper.pdf) in PyTorch.\n\nIn this paper, style transfer uses the features found in the 19-layer VGG Network, which is comprised of a series of convolutional and pooling layers, and a few fully-connected layers. In the image below, the convolutional layers are named by stack and their order in the stack. Conv_1_1 is the first convolutional layer that an image is passed through, in the first stack. Conv_2_1 is the first convolutional layer in the *second* stack. The deepest convolutional layer in the network is conv_5_4.\n\n<img src='notebook_ims/vgg19_convlayers.png' width=80% />\n\n### Separating Style and Content\n\nStyle transfer relies on separating the content and style of an image. Given one content image and one style image, we aim to create a new, _target_ image which should contain our desired content and style components:\n* objects and their arrangement are similar to that of the **content image**\n* style, colors, and textures are similar to that of the **style image**\n\nAn example is shown below, where the content image is of a cat, and the style image is of [Hokusai's Great Wave](https://en.wikipedia.org/wiki/The_Great_Wave_off_Kanagawa). The generated target image still contains the cat but is stylized with the waves, blue and beige colors, and block print textures of the style image!\n\n<img src='notebook_ims/style_tx_cat.png' width=80% />\n\nIn this notebook, we'll use a pre-trained VGG19 Net to extract content or style features from a passed in image. We'll then formalize the idea of content and style _losses_ and use those to iteratively update our target image until we get a result that we want. You are encouraged to use a style and content image of your own and share your work on Twitter with @udacity; we'd love to see what you come up with!"},{"metadata":{"collapsed":true,"trusted":false},"cell_type":"code","source":"# import resources\n%matplotlib inline\n\nfrom PIL import Image\nfrom io import BytesIO\nimport matplotlib.pyplot as plt\nimport numpy as np\n\nimport torch\nimport torch.optim as optim\nimport requests\nfrom torchvision import transforms, models","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Load in VGG19 (features)\n\nVGG19 is split into two portions:\n* `vgg19.features`, which are all the convolutional and pooling layers\n* `vgg19.classifier`, which are the three linear, classifier layers at the end\n\nWe only need the `features` portion, which we're going to load in and \"freeze\" the weights of, below."},{"metadata":{"collapsed":true,"trusted":false},"cell_type":"code","source":"# get the \"features\" portion of VGG19 (we will not need the \"classifier\" portion)\nvgg = models.vgg19(pretrained=True).features\n\n# freeze all VGG parameters since we're only optimizing the target image\nfor param in vgg.parameters():\n    param.requires_grad_(False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"# move the model to GPU, if available\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nvgg.to(device)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Load in Content and Style Images\n\nYou can load in any images you want! Below, we've provided a helper function for loading in any type and size of image. The `load_image` function also converts images to normalized Tensors.\n\nAdditionally, it will be easier to have smaller images and to squish the content and style images so that they are of the same size."},{"metadata":{"collapsed":true,"trusted":false},"cell_type":"code","source":"def load_image(img_path, max_size=400, shape=None):\n    ''' Load in and transform an image, making sure the image\n       is <= 400 pixels in the x-y dims.'''\n    if \"http\" in img_path:\n        response = requests.get(img_path)\n        image = Image.open(BytesIO(response.content)).convert('RGB')\n    else:\n        image = Image.open(img_path).convert('RGB')\n    \n    # large images will slow down processing\n    if max(image.size) > max_size:\n        size = max_size\n    else:\n        size = max(image.size)\n    \n    if shape is not None:\n        size = shape\n        \n    in_transform = transforms.Compose([\n                        transforms.Resize(size),\n                        transforms.ToTensor(),\n                        transforms.Normalize((0.485, 0.456, 0.406), \n                                             (0.229, 0.224, 0.225))])\n\n    # discard the transparent, alpha channel (that's the :3) and add the batch dimension\n    image = in_transform(image)[:3,:,:].unsqueeze(0)\n    \n    return image","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Next, I'm loading in images by file name and forcing the style image to be the same size as the content image."},{"metadata":{"collapsed":true,"trusted":false},"cell_type":"code","source":"# load in content and style image\ncontent = load_image('images/octopus.jpg').to(device)\n# Resize style to match content, makes code easier\nstyle = load_image('images/hockney.jpg', shape=content.shape[-2:]).to(device)","execution_count":null,"outputs":[]},{"metadata":{"collapsed":true,"trusted":false},"cell_type":"code","source":"# helper function for un-normalizing an image \n# and converting it from a Tensor image to a NumPy image for display\ndef im_convert(tensor):\n    \"\"\" Display a tensor as an image. \"\"\"\n    \n    image = tensor.to(\"cpu\").clone().detach()\n    image = image.numpy().squeeze()\n    image = image.transpose(1,2,0)\n    image = image * np.array((0.229, 0.224, 0.225)) + np.array((0.485, 0.456, 0.406))\n    image = image.clip(0, 1)\n\n    return image","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"# display the images\nfig, (ax1, ax2) = plt.subplots(1, 2, figsize=(20, 10))\n# content and style ims side-by-side\nax1.imshow(im_convert(content))\nax2.imshow(im_convert(style))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"---\n## VGG19 Layers\n\nTo get the content and style representations of an image, we have to pass an image forward throug the VGG19 network until we get to the desired layer(s) and then get the output from that layer."},{"metadata":{"collapsed":true,"trusted":false},"cell_type":"code","source":"# print out VGG19 structure so you can see the names of various layers\n# print(vgg)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Content and Style Features\n\nBelow, complete the mapping of layer names to the names found in the paper for the _content representation_ and the _style representation_."},{"metadata":{"collapsed":true,"trusted":false},"cell_type":"code","source":"def get_features(image, model, layers=None):\n    \"\"\" Run an image forward through a model and get the features for \n        a set of layers. Default layers are for VGGNet matching Gatys et al (2016)\n    \"\"\"\n    \n    ## TODO: Complete mapping layer names of PyTorch's VGGNet to names from the paper\n    ## Need the layers for the content and style representations of an image\n    if layers is None:\n        layers = {'0': 'conv1_1',\n                  '5': 'conv2_1', \n                  '10': 'conv3_1', \n                  '19': 'conv4_1',\n                  '21': 'conv4_2',  ## content representation\n                  '28': 'conv5_1'}\n        \n    features = {}\n    x = image\n    # model._modules is a dictionary holding each module in the model\n    for name, layer in model._modules.items():\n        x = layer(x)\n        if name in layers:\n            features[layers[name]] = x\n            \n    return features","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"---\n## Gram Matrix \n\nThe output of every convolutional layer is a Tensor with dimensions associated with the `batch_size`, a depth, `d` and some height and width (`h`, `w`). The Gram matrix of a convolutional layer can be calculated as follows:\n* Get the depth, height, and width of a tensor using `batch_size, d, h, w = tensor.size`\n* Reshape that tensor so that the spatial dimensions are flattened\n* Calculate the gram matrix by multiplying the reshaped tensor by it's transpose \n\n*Note: You can multiply two matrices using `torch.mm(matrix1, matrix2)`.*"},{"metadata":{"collapsed":true,"trusted":false},"cell_type":"code","source":"def gram_matrix(tensor):\n    \"\"\" Calculate the Gram Matrix of a given tensor \n        Gram Matrix: https://en.wikipedia.org/wiki/Gramian_matrix\n    \"\"\"\n    \n    # get the batch_size, depth, height, and width of the Tensor\n    _, d, h, w = tensor.size()\n    \n    # reshape so we're multiplying the features for each channel\n    tensor = tensor.view(d, h * w)\n    \n    # calculate the gram matrix\n    gram = torch.mm(tensor, tensor.t())\n    \n    return gram ","execution_count":null,"outputs":[]},{"metadata":{"collapsed":true},"cell_type":"markdown","source":"## Putting it all Together\n\nNow that we've written functions for extracting features and computing the gram matrix of a given convolutional layer; let's put all these pieces together! We'll extract our features from our images and calculate the gram matrices for each layer in our style representation."},{"metadata":{"collapsed":true,"trusted":false},"cell_type":"code","source":"# get content and style features only once before training\ncontent_features = get_features(content, vgg)\nstyle_features = get_features(style, vgg)\n\n# calculate the gram matrices for each layer of our style representation\nstyle_grams = {layer: gram_matrix(style_features[layer]) for layer in style_features}\n\n# create a third \"target\" image and prep it for change\n# it is a good idea to start off with the target as a copy of our *content* image\n# then iteratively change its style\ntarget = content.clone().requires_grad_(True).to(device)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"---\n## Loss and Weights\n\n#### Individual Layer Style Weights\n\nBelow, you are given the option to weight the style representation at each relevant layer. It's suggested that you use a range between 0-1 to weight these layers. By weighting earlier layers (`conv1_1` and `conv2_1`) more, you can expect to get _larger_ style artifacts in your resulting, target image. Should you choose to weight later layers, you'll get more emphasis on smaller features. This is because each layer is a different size and together they create a multi-scale style representation!\n\n#### Content and Style Weight\n\nJust like in the paper, we define an alpha (`content_weight`) and a beta (`style_weight`). This ratio will affect how _stylized_ your final image is. It's recommended that you leave the content_weight = 1 and set the style_weight to achieve the ratio you want."},{"metadata":{"collapsed":true,"trusted":false},"cell_type":"code","source":"# weights for each style layer \n# weighting earlier layers more will result in *larger* style artifacts\n# notice we are excluding `conv4_2` our content representation\nstyle_weights = {'conv1_1': 1.,\n                 'conv2_1': 0.75,\n                 'conv3_1': 0.2,\n                 'conv4_1': 0.2,\n                 'conv5_1': 0.2}\n\ncontent_weight = 1  # alpha\nstyle_weight = 1e6  # beta","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Updating the Target & Calculating Losses\n\nYou'll decide on a number of steps for which to update your image, this is similar to the training loop that you've seen before, only we are changing our _target_ image and nothing else about VGG19 or any other image. Therefore, the number of steps is really up to you to set! **I recommend using at least 2000 steps for good results.** But, you may want to start out with fewer steps if you are just testing out different weight values or experimenting with different images.\n\nInside the iteration loop, you'll calculate the content and style losses and update your target image, accordingly.\n\n#### Content Loss\n\nThe content loss will be the mean squared difference between the target and content features at layer `conv4_2`. This can be calculated as follows: \n```\ncontent_loss = torch.mean((target_features['conv4_2'] - content_features['conv4_2'])**2)\n```\n\n#### Style Loss\n\nThe style loss is calculated in a similar way, only you have to iterate through a number of layers, specified by name in our dictionary `style_weights`. \n> You'll calculate the gram matrix for the target image, `target_gram` and style image `style_gram` at each of these layers and compare those gram matrices, calculating the `layer_style_loss`. \n> Later, you'll see that this value is normalized by the size of the layer.\n\n#### Total Loss\n\nFinally, you'll create the total loss by adding up the style and content losses and weighting them with your specified alpha and beta!\n\nIntermittently, we'll print out this loss; don't be alarmed if the loss is very large. It takes some time for an image's style to change and you should focus on the appearance of your target image rather than any loss value. Still, you should see that this loss decreases over some number of iterations."},{"metadata":{"scrolled":true,"trusted":false},"cell_type":"code","source":"# for displaying the target image, intermittently\nshow_every = 400\n\n# iteration hyperparameters\noptimizer = optim.Adam([target], lr=0.003)\nsteps = 2000  # decide how many iterations to update your image (5000)\n\nfor ii in range(1, steps+1):\n    \n    # get the features from your target image\n    target_features = get_features(target, vgg)\n    \n    # the content loss\n    content_loss = torch.mean((target_features['conv4_2'] - content_features['conv4_2'])**2)\n    \n    # the style loss\n    # initialize the style loss to 0\n    style_loss = 0\n    # then add to it for each layer's gram matrix loss\n    for layer in style_weights:\n        # get the \"target\" style representation for the layer\n        target_feature = target_features[layer]\n        target_gram = gram_matrix(target_feature)\n        _, d, h, w = target_feature.shape\n        # get the \"style\" style representation\n        style_gram = style_grams[layer]\n        # the style loss for one layer, weighted appropriately\n        layer_style_loss = style_weights[layer] * torch.mean((target_gram - style_gram)**2)\n        # add to the style loss\n        style_loss += layer_style_loss / (d * h * w)\n        \n    # calculate the *total* loss\n    total_loss = content_weight * content_loss + style_weight * style_loss\n    \n    # update your target image\n    optimizer.zero_grad()\n    total_loss.backward()\n    optimizer.step()\n    \n    # display intermediate images and print the loss\n    if  ii % show_every == 0:\n        print('Total loss: ', total_loss.item())\n        plt.imshow(im_convert(target))\n        plt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Display the Target Image"},{"metadata":{"trusted":false},"cell_type":"code","source":"# display content and final, target image\nfig, (ax1, ax2) = plt.subplots(1, 2, figsize=(20, 10))\nax1.imshow(im_convert(content))\nax2.imshow(im_convert(target))","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python [default]","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.6.4"}},"nbformat":4,"nbformat_minor":1}