notebook.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Image impainting",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python_defaultSpec_1606826343952",
      "display_name": "Python 3.7.7 64-bit"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "AqLspm657HvV"
      },
      "source": [
        "import torch\n",
        "from torch import nn, optim\n",
        "from torchsummary import summary\n",
        "from torchvision.utils import make_grid\n",
        "\n",
        "import os\n",
        "from shutil import rmtree\n",
        "import imageio\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "import matplotlib.pyplot as plt"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xnt86T807H1A"
      },
      "source": [
        "def pil_to_np_array(pil_image):\n",
        "    array = np.array(pil_image).transpose(2,0,1)\n",
        "    return array.astype(np.float32) / 255.\n",
        "\n",
        "def np_to_torch_array(np_array):\n",
        "    return torch.from_numpy(np_array)[None, :]\n",
        "\n",
        "def torch_to_np_array(torch_array):\n",
        "    return torch_array.detach().cpu().numpy()[0]\n",
        "\n",
        "def read_image(path, image_size = -1):\n",
        "    pil_image = Image.open(path)\n",
        "    return pil_image\n",
        "\n",
        "def save_image(np_array, step):\n",
        "    pil_image = Image.fromarray((np_array * 255.0).transpose(1, 2, 0).astype('uint8'), 'RGB')\n",
        "    pil_image.save(f'progress/{str(step).zfill(len(str(TRAINING_STEPS)))}.png')\n",
        "\n",
        "def get_image_grid(images, nrow = 3):\n",
        "    torch_images = [torch.from_numpy(x) for x in images]\n",
        "    grid = make_grid(torch_images, nrow)\n",
        "    return grid.numpy()\n",
        "    \n",
        "def visualize_sample(*images_np, nrow = 3, size_factor = 10):\n",
        "    c = max(x.shape[0] for x in images_np)\n",
        "    images_np = [x if (x.shape[0] == c) else np.concatenate([x, x, x], axis=0) for x in images_np]\n",
        "    grid = get_image_grid(images_np, nrow)\n",
        "    plt.figure(figsize = (len(images_np) + size_factor, 12 + size_factor))\n",
        "    plt.axis('off')\n",
        "    plt.imshow(grid.transpose(1, 2, 0))\n",
        "    plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ScS48lQN7H_-"
      },
      "source": [
        "DTYPE = torch.cuda.FloatTensor\n",
        "INPUT_DEPTH = 32\n",
        "LR = 0.01 \n",
        "TRAINING_STEPS = 10000\n",
        "SHOW_STEP = 50\n",
        "REG_NOISE = 0.03\n",
        "MAX_DIM = 128"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 341
        },
        "id": "mM4-mH_l7H45",
        "outputId": "70f8881d-6244-40ac-dd5d-975891b282c6"
      },
      "source": [
        "original_image_path  = os.path.join('data/watermark-available/image1.png')\n",
        "watermark_path = os.path.join('data/watermark-available/watermark.png')\n",
        "\n",
        "original_image_pil = read_image(original_image_path)\n",
        "original_image_pil = original_image_pil.convert('RGB')\n",
        "original_image_pil = original_image_pil.resize((128, 128))\n",
        "\n",
        "w, h = original_image_pil.size\n",
        "aspect_ratio = w / h\n",
        "if w > MAX_DIM and w > h:\n",
        "    h = int((h / w) * MAX_DIM)\n",
        "    w = MAX_DIM\n",
        "elif h > MAX_DIM and h > w:\n",
        "    w = int((w / h) * MAX_DIM)\n",
        "    h = MAX_DIM\n",
        "\n",
        "original_image_pil = original_image_pil.resize((w, h))\n",
        "\n",
        "watermark_pil = read_image(watermark_path)\n",
        "watermark_pil = watermark_pil.convert('RGB')\n",
        "watermark_pil = watermark_pil.resize((original_image_pil.size[0], original_image_pil.size[1]))\n",
        "\n",
        "original_image_np = pil_to_np_array(original_image_pil)\n",
        "watermark_np = pil_to_np_array(watermark_pil)\n",
        "# watermark_np[watermark_np == 0.0] = 1.0\n",
        "\n",
        "watermarked_np = original_image_np * watermark_np\n",
        "image_mask_var = np_to_torch_array(watermark_np).type(DTYPE)\n",
        "\n",
        "visualize_sample(original_image_np, watermark_np, watermarked_np, nrow = 3, size_factor = 10)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zSEdxRbN9Vyr"
      },
      "source": [
        "class Conv2dBlock(nn.Module):\n",
        "    def __init__(self, in_channels, out_channels, kernel_size, stride = 1, bias = False):\n",
        "        super(Conv2dBlock, self).__init__()\n",
        "\n",
        "        self.model = nn.Sequential(\n",
        "            nn.ReflectionPad2d(int((kernel_size - 1) / 2)),\n",
        "            nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding = 0, bias = bias),\n",
        "            nn.BatchNorm2d(out_channels),\n",
        "            nn.LeakyReLU(0.2)\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.model(x)\n",
        "\n",
        "class Concat(nn.Module):\n",
        "    def __init__(self, dim, *args):\n",
        "        super(Concat, self).__init__()\n",
        "        self.dim = dim\n",
        "\n",
        "        for idx, module in enumerate(args):\n",
        "            self.add_module(str(idx), module)\n",
        "\n",
        "    def forward(self, input):\n",
        "        inputs = []\n",
        "        for module in self._modules.values():\n",
        "            inputs.append(module(input))\n",
        "\n",
        "        inputs_shapes2 = [x.shape[2] for x in inputs]\n",
        "        inputs_shapes3 = [x.shape[3] for x in inputs]        \n",
        "\n",
        "        if np.all(np.array(inputs_shapes2) == min(inputs_shapes2)) and np.all(np.array(inputs_shapes3) == min(inputs_shapes3)):\n",
        "            inputs_ = inputs\n",
        "        else:\n",
        "            target_shape2 = min(inputs_shapes2)\n",
        "            target_shape3 = min(inputs_shapes3)\n",
        "\n",
        "            inputs_ = []\n",
        "            for inp in inputs: \n",
        "                diff2 = (inp.size(2) - target_shape2) // 2 \n",
        "                diff3 = (inp.size(3) - target_shape3) // 2 \n",
        "                inputs_.append(inp[:, :, diff2: diff2 + target_shape2, diff3:diff3 + target_shape3])\n",
        "\n",
        "        return torch.cat(inputs_, dim=self.dim)\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self._modules)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YNd3-5uE9S9O"
      },
      "source": [
        "class SkipEncoderDecoder(nn.Module):\n",
        "    def __init__(self, input_depth, num_channels_down = [128] * 5, num_channels_up = [128] * 5, num_channels_skip = [128] * 5):\n",
        "        super(SkipEncoderDecoder, self).__init__()\n",
        "\n",
        "        self.model = nn.Sequential()\n",
        "        model_tmp = self.model\n",
        "\n",
        "        for i in range(len(num_channels_down)):\n",
        "\n",
        "            deeper = nn.Sequential()\n",
        "            skip = nn.Sequential()\n",
        "\n",
        "            if num_channels_skip[i] != 0:\n",
        "                model_tmp.add_module(str(len(model_tmp) + 1), Concat(1, skip, deeper))\n",
        "            else:\n",
        "                model_tmp.add_module(str(len(model_tmp) + 1), deeper)\n",
        "            \n",
        "            model_tmp.add_module(str(len(model_tmp) + 1), nn.BatchNorm2d(num_channels_skip[i] + (num_channels_up[i + 1] if i < (len(num_channels_down) - 1) else num_channels_down[i])))\n",
        "\n",
        "            if num_channels_skip[i] != 0:\n",
        "                skip.add_module(str(len(skip) + 1), Conv2dBlock(input_depth, num_channels_skip[i], 1, bias = False))\n",
        "                \n",
        "            deeper.add_module(str(len(deeper) + 1), Conv2dBlock(input_depth, num_channels_down[i], 3, 2, bias = False))\n",
        "            deeper.add_module(str(len(deeper) + 1), Conv2dBlock(num_channels_down[i], num_channels_down[i], 3, bias = False))\n",
        "\n",
        "            deeper_main = nn.Sequential()\n",
        "\n",
        "            if i == len(num_channels_down) - 1:\n",
        "                k = num_channels_down[i]\n",
        "            else:\n",
        "                deeper.add_module(str(len(deeper) + 1), deeper_main)\n",
        "                k = num_channels_up[i + 1]\n",
        "\n",
        "            deeper.add_module(str(len(deeper) + 1), nn.Upsample(scale_factor = 2, mode = 'nearest'))\n",
        "\n",
        "            model_tmp.add_module(str(len(model_tmp) + 1), Conv2dBlock(num_channels_skip[i] + k, num_channels_up[i], 3, 1, bias = False))\n",
        "            model_tmp.add_module(str(len(model_tmp) + 1), Conv2dBlock(num_channels_up[i], num_channels_up[i], 1, bias = False))\n",
        "\n",
        "            input_depth = num_channels_down[i]\n",
        "            model_tmp = deeper_main\n",
        "\n",
        "        self.model.add_module(str(len(self.model) + 1), nn.Conv2d(num_channels_up[0], 3, 1, bias = True))\n",
        "        self.model.add_module(str(len(self.model) + 1), nn.Sigmoid())\n",
        "    \n",
        "    def forward(self, x):\n",
        "        return self.model(x)\n",
        "\n",
        "def input_noise(input_depth, spatial_size, scale = 1.0/10):\n",
        "    shape = [1, input_depth, spatial_size[0], spatial_size[1]]\n",
        "    return torch.rand(*shape) * scale"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "quKNwM0D8uaD",
        "outputId": "3ddb875a-5303-4da0-985d-5cad3107ea10",
        "tags": []
      },
      "source": [
        "generator = SkipEncoderDecoder(\n",
        "    INPUT_DEPTH,\n",
        "    num_channels_down = [128] * 5,\n",
        "    num_channels_up = [128] * 5,\n",
        "    num_channels_skip = [128] * 5\n",
        ").type(DTYPE)\n",
        "generator_input = input_noise(INPUT_DEPTH, watermarked_np.shape[1:]).type(DTYPE)\n",
        "summary(generator, generator_input.shape[1:])\n",
        "\n",
        "objective = torch.nn.MSELoss().type(DTYPE)\n",
        "optimizer = optim.Adam(generator.parameters(), LR)\n",
        "\n",
        "watermarked_var = np_to_torch_array(watermarked_np).type(DTYPE)\n",
        "watermark_var = np_to_torch_array(watermark_np).type(DTYPE)\n",
        "\n",
        "generator_input_saved = generator_input.detach().clone()\n",
        "noise = generator_input.detach().clone()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pEe6N-ve8ufr",
        "tags": []
      },
      "source": [
        "if os.path.isdir('progress'):\n",
        "    rmtree('progress')\n",
        "os.mkdir('progress')\n",
        "\n",
        "for step in range(TRAINING_STEPS):\n",
        "    optimizer.zero_grad()\n",
        "    generator_input = generator_input_saved\n",
        "\n",
        "    if REG_NOISE > 0:\n",
        "        generator_input = generator_input_saved + (noise.normal_() * REG_NOISE)\n",
        "        \n",
        "    output = generator(generator_input)\n",
        "   \n",
        "    loss = objective(output * watermark_var, watermarked_var)\n",
        "    loss.backward()\n",
        "        \n",
        "    if step % SHOW_STEP == 0:\n",
        "        output_image = torch_to_np_array(output)\n",
        "        save_image(output_image, step)\n",
        "        visualize_sample(output_image, nrow = 1, size_factor = 5)\n",
        "        print(f'Step: {step} | Loss: {loss.item()}')\n",
        "        \n",
        "    optimizer.step()"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}