NeuralStyleTransfer.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "mount_file_id": "121TZP9tdEjwtEEXm1ey3VAh9LneORWri",
      "authorship_tag": "ABX9TyOEXw62VhIolJQK05CPwV1X"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CDNgjTzCV6SN"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "import torchvision\n",
        "from PIL import Image\n",
        "import torchvision.transforms as transforms\n",
        "import torchvision.models as models\n",
        "from torchvision.utils import save_image\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "model=models.vgg19(pretrained=True)"
      ],
      "metadata": {
        "id": "M-zu_H3XWolo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model=model.features"
      ],
      "metadata": {
        "id": "cROG4lZnW3iB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "0,5,10,19,28-> req. conv layers"
      ],
      "metadata": {
        "id": "QW9guYDOXOJL"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class VGG(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(VGG, self).__init__()\n",
        "        self.chosenlayers=[0,5,10,19,28]\n",
        "        self.model=models.vgg19(pretrained=True).features[:29]\n",
        "\n",
        "    def forward(self,x):\n",
        "      ftrs=[]\n",
        "\n",
        "      for num,layer in enumerate(self.model):\n",
        "        x=layer(x) #Applying conv layer on x\n",
        "\n",
        "        if num in self.chosenlayers:\n",
        "          ftrs.append(x)\n",
        "      return ftrs\n",
        "\n",
        "\n",
        "def load_img(img_name):\n",
        "  img=Image.open(img_name)\n",
        "  img=loader(img).unsqueeze(0)\n",
        "\n",
        "  return img.to(device)\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "p64hV64jXXMo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "device=torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "img_size=356\n",
        "\n",
        "loader=transforms.Compose([\n",
        "    transforms.Resize((img_size,img_size)),\n",
        "    transforms.ToTensor(),\n",
        "])\n"
      ],
      "metadata": {
        "id": "Dije_fGmZtSu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import cv2\n",
        "cap= cv2.VideoCapture('cont_vid.mp4')\n",
        "i=0\n",
        "print(1)\n",
        "while 1:\n",
        "  ret,frame=cap.read()\n",
        "\n",
        "  if ret:\n",
        "    cv2.imwrite(r\"./data/cvid\"+str(i)+\".jpg\",frame)\n",
        "    i+=1\n",
        "    print(i)\n",
        "  else:\n",
        "    print(\"done\")\n",
        "    break\n",
        "cap.release()"
      ],
      "metadata": {
        "id": "lJ9zR67oPbuh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "cimg=load_img(\"contn.jpg\")\n",
        "simg=load_img(\"style4.jpeg\")\n",
        "\n",
        "gen=cimg.clone().requires_grad_(True)"
      ],
      "metadata": {
        "id": "-w8sO42Haaks"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model= VGG().to(device).eval() #eval is to freeze weights"
      ],
      "metadata": {
        "id": "JCNJVTcwdbik"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "tot_steps=6000\n",
        "lr=0.001\n",
        "alpha=0.9\n",
        "beta=0.01\n",
        "optimizer=optim.Adam([gen],lr=lr)"
      ],
      "metadata": {
        "id": "3v2p59KHd0eb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Losses- content loss=mean((gen-cont)^2)\n",
        "        style loss="
      ],
      "metadata": {
        "id": "44ygO6HGfGY9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "for step in range(tot_steps):\n",
        "\n",
        "  optimizer.zero_grad()\n",
        "  gen_features=model(gen)\n",
        "  content_features=model(cimg)\n",
        "  style_features=model(simg)\n",
        "\n",
        "  s_loss=c_loss=0\n",
        "\n",
        "  for gf,cf,sf in zip(gen_features,content_features,style_features):\n",
        "    b_size,c,h,w=gf.shape\n",
        "    c_loss+=torch.mean((gf-cf)**2)\n",
        "\n",
        "    G= lambda x : x.view(c,h*w).mm(x.view(c,h*w).t())\n",
        "    gen_gram=G(gf)\n",
        "    style_gram=G(sf)\n",
        "\n",
        "    s_loss+=torch.mean((gen_gram- style_gram)**2)\n",
        "\n",
        "  tot_loss=alpha*c_loss+beta*s_loss\n",
        "  optimizer.zero_grad()\n",
        "  tot_loss.backward()\n",
        "  optimizer.step()\n",
        "\n",
        "  if step %200 ==0:\n",
        "    print(tot_loss)\n",
        "    save_image(gen,'gen.png')\n"
      ],
      "metadata": {
        "id": "WHWkuSQ0ee01"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#FOR VIDEO>:<:\n",
        "tot_steps=1000\n",
        "lr=0.001\n",
        "alpha=0.9\n",
        "beta=0.01\n",
        "\n",
        "for i in range(90):\n",
        "  cimg=load_img(r\"./data/cvid\"+str(i)+\".jpg\")\n",
        "  simg=load_img(\"style4.jpeg\")\n",
        "\n",
        "  gen=cimg.clone().requires_grad_(True)\n",
        "\n",
        "  model= VGG().to(device).eval() #eval is to freeze weights\n",
        "\n",
        "\n",
        "  optimizer=optim.Adam([gen],lr=lr)\n",
        "\n",
        "  for step in range(tot_steps):\n",
        "\n",
        "    optimizer.zero_grad()\n",
        "    gen_features=model(gen)\n",
        "    content_features=model(cimg)\n",
        "    style_features=model(simg)\n",
        "\n",
        "    s_loss=c_loss=0\n",
        "\n",
        "    for gf,cf,sf in zip(gen_features,content_features,style_features):\n",
        "      b_size,c,h,w=gf.shape\n",
        "      c_loss+=torch.mean((gf-cf)**2)\n",
        "\n",
        "      G= lambda x : x.view(c,h*w).mm(x.view(c,h*w).t())\n",
        "      gen_gram=G(gf)\n",
        "      style_gram=G(sf)\n",
        "\n",
        "      s_loss+=torch.mean((gen_gram- style_gram)**2)\n",
        "\n",
        "    tot_loss=alpha*c_loss+beta*s_loss\n",
        "    optimizer.zero_grad()\n",
        "    tot_loss.backward()\n",
        "    optimizer.step()\n",
        "\n",
        "    if step %200 ==0:\n",
        "      print(tot_loss)\n",
        "      save_image(gen,r'./out/gen'+str(i)+'.png')\n"
      ],
      "metadata": {
        "id": "1CgzWJrcRyHn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!zip -r out.zip out"
      ],
      "metadata": {
        "id": "wHmI8K-VbAU5"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}