Image2Style_Implementation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Image2Style-Implementation.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyMslSjvw8dAgn8L2Y3lgvDi",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/zaidbhat1234/Image2StyleGAN/blob/main/Image2Style_Implementation.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7OfChTfTGuFN"
      },
      "source": [
        "##Mounting your google drive containing the code files.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Of4YWrPBZ1VV"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')\n",
        "%cd /content/gdrive/My\\ Drive/"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cGdy5d6RGw0a"
      },
      "source": [
        "##The code should be arranged in this order of directories"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Wp2xm7lKa_rX"
      },
      "source": [
        "%cd KAUST_Internship/StyleGAN_LatentEditor"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EOkD3m5DHLmd"
      },
      "source": [
        "##Resolving dependencies and importing libraries"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PdDq-9t5bMHs"
      },
      "source": [
        "from stylegan_layers import  G_mapping,G_synthesis\n",
        "import torch\n",
        "import torch.optim as optim\n",
        "import torch.nn as nn\n",
        "from collections import OrderedDict\n",
        "import torch.nn.functional as F\n",
        "from PIL import Image\n",
        "from torchvision import transforms\n",
        "import torchvision\n",
        "from torchvision import models\n",
        "from torchvision.utils import save_image\n",
        "import numpy as np\n",
        "from math import log10\n",
        "import matplotlib.pyplot as plt\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YFohF6yxHPfT"
      },
      "source": [
        "##Setup up the network on GPU"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Lpzc3ssEbMvj"
      },
      "source": [
        "device = 'cuda:0' if torch.cuda.is_available() else 'cpu'\n",
        "g_all = nn.Sequential(OrderedDict([('g_mapping', G_mapping()),\n",
        "    #('truncation', Truncation(avg_latent)),\n",
        "    ('g_synthesis', G_synthesis(resolution=1024))    \n",
        "    ]))\n",
        "\n",
        "#Load the pre-trained model\n",
        "g_all.load_state_dict(torch.load('weight_files/pytorch/karras2019stylegan-ffhq-1024x1024.pt', map_location=device))\n",
        "g_all.eval()\n",
        "g_all.to(device)\n",
        "g_mapping, g_synthesis = g_all[0],g_all[1]\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WFFdq4w5WKju"
      },
      "source": [
        "#%cd Images2/Hierar_optim #When testing Hierarchical Optimisation"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hNlNsmtNI-eU"
      },
      "source": [
        "##Read Images for inverting to latent space W+ and using in various experiments"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "64WmzCHXcPru"
      },
      "source": [
        "#Read a sample image we want to find a latent vector for\n",
        "# IMages : Portrait00002.png\n",
        "img_path = 'images/obama.png'\n",
        "with open(img_path,\"rb\") as f: \n",
        "  image=Image.open(f)\n",
        "  image=image.convert(\"RGB\")\n",
        "transform = transforms.Compose([transforms.ToTensor()])\n",
        "image = transform(image)\n",
        "image = image.unsqueeze(0)\n",
        "image = image.to(device)\n",
        "print(image.shape)\n",
        "\n",
        "path = img_path[3:-4]\n",
        "print(path)\n",
        "\n",
        "#Read a sample image we want to find a latent vector for\n",
        "img_path = 'images/ryan_01.png'\n",
        "with open(img_path,\"rb\") as f: \n",
        "  image1=Image.open(f)\n",
        "  image1=image1.convert(\"RGB\")\n",
        "transform = transforms.Compose([transforms.ToTensor()])\n",
        "image1 = transform(image1)\n",
        "image1 = image1.unsqueeze(0)\n",
        "image1 = image1.to(device)\n",
        "print(image1.shape)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Yb34Vd9_JGDH"
      },
      "source": [
        "##Generate random images using random W vectors"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yfYtdNLm4G4j"
      },
      "source": [
        "for i in range(20):\n",
        "  z = torch.randn(1,512,device = device)\n",
        "  img = g_all(z)\n",
        "  img = (img +1.0)/2.0\n",
        "  save_image(img.clamp(0,1),\"save_image/random_SG1-{}.png\".format(i+1))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FwnWHcnQJYSx"
      },
      "source": [
        "##VGG Perceptual loss network to give feature vectors from 4 parts of the pre-trained VGG-16 from 2,4,14,21"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KbpuWsfsCeA5"
      },
      "source": [
        "class VGG16_perceptual(torch.nn.Module):\n",
        "    def __init__(self, requires_grad=False):\n",
        "        super(VGG16_perceptual, self).__init__()\n",
        "        vgg_pretrained_features = models.vgg16(pretrained=True).features\n",
        "        self.slice1 = torch.nn.Sequential()\n",
        "        self.slice2 = torch.nn.Sequential()\n",
        "        self.slice3 = torch.nn.Sequential()\n",
        "        self.slice4 = torch.nn.Sequential()\n",
        "        for x in range(2):\n",
        "            self.slice1.add_module(str(x), vgg_pretrained_features[x])\n",
        "        for x in range(2, 4):\n",
        "            self.slice2.add_module(str(x), vgg_pretrained_features[x])\n",
        "        for x in range(4, 14):\n",
        "            self.slice3.add_module(str(x), vgg_pretrained_features[x])\n",
        "        for x in range(14, 21):\n",
        "            self.slice4.add_module(str(x), vgg_pretrained_features[x])\n",
        "        if not requires_grad:\n",
        "            for param in self.parameters():\n",
        "                param.requires_grad = False\n",
        "\n",
        "    def forward(self, X):\n",
        "        h = self.slice1(X)\n",
        "        h_relu1_1 = h\n",
        "        h = self.slice2(h)\n",
        "        h_relu1_2 = h\n",
        "        h = self.slice3(h)\n",
        "        h_relu3_2 = h\n",
        "        h = self.slice4(h)\n",
        "        h_relu4_2 = h\n",
        "        return h_relu1_1, h_relu1_2, h_relu3_2, h_relu4_2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NRMUGjIPJcur"
      },
      "source": [
        "##Loss function to calculate MSE and Perceptual losses"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HQKY9pcRChN0"
      },
      "source": [
        "def loss_function(syn_img, img, img_p, MSE_loss, upsample, perceptual):\n",
        "\n",
        "  #UpSample synthesized image to match the input size of VGG-16 input. \n",
        "  #Extract mid level features for real and synthesized image and find the MSE loss between them for perceptual loss. \n",
        "  #Find MSE loss between the real and synthesized images of actual size\n",
        "  syn_img_p = upsample(syn_img)\n",
        "  syn0, syn1, syn2, syn3 = perceptual(syn_img_p)\n",
        "  r0, r1, r2, r3 = perceptual(img_p)\n",
        "  mse = MSE_loss(syn_img,img)\n",
        "\n",
        "  per_loss = 0\n",
        "  per_loss += MSE_loss(syn0,r0)\n",
        "  per_loss += MSE_loss(syn1,r1)\n",
        "  per_loss += MSE_loss(syn2,r2)\n",
        "  per_loss += MSE_loss(syn3,r3)\n",
        "\n",
        "  return mse, per_loss"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I0qEEX6uJjIW"
      },
      "source": [
        "##Calculate PSNR"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z6HTTZUQfvBq"
      },
      "source": [
        "def PSNR(mse, flag = 0):\n",
        "  #flag = 0 if a single image is used and 1 if loss for a batch of images is to be calculated\n",
        "  if flag == 0:\n",
        "    psnr = 10 * log10(1 / mse.item())\n",
        "  return psnr"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4XK765XH4wsf"
      },
      "source": [
        "psnr_total = []"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xDxxua8kJobM"
      },
      "source": [
        "##Embedding Function to optimise the latent code W+ for GAN inversion."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_rqFu_EZFbum"
      },
      "source": [
        "def embedding_function(image):\n",
        "  upsample = torch.nn.Upsample(scale_factor = 256/1024, mode = 'bilinear')\n",
        "  img_p = image.clone()\n",
        "  img_p = upsample(img_p)\n",
        "  #Perceptual loss initialise object\n",
        "  perceptual = VGG16_perceptual().to(device)\n",
        "  \n",
        "  #MSE loss object\n",
        "  MSE_loss = nn.MSELoss(reduction=\"mean\")\n",
        "  #since the synthesis network expects 18 w vectors of size 1x512 thus we take latent vector of the same size\n",
        "  latents = torch.zeros((1,18,512), requires_grad = True, device = device)\n",
        "  #Optimizer to change latent code in each backward step\n",
        "  optimizer = optim.Adam({latents},lr=0.01,betas=(0.9,0.999),eps=1e-8)\n",
        "\n",
        "\n",
        "  #Loop to optimise latent vector to match the generated image to input image\n",
        "  loss_ = []\n",
        "  loss_psnr = []\n",
        "  for e in range(1500):\n",
        "    optimizer.zero_grad()\n",
        "    syn_img = g_synthesis(latents)\n",
        "    syn_img = (syn_img+1.0)/2.0\n",
        "    mse, per_loss = loss_function(syn_img, image, img_p, MSE_loss, upsample, perceptual)\n",
        "    psnr = PSNR(mse, flag = 0)\n",
        "    loss = per_loss +mse\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "    loss_np=loss.detach().cpu().numpy()\n",
        "    loss_p=per_loss.detach().cpu().numpy()\n",
        "    loss_m=mse.detach().cpu().numpy()\n",
        "    loss_psnr.append(psnr)\n",
        "    loss_.append(loss_np)\n",
        "    if (e+1)%500==0 :\n",
        "      print(\"iter{}: loss -- {},  mse_loss --{},  percep_loss --{}, psnr --{}\".format(e+1,loss_np,loss_m,loss_p,psnr))\n",
        "      save_image(syn_img.clamp(0,1),\"obama{}.png\".format(e+1))\n",
        "      #np.save(\"loss_list.npy\",loss_)\n",
        "      #np.save(\"latent_W.npy\".format(),latents.detach().cpu().numpy())\n",
        "\n",
        "  plt.plot(loss_, label = 'Loss = MSELoss + Perceptual')\n",
        "  plt.plot(loss_psnr, label = 'PSNR')\n",
        "  plt.legend()\n",
        "  return latents"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vJmDbawVF1jt"
      },
      "source": [
        "#Check generated images from 10 randomly generated faces \n",
        "\"\"\"Code in embed algo\n",
        "rnd = np.random.RandomState(7)\n",
        "z = torch.randn((1,512), requires_grad= True, device= device)\n",
        "img1 = g_all(z)\n",
        "img1 = (img1+1.0)/2.0\n",
        "latents = g_mapping(z)\n",
        "latents = torch.tensor(latents, requires_grad=True)\n",
        "save_image(img1.clamp(0,1),\"Random/random_image{}.png\".format(i))\"\"\"\n",
        "\n",
        "latent1 = embedding_function(image1)\n",
        "#latent2 = embedding_function(image1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qDiQaxuXJynP"
      },
      "source": [
        "##Image Morphing Experiment"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zBQaRIkSPNvH"
      },
      "source": [
        "for i in range(20):\n",
        "  a = (1/20)*i\n",
        "  w = latent1 * (1-a)+ latent2 * a\n",
        "  syn_img = g_synthesis(w)\n",
        "  syn_img = (syn_img+1.0)/2.0\n",
        "  save_image(syn_img.clamp(0,1),\"save_image/encode1/Morphed{}.png\".format(i))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TdHUYYGwKEbK"
      },
      "source": [
        "##Hierarchical Optimisation by optimising W vector first followed by optimising W+ in second pass"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5ij5gC03-lNO"
      },
      "source": [
        "def embedding_Hierarchical(image):\n",
        "  upsample = torch.nn.Upsample(scale_factor = 256/1024, mode = 'bilinear')\n",
        "  img_p = image.clone()\n",
        "  img_p = upsample(img_p)\n",
        "  \n",
        "  #Perceptual loss initialise object\n",
        "  perceptual = VGG16_perceptual().to(device)\n",
        "  \n",
        "  #MSE loss object\n",
        "  MSE_loss = nn.MSELoss(reduction=\"mean\")\n",
        "  #since the synthesis network expects 18 w vectors of size 1x512 thus we take latent vector of the same size\n",
        "  latent_w = torch.zeros((1,512), requires_grad = True, device = device)\n",
        "  \n",
        "  #Optimizer to change latent code in each backward step\n",
        "  optimizer = optim.Adam({latent_w},lr=0.01,betas=(0.9,0.999),eps=1e-8)\n",
        "\n",
        "\n",
        "  #Loop to optimise latent vector to match the generated image to input image\n",
        "  loss_ = []\n",
        "  loss_psnr = []\n",
        "  for e in range(1000):\n",
        "    optimizer.zero_grad()\n",
        "    latent_w1 = latent_w.unsqueeze(1).expand(-1, 18, -1)\n",
        "    syn_img = g_synthesis(latent_w1)\n",
        "    syn_img = (syn_img+1.0)/2.0\n",
        "    mse, per_loss = loss_function(syn_img, image, img_p, MSE_loss, upsample, perceptual)\n",
        "    psnr = PSNR(mse, flag = 0)\n",
        "    loss = per_loss +mse\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "    loss_np=loss.detach().cpu().numpy()\n",
        "    loss_p=per_loss.detach().cpu().numpy()\n",
        "    loss_m=mse.detach().cpu().numpy()\n",
        "    loss_psnr.append(psnr)\n",
        "    loss_.append(loss_np)\n",
        "    if (e+1)%500==0 :\n",
        "      print(\"iter{}: loss -- {},  mse_loss --{},  percep_loss --{}, psnr --{}\".format(e+1,loss_np,loss_m,loss_p,psnr))\n",
        "      save_image(syn_img.clamp(0,1),\"Hier_pass_morphP1-{}.png\".format(e+1))\n",
        "      #np.save(\"loss_list.npy\",loss_)\n",
        "      #np.save(\"latent_W.npy\".format(),latents.detach().cpu().numpy())\n",
        "\n",
        "  \n",
        "  latent_w1 = latent_w.unsqueeze(1).expand(-1, 18, -1)\n",
        "  latent_w1 = torch.tensor(latent_w1, requires_grad=True)\n",
        "  optimizer = optim.Adam({latent_w1},lr=0.01,betas=(0.9,0.999),eps=1e-8)\n",
        "  for e in range(1000):  \n",
        "    optimizer.zero_grad()\n",
        "    syn_img = g_synthesis(latent_w1)\n",
        "    syn_img = (syn_img+1.0)/2.0\n",
        "    mse, per_loss = loss_function(syn_img, image, img_p, MSE_loss, upsample, perceptual)\n",
        "    psnr = PSNR(mse, flag = 0)\n",
        "    loss = per_loss +mse\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "    loss_np=loss.detach().cpu().numpy()\n",
        "    loss_p=per_loss.detach().cpu().numpy()\n",
        "    loss_m=mse.detach().cpu().numpy()\n",
        "    loss_psnr.append(psnr)\n",
        "    loss_.append(loss_np)\n",
        "    if (e+1)%500==0 :\n",
        "      print(\"iter{}: loss -- {},  mse_loss --{},  percep_loss --{}, psnr --{}\".format(e+1,loss_np,loss_m,loss_p,psnr))\n",
        "      save_image(syn_img.clamp(0,1),\"Hier_pass_morphP2-{}.png\".format(e+1))\n",
        "\n",
        "\n",
        "  plt.plot(loss_, label = 'Loss = MSELoss + Perceptual')\n",
        "  plt.plot(loss_psnr, label = 'PSNR')\n",
        "  plt.legend()\n",
        "  return latent_w1"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "esrDWGdZKO34"
      },
      "source": [
        "##Obtaining latent codes for 2 images using Hierarchical Optimisation to check Image Morphing"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gkvuIrWSBeSD"
      },
      "source": [
        "latent1 = embedding_Hierarchical(image)\n",
        "latent2 = embedding_Hierarchical(image1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4H3xEUPqC4fm"
      },
      "source": [
        "#Image Morphing Experiment\n",
        "for i in range(20):\n",
        "  a = (1/20)*i\n",
        "  w = latent2 * (1-a)+ latent1 * a\n",
        "  syn_img = g_synthesis(w)\n",
        "  syn_img = (syn_img+1.0)/2.0\n",
        "  save_image(syn_img.clamp(0,1),\"Hier-Morphed{}.png\".format(i))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XXQCxHx-qxuJ"
      },
      "source": [
        "Downsamplimg to use a deeper network as it reduces computational cost as well as to increase receptive field. As downsampled image each pixel in it will have a wider receptive area it will affect in the actual higher dimensional image.\n",
        "\n",
        "Instead of pixel wise similarity we push for similarity in feature representations of real and synthesized images.\n",
        "\n",
        "Downsample to match the input size expectation of VGG - 256 or 224\n",
        "\n",
        "Taken Perceptual loss online official implementation\n",
        "\n",
        "We dont need the mapping part of the StyleGAN but only synthesis part which expects input as 1 x 18 x 512 i.e, 18 w's each input into the layers of Synthesis part."
      ]
    }
  ]
}