Fast-Style-Transfer-PyTorch.ipynb

{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"name":"Fast-Style-Transfer-PyTorch.ipynb","version":"0.3.2","provenance":[],"collapsed_sections":[]},"kernelspec":{"name":"python3","display_name":"Python 3"},"accelerator":"GPU"},"cells":[{"metadata":{"id":"GGwfDV9ytcS-","colab_type":"code","colab":{}},"cell_type":"code","source":["from __future__ import print_function\n","\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","import torch.optim as optim\n","import os\n","from PIL import Image\n","from IPython.display import display \n","import matplotlib.pyplot as plt\n","from torchvision import datasets\n","import torchvision.transforms as transforms\n","import torchvision.models as models\n","import torchsummary\n","import numpy as np\n","from collections import namedtuple\n","\n","import copy\n","import time\n","import glob\n","import re\n","import cv2"],"execution_count":0,"outputs":[]},{"metadata":{"id":"Wc-MT-aTyh2d","colab_type":"code","colab":{}},"cell_type":"code","source":["device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n","print(device)"],"execution_count":0,"outputs":[]},{"metadata":{"id":"TAcNiZX0yno0","colab_type":"code","colab":{}},"cell_type":"code","source":["from google.colab import drive\n","drive.mount(\"/content/gdrive\")"],"execution_count":0,"outputs":[]},{"metadata":{"id":"H1VUep-868As","colab_type":"text"},"cell_type":"markdown","source":["Download COCO Dataset - http://images.cocodataset.org/zips/val2017.zip\n","\n","some differences:\n","\n","Original paper: COCO 2014 train --> 80000 images & 13GB\n","\n","My implementation: COCO 2017 val --> 5000 images / 1GB --> x16 epochs"]},{"metadata":{"id":"T8Q10broPFty","colab_type":"code","colab":{}},"cell_type":"code","source":["style_image_location = \"/content/gdrive/My Drive/Colab_Notebooks/data/vikendi.jpg\" #FIXME\n","\n","style_image_sample = Image.open(style_image_location, 'r')\n","display(style_image_sample)"],"execution_count":0,"outputs":[]},{"metadata":{"id":"zWFReczp6ckL","colab_type":"code","colab":{}},"cell_type":"code","source":["batch_size = 8\n","random_seed = 10\n","num_epochs = 64 \n","initial_lr = 1e-3\n","checkpoint_dir = \"/content/gdrive/My Drive/Colab_Notebooks/data/\" #FIXME\n","\n","content_weight = 1e5\n","style_weight = 1e10\n","log_interval = 50\n","checkpoint_interval = 500\n","\n","#running_option = \"test\" #FIXME \n","running_option = \"test_video\" #FIXME\n","#running_option = \"training\" #FIXME"],"execution_count":0,"outputs":[]},{"metadata":{"id":"MHjz8SjkPyza","colab_type":"code","colab":{}},"cell_type":"code","source":["class VGG16(torch.nn.Module):\n","    def __init__(self, requires_grad=False):\n","        super(VGG16, 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(4):\n","            self.slice1.add_module(str(x), vgg_pretrained_features[x])\n","        for x in range(4, 9):\n","            self.slice2.add_module(str(x), vgg_pretrained_features[x])\n","        for x in range(9, 16):\n","            self.slice3.add_module(str(x), vgg_pretrained_features[x])\n","        for x in range(16, 23):\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_2 = h\n","        h = self.slice2(h)\n","        h_relu2_2 = h\n","        h = self.slice3(h)\n","        h_relu3_3 = h\n","        h = self.slice4(h)\n","        h_relu4_3 = h\n","        vgg_outputs = namedtuple(\"VggOutputs\", ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3'])\n","        out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3)\n","        return out"],"execution_count":0,"outputs":[]},{"metadata":{"id":"ZAfafcBIsvzg","colab_type":"code","colab":{}},"cell_type":"code","source":["class ConvLayer(nn.Module):\n","    def __init__(self, in_channels, out_channels, kernel_size, stride):\n","        super(ConvLayer, self).__init__()\n","        reflection_padding = kernel_size // 2\n","        self.reflection_pad = nn.ReflectionPad2d(reflection_padding)\n","        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride)\n","\n","    def forward(self, x):\n","        out = self.reflection_pad(x)\n","        out = self.conv2d(out)\n","        return out"],"execution_count":0,"outputs":[]},{"metadata":{"id":"iLkW4JFougnf","colab_type":"code","colab":{}},"cell_type":"code","source":["class ResidualBlock(nn.Module):\n","    \"\"\"ResidualBlock\n","    introduced in: https://arxiv.org/abs/1512.03385\n","    recommended architecture: http://torch.ch/blog/2016/02/04/resnets.html\n","    \"\"\"\n","\n","    def __init__(self, channels):\n","        super(ResidualBlock, self).__init__()\n","        self.conv1 = ConvLayer(channels, channels, kernel_size=3, stride=1)\n","        self.in1 = nn.InstanceNorm2d(channels, affine=True)\n","        self.conv2 = ConvLayer(channels, channels, kernel_size=3, stride=1)\n","        self.in2 = nn.InstanceNorm2d(channels, affine=True)\n","        self.relu = nn.ReLU()\n","\n","    def forward(self, x):\n","        residual = x\n","        out = self.relu(self.in1(self.conv1(x)))\n","        out = self.in2(self.conv2(out))\n","        out = out + residual\n","        return out"],"execution_count":0,"outputs":[]},{"metadata":{"id":"-CryjuYwwJKg","colab_type":"code","colab":{}},"cell_type":"code","source":["class UpsampleConvLayer(nn.Module):\n","    \"\"\"UpsampleConvLayer\n","    Upsamples the input and then does a convolution. This method gives better results\n","    compared to ConvTranspose2d.\n","    ref: http://distill.pub/2016/deconv-checkerboard/\n","    \"\"\"\n","\n","    def __init__(self, in_channels, out_channels, kernel_size, stride, upsample=None):\n","        super(UpsampleConvLayer, self).__init__()\n","        self.upsample = upsample\n","        reflection_padding = kernel_size // 2\n","        self.reflection_pad = nn.ReflectionPad2d(reflection_padding)\n","        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride)\n","\n","    def forward(self, x):\n","        x_in = x\n","        if self.upsample:\n","            x_in = nn.functional.interpolate(x_in, mode='nearest', scale_factor=self.upsample)\n","        out = self.reflection_pad(x_in)\n","        out = self.conv2d(out)\n","        return out"],"execution_count":0,"outputs":[]},{"metadata":{"id":"R8dYDDEzxX9l","colab_type":"code","colab":{}},"cell_type":"code","source":["class TransformerNet(nn.Module):\n","    def __init__(self):\n","        super(TransformerNet, self).__init__()\n","        # Initial convolution layers\n","        self.encoder = nn.Sequential()\n","        \n","        self.encoder.add_module('conv1', ConvLayer(3, 32, kernel_size=9, stride=1))\n","        self.encoder.add_module('in1', nn.InstanceNorm2d(32, affine=True))\n","        self.encoder.add_module('relu1', nn.ReLU())\n","        \n","        self.encoder.add_module('conv2', ConvLayer(32, 64, kernel_size=3, stride=2))\n","        self.encoder.add_module('in2', nn.InstanceNorm2d(64, affine=True))\n","        self.encoder.add_module('relu2', nn.ReLU())\n","        \n","        self.encoder.add_module('conv3', ConvLayer(64, 128, kernel_size=3, stride=2))\n","        self.encoder.add_module('in3', nn.InstanceNorm2d(128, affine=True))\n","        self.encoder.add_module('relu3', nn.ReLU())\n","\n","        # Residual layers\n","        self.residual = nn.Sequential()\n","        \n","        for i in range(5):\n","          self.residual.add_module('resblock_%d' %(i+1), ResidualBlock(128))\n","        \n","        # Upsampling Layers\n","        self.decoder = nn.Sequential()\n","        self.decoder.add_module('deconv1', UpsampleConvLayer(128, 64, kernel_size=3, stride=1, upsample=2))\n","        self.decoder.add_module('in4', nn.InstanceNorm2d(64, affine=True))\n","        self.encoder.add_module('relu4', nn.ReLU())\n","\n","        self.decoder.add_module('deconv2', UpsampleConvLayer(64, 32, kernel_size=3, stride=1, upsample=2))\n","        self.decoder.add_module('in5', nn.InstanceNorm2d(32, affine=True))\n","        self.encoder.add_module('relu5', nn.ReLU())\n","\n","        self.decoder.add_module('deconv3', ConvLayer(32, 3, kernel_size=9, stride=1))\n","\n","\n","    def forward(self, x):\n","        encoder_output = self.encoder(x)\n","        residual_output = self.residual(encoder_output)\n","        decoder_output = self.decoder(residual_output)\n","        \n","        return decoder_output"],"execution_count":0,"outputs":[]},{"metadata":{"id":"uCWbtGj6DrdP","colab_type":"code","colab":{}},"cell_type":"code","source":["# Video Capture Sanity Check\n","cap = cv2.VideoCapture(\"/content/gdrive/My Drive/Colab_Notebooks/data/mirama_demo.mp4\")\n","\n","while(cap.isOpened()):\n","    ret, frame = cap.read()\n","    frame = frame[:,:,::-1]\n","    \n","    print(frame.shape)\n","    \n","    plt.imshow(frame)\n","    \n","    break\n","cap.release()\n","cv2.destroyAllWindows()"],"execution_count":0,"outputs":[]},{"metadata":{"id":"6SBcm-tcwPYn","colab_type":"code","colab":{}},"cell_type":"code","source":["\"\"\" Util Functions \"\"\"\n","def load_image(filename, size=None, scale=None):\n","    img = Image.open(filename)\n","    if size is not None:\n","        img = img.resize((size, size), Image.ANTIALIAS)\n","    elif scale is not None:\n","        img = img.resize((int(img.size[0] / scale), int(img.size[1] / scale)), Image.ANTIALIAS)\n","    return img\n","    \n","\n","def save_image(filename, data):\n","    img = data.clone().clamp(0, 255).numpy()\n","    img = img.transpose(1, 2, 0).astype(\"uint8\")\n","    img = Image.fromarray(img)\n","    display(img)\n","    img.save(filename)\n","\n","def post_process_image(data):\n","    img = data.clone().clamp(0, 255).numpy()\n","    img = img.transpose(1, 2, 0).astype(\"uint8\")\n","    #img = Image.fromarray(img)\n","    \n","    return img\n","    \n","\n","def gram_matrix(y):\n","    (b, ch, h, w) = y.size()\n","    features = y.view(b, ch, w * h)\n","    features_t = features.transpose(1, 2)\n","    gram = features.bmm(features_t) / (ch * h * w)\n","    return gram\n","\n","\n","def normalize_batch(batch):\n","    # normalize using imagenet mean and std\n","    mean = batch.new_tensor([0.485, 0.456, 0.406]).view(-1, 1, 1)\n","    std = batch.new_tensor([0.229, 0.224, 0.225]).view(-1, 1, 1)\n","    batch = batch.div_(255.0)\n","    return (batch - mean) / std"],"execution_count":0,"outputs":[]},{"metadata":{"id":"2e6wKOLrxNcn","colab_type":"code","colab":{}},"cell_type":"code","source":["np.random.seed(random_seed)\n","torch.manual_seed(random_seed)\n","\n","transform = transforms.Compose([\n","    transforms.Resize(256),\n","    transforms.CenterCrop(256),\n","    transforms.ToTensor(),\n","    transforms.Lambda(lambda x: x.mul(255))\n","])\n","\n","style_transform = transforms.Compose([\n","        transforms.ToTensor(),\n","        transforms.Lambda(lambda x: x.mul(255))\n","])\n","\n","print(glob.glob(\"/content/gdrive/My Drive/Colab_Notebooks/data/COCO/val2017/*\"))\n","\n","train_dataset = datasets.ImageFolder(\"/content/gdrive/My Drive/Colab_Notebooks/data/COCO\", transform) #FIXME\n","train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size)"],"execution_count":0,"outputs":[]},{"metadata":{"id":"LJt1Q7t8KdqN","colab_type":"code","colab":{}},"cell_type":"code","source":["transformer = TransformerNet()\n","vgg = VGG16(requires_grad=False).to(device)"],"execution_count":0,"outputs":[]},{"metadata":{"id":"mK8InidkOFKd","colab_type":"code","colab":{}},"cell_type":"code","source":["optimizer = torch.optim.Adam(transformer.parameters(), initial_lr)\n","mse_loss = nn.MSELoss()\n","\n","style = load_image(filename=style_image_location, size=None, scale=None)\n","style = style_transform(style)\n","style = style.repeat(batch_size, 1, 1, 1).to(device)\n","\n","features_style = vgg(normalize_batch(style))\n","gram_style = [gram_matrix(y) for y in features_style]"],"execution_count":0,"outputs":[]},{"metadata":{"id":"NptYn-DTBDiQ","colab_type":"code","colab":{}},"cell_type":"code","source":["transfer_learning = False # inference or training first --> False / Transfer learning --> True\n","ckpt_model_path = os.path.join(checkpoint_dir, \"ckpt_epoch_63_batch_id_500.pth\") #FIXME\n","\n","if transfer_learning:\n","  checkpoint = torch.load(ckpt_model_path, map_location=device)\n","  transformer.load_state_dict(checkpoint['model_state_dict'])\n","  transformer.to(device)"],"execution_count":0,"outputs":[]},{"metadata":{"id":"8G5CJq0_QxB1","colab_type":"code","colab":{}},"cell_type":"code","source":["if running_option == \"training\":\n","  if transfer_learning:\n","      transfer_learning_epoch = checkpoint['epoch'] \n","  else:\n","      transfer_learning_epoch = 0\n","\n","  for epoch in range(transfer_learning_epoch, num_epochs):\n","        transformer.train()\n","        agg_content_loss = 0.\n","        agg_style_loss = 0.\n","        count = 0\n","\n","        for batch_id, (x, _) in enumerate(train_loader):\n","            n_batch = len(x)\n","            count += n_batch\n","            optimizer.zero_grad()\n","\n","            x = x.to(device)\n","            y = transformer(x)\n","\n","            y = normalize_batch(y)\n","            x = normalize_batch(x)\n","\n","            features_y = vgg(y)\n","            features_x = vgg(x)\n","\n","            content_loss = content_weight * mse_loss(features_y.relu2_2, features_x.relu2_2)\n","\n","            style_loss = 0.\n","            for ft_y, gm_s in zip(features_y, gram_style):\n","                gm_y = gram_matrix(ft_y)\n","                style_loss += mse_loss(gm_y, gm_s[:n_batch, :, :])\n","            style_loss *= style_weight\n","\n","            total_loss = content_loss + style_loss\n","            total_loss.backward()\n","            optimizer.step()\n","\n","            agg_content_loss += content_loss.item()\n","            agg_style_loss += style_loss.item()\n","\n","            if (batch_id + 1) % log_interval == 0:\n","                mesg = \"{}\\tEpoch {}:\\t[{}/{}]\\tcontent: {:.6f}\\tstyle: {:.6f}\\ttotal: {:.6f}\".format(\n","                    time.ctime(), epoch + 1, count, len(train_dataset),\n","                                  agg_content_loss / (batch_id + 1),\n","                                  agg_style_loss / (batch_id + 1),\n","                                  (agg_content_loss + agg_style_loss) / (batch_id + 1)\n","                )\n","                print(mesg)\n","\n","            if checkpoint_dir is not None and (batch_id + 1) % checkpoint_interval == 0:\n","                transformer.eval().cpu()\n","                ckpt_model_filename = \"ckpt_epoch_\" + str(epoch) + \"_batch_id_\" + str(batch_id + 1) + \".pth\"\n","                print(str(epoch), \"th checkpoint is saved!\")\n","                ckpt_model_path = os.path.join(checkpoint_dir, ckpt_model_filename)\n","                torch.save({\n","                'epoch': epoch,\n","                'model_state_dict': transformer.state_dict(),\n","                'optimizer_state_dict': optimizer.state_dict(),\n","                'loss': total_loss\n","                }, ckpt_model_path)\n","\n","                transformer.to(device).train()  "],"execution_count":0,"outputs":[]},{"metadata":{"id":"VdGrXA2hSXD9","colab_type":"code","colab":{}},"cell_type":"code","source":["if running_option == \"test\":\n","    content_image = load_image(\"/content/gdrive/My Drive/Colab_Notebooks/data/sanok_test.png\", scale=2) #FIXME\n","    content_transform = transforms.Compose([\n","        transforms.ToTensor(),\n","        transforms.Lambda(lambda x: x.mul(255))\n","    ])\n","    content_image = content_transform(content_image)\n","    content_image = content_image.unsqueeze(0).to(device)\n","\n","    with torch.no_grad():\n","        style_model = TransformerNet()\n","        \n","        ckpt_model_path = os.path.join(checkpoint_dir, \"ckpt_epoch_63_batch_id_500.pth\") #FIXME\n","        checkpoint = torch.load(ckpt_model_path, map_location=device)\n","        \n","        # remove saved deprecated running_* keys in InstanceNorm from the checkpoint\n","        for k in list(checkpoint.keys()):\n","            if re.search(r'in\\d+\\.running_(mean|var)$', k):\n","                del checkpoint[k]\n","        \n","        style_model.load_state_dict(checkpoint['model_state_dict'])\n","        style_model.to(device)\n","        \n","        output = style_model(content_image).cpu()\n","        \n","    save_image(\"/content/gdrive/My Drive/Colab_Notebooks/data/sanok_result.png\", output[0]) #FIXME"],"execution_count":0,"outputs":[]},{"metadata":{"id":"YfZDvDfdBunq","colab_type":"code","colab":{}},"cell_type":"code","source":["if running_option == \"test_video\":\n","    \n","    with torch.no_grad():\n","        style_model = TransformerNet()\n","\n","        ckpt_model_path = os.path.join(checkpoint_dir, \"ckpt_epoch_63_batch_id_500.pth\") #FIXME\n","        checkpoint = torch.load(ckpt_model_path, map_location=device)\n","\n","        # remove saved deprecated running_* keys in InstanceNorm from the checkpoint\n","        for k in list(checkpoint.keys()):\n","            if re.search(r'in\\d+\\.running_(mean|var)$', k):\n","                del checkpoint[k]\n","\n","        style_model.load_state_dict(checkpoint['model_state_dict'])\n","        style_model.to(device)\n","\n","        cap = cv2.VideoCapture(\"/content/gdrive/My Drive/Colab_Notebooks/data/mirama_demo.mp4\") #FIXME\n","\n","        frame_cnt = 0\n","        \n","        fourcc = cv2.VideoWriter_fourcc(*'XVID') \n","        out = cv2.VideoWriter('/content/gdrive/My Drive/Colab_Notebooks/data/mirama_demo_result.avi', fourcc, 60.0, (1920,1080)) #FIXME\n","\n","        \n","        while(cap.isOpened()):\n","            ret, frame = cap.read()\n","            \n","            try:\n","              frame = frame[:,:,::-1] - np.zeros_like(frame)\n","            except:\n","              break\n","              \n","            print(frame_cnt, \"th frame is loaded!\")\n","\n","            content_image = frame\n","            content_transform = transforms.Compose([\n","                transforms.ToTensor(),\n","                transforms.Lambda(lambda x: x.mul(255))\n","            ])\n","            content_image = content_transform(content_image)\n","            content_image = content_image.unsqueeze(0).to(device)\n","\n","            output = style_model(content_image).cpu()\n","            #save_image(\"/content/gdrive/My Drive/Colab_Notebooks/data/vikendi_video_result/\" + str(frame_cnt) +\".png\", output[0]) #FIXME\n","            out.write(post_process_image(output[0]))\n","            frame_cnt += 1\n","            \n","\n","        cap.release()\n","        out.release()\n","        cv2.destroyAllWindows()\n","  \n"],"execution_count":0,"outputs":[]},{"metadata":{"id":"-BgPhiVmjoNY","colab_type":"code","colab":{}},"cell_type":"code","source":[""],"execution_count":0,"outputs":[]}]}