MasksGAN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "MasksGAN.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "metadata": {
        "id": "ZkMW2EIv1QH4",
        "colab_type": "code",
        "outputId": "b7b70609-1abe-4371-b49f-52e2daffe03c",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 179
        }
      },
      "cell_type": "code",
      "source": [
        "!pip3 install -q torch torchvision\n",
        "!pip install --no-cache-dir -I pillow"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Collecting pillow\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/62/94/5430ebaa83f91cc7a9f687ff5238e26164a779cca2ef9903232268b0a318/Pillow-5.3.0-cp36-cp36m-manylinux1_x86_64.whl (2.0MB)\n",
            "\r\u001b[K    0% |▏                               | 10kB 12.3MB/s eta 0:00:01\r\u001b[K    1% |▎                               | 20kB 1.9MB/s eta 0:00:02\r\u001b[K    1% |▌                               | 30kB 2.3MB/s eta 0:00:01\r\u001b[K    2% |▋                               | 40kB 2.1MB/s eta 0:00:01\r\u001b[K    2% |▉                               | 51kB 2.3MB/s eta 0:00:01\r\u001b[K    3% |█                               | 61kB 2.8MB/s eta 0:00:01\r\u001b[K    3% |█▏                              | 71kB 2.8MB/s eta 0:00:01\r\u001b[K    4% |█▎                              | 81kB 2.8MB/s eta 0:00:01\r\u001b[K    4% |█▌                              | 92kB 3.1MB/s eta 0:00:01\r\u001b[K    5% |█▋                              | 102kB 3.1MB/s eta 0:00:01\r\u001b[K    5% |█▉                              | 112kB 3.2MB/s eta 0:00:01\r\u001b[K    6% |██                              | 122kB 4.0MB/s eta 0:00:01\r\u001b[K    6% |██▏                             | 133kB 3.8MB/s eta 0:00:01\r\u001b[K    7% |██▎                             | 143kB 5.0MB/s eta 0:00:01\r\u001b[K    7% |██▌                             | 153kB 4.8MB/s eta 0:00:01\r\u001b[K    8% |██▋                             | 163kB 4.2MB/s eta 0:00:01\r\u001b[K    8% |██▉                             | 174kB 4.8MB/s eta 0:00:01\r\u001b[K    9% |███                             | 184kB 5.0MB/s eta 0:00:01\r\u001b[K    9% |███▏                            | 194kB 5.0MB/s eta 0:00:01\r\u001b[K    10% |███▎                            | 204kB 5.4MB/s eta 0:00:01\r\u001b[K    10% |███▌                            | 215kB 5.0MB/s eta 0:00:01\r\u001b[K    11% |███▋                            | 225kB 5.2MB/s eta 0:00:01\r\u001b[K    11% |███▉                            | 235kB 5.1MB/s eta 0:00:01\r\u001b[K    12% |████                            | 245kB 5.1MB/s eta 0:00:01\r\u001b[K    12% |████▏                           | 256kB 5.1MB/s eta 0:00:01\r\u001b[K    13% |████▎                           | 266kB 6.1MB/s eta 0:00:01\r\u001b[K    13% |████▌                           | 276kB 5.5MB/s eta 0:00:01\r\u001b[K    14% |████▋                           | 286kB 5.4MB/s eta 0:00:01\r\u001b[K    15% |████▉                           | 296kB 5.4MB/s eta 0:00:01\r\u001b[K    15% |█████                           | 307kB 5.1MB/s eta 0:00:01\r\u001b[K    16% |█████▏                          | 317kB 5.2MB/s eta 0:00:01\r\u001b[K    16% |█████▎                          | 327kB 5.9MB/s eta 0:00:01\r\u001b[K    17% |█████▌                          | 337kB 6.4MB/s eta 0:00:01\r\u001b[K    17% |█████▋                          | 348kB 6.4MB/s eta 0:00:01\r\u001b[K    18% |█████▉                          | 358kB 6.0MB/s eta 0:00:01\r\u001b[K    18% |██████                          | 368kB 5.4MB/s eta 0:00:01\r\u001b[K    19% |██████▏                         | 378kB 5.9MB/s eta 0:00:01\r\u001b[K    19% |██████▎                         | 389kB 6.8MB/s eta 0:00:01\r\u001b[K    20% |██████▌                         | 399kB 5.7MB/s eta 0:00:01\r\u001b[K    20% |██████▋                         | 409kB 6.7MB/s eta 0:00:01\r\u001b[K    21% |██████▉                         | 419kB 6.2MB/s eta 0:00:01\r\u001b[K    21% |███████                         | 430kB 6.1MB/s eta 0:00:01\r\u001b[K    22% |███████▏                        | 440kB 6.5MB/s eta 0:00:01\r\u001b[K    22% |███████▎                        | 450kB 5.4MB/s eta 0:00:01\r\u001b[K    23% |███████▌                        | 460kB 6.6MB/s eta 0:00:01\r\u001b[K    23% |███████▋                        | 471kB 6.8MB/s eta 0:00:01\r\u001b[K    24% |███████▉                        | 481kB 5.8MB/s eta 0:00:01\r\u001b[K    24% |████████                        | 491kB 6.3MB/s eta 0:00:01\r\u001b[K    25% |████████▏                       | 501kB 7.0MB/s eta 0:00:01\r\u001b[K    25% |████████▎                       | 512kB 6.9MB/s eta 0:00:01\r\u001b[K    26% |████████▌                       | 522kB 6.7MB/s eta 0:00:01\r\u001b[K    26% |████████▋                       | 532kB 6.1MB/s eta 0:00:01\r\u001b[K    27% |████████▉                       | 542kB 6.6MB/s eta 0:00:01\r\u001b[K    27% |█████████                       | 552kB 7.0MB/s eta 0:00:01\r\u001b[K    28% |█████████▏                      | 563kB 6.5MB/s eta 0:00:01\r\u001b[K    29% |█████████▎                      | 573kB 7.2MB/s eta 0:00:01\r\u001b[K    29% |█████████▌                      | 583kB 7.3MB/s eta 0:00:01\r\u001b[K    30% |█████████▋                      | 593kB 7.4MB/s eta 0:00:01\r\u001b[K    30% |█████████▉                      | 604kB 6.7MB/s eta 0:00:01\r\u001b[K    31% |██████████                      | 614kB 6.2MB/s eta 0:00:01\r\u001b[K    31% |██████████▏                     | 624kB 7.5MB/s eta 0:00:01\r\u001b[K    32% |██████████▎                     | 634kB 7.3MB/s eta 0:00:01\r\u001b[K    32% |██████████▌                     | 645kB 6.6MB/s eta 0:00:01\r\u001b[K    33% |██████████▋                     | 655kB 7.8MB/s eta 0:00:01\r\u001b[K    33% |██████████▊                     | 665kB 7.3MB/s eta 0:00:01\r\u001b[K    34% |███████████                     | 675kB 7.3MB/s eta 0:00:01\r\u001b[K    34% |███████████                     | 686kB 7.6MB/s eta 0:00:01\r\u001b[K    35% |███████████▎                    | 696kB 6.5MB/s eta 0:00:01\r\u001b[K    35% |███████████▍                    | 706kB 8.0MB/s eta 0:00:01\r\u001b[K    36% |███████████▋                    | 716kB 7.8MB/s eta 0:00:01\r\u001b[K    36% |███████████▊                    | 727kB 6.9MB/s eta 0:00:01\r\u001b[K    37% |████████████                    | 737kB 8.0MB/s eta 0:00:01\r\u001b[K    37% |████████████                    | 747kB 8.3MB/s eta 0:00:01\r\u001b[K    38% |████████████▎                   | 757kB 8.3MB/s eta 0:00:01\r\u001b[K    38% |████████████▍                   | 768kB 7.6MB/s eta 0:00:01\r\u001b[K    39% |████████████▋                   | 778kB 7.0MB/s eta 0:00:01\r\u001b[K    39% |████████████▊                   | 788kB 8.0MB/s eta 0:00:01\r\u001b[K    40% |█████████████                   | 798kB 8.5MB/s eta 0:00:01\r\u001b[K    40% |█████████████                   | 808kB 7.5MB/s eta 0:00:01\r\u001b[K    41% |█████████████▎                  | 819kB 8.6MB/s eta 0:00:01\r\u001b[K    41% |█████████████▍                  | 829kB 9.1MB/s eta 0:00:01\r\u001b[K    42% |█████████████▋                  | 839kB 9.1MB/s eta 0:00:01\r\u001b[K    43% |█████████████▊                  | 849kB 7.6MB/s eta 0:00:01\r\u001b[K    43% |██████████████                  | 860kB 7.0MB/s eta 0:00:01\r\u001b[K    44% |██████████████                  | 870kB 8.9MB/s eta 0:00:01\r\u001b[K    44% |██████████████▎                 | 880kB 8.7MB/s eta 0:00:01\r\u001b[K    45% |██████████████▍                 | 890kB 7.3MB/s eta 0:00:01\r\u001b[K    45% |██████████████▋                 | 901kB 8.1MB/s eta 0:00:01\r\u001b[K    46% |██████████████▊                 | 911kB 9.1MB/s eta 0:00:01\r\u001b[K    46% |███████████████                 | 921kB 9.1MB/s eta 0:00:01\r\u001b[K    47% |███████████████                 | 931kB 7.5MB/s eta 0:00:01\r\u001b[K    47% |███████████████▎                | 942kB 7.2MB/s eta 0:00:01\r\u001b[K    48% |███████████████▍                | 952kB 9.4MB/s eta 0:00:01\r\u001b[K    48% |███████████████▋                | 962kB 8.7MB/s eta 0:00:01\r\u001b[K    49% |███████████████▊                | 972kB 8.3MB/s eta 0:00:01\r\u001b[K    49% |████████████████                | 983kB 9.3MB/s eta 0:00:01\r\u001b[K    50% |████████████████                | 993kB 10.3MB/s eta 0:00:01\r\u001b[K    50% |████████████████▎               | 1.0MB 10.3MB/s eta 0:00:01\r\u001b[K    51% |████████████████▍               | 1.0MB 8.3MB/s eta 0:00:01\r\u001b[K    51% |████████████████▋               | 1.0MB 7.2MB/s eta 0:00:01\r\u001b[K    52% |████████████████▊               | 1.0MB 9.4MB/s eta 0:00:01\r\u001b[K    52% |█████████████████               | 1.0MB 8.7MB/s eta 0:00:01\r\u001b[K    53% |█████████████████               | 1.1MB 8.4MB/s eta 0:00:01\r\u001b[K    53% |█████████████████▎              | 1.1MB 10.0MB/s eta 0:00:01\r\u001b[K    54% |█████████████████▍              | 1.1MB 9.6MB/s eta 0:00:01\r\u001b[K    54% |█████████████████▋              | 1.1MB 9.6MB/s eta 0:00:01\r\u001b[K    55% |█████████████████▊              | 1.1MB 9.4MB/s eta 0:00:01\r\u001b[K    55% |██████████████████              | 1.1MB 7.3MB/s eta 0:00:01\r\u001b[K    56% |██████████████████              | 1.1MB 9.0MB/s eta 0:00:01\r\u001b[K    56% |██████████████████▎             | 1.1MB 10.3MB/s eta 0:00:01\r\u001b[K    57% |██████████████████▍             | 1.1MB 9.5MB/s eta 0:00:01\r\u001b[K    58% |██████████████████▋             | 1.1MB 10.9MB/s eta 0:00:01\r\u001b[K    58% |██████████████████▊             | 1.2MB 10.0MB/s eta 0:00:01\r\u001b[K    59% |███████████████████             | 1.2MB 10.0MB/s eta 0:00:01\r\u001b[K    59% |███████████████████             | 1.2MB 9.3MB/s eta 0:00:01\r\u001b[K    60% |███████████████████▎            | 1.2MB 8.7MB/s eta 0:00:01\r\u001b[K    60% |███████████████████▍            | 1.2MB 10.0MB/s eta 0:00:01\r\u001b[K    61% |███████████████████▋            | 1.2MB 11.6MB/s eta 0:00:01\r\u001b[K    61% |███████████████████▊            | 1.2MB 9.8MB/s eta 0:00:01\r\u001b[K    62% |████████████████████            | 1.2MB 10.1MB/s eta 0:00:01\r\u001b[K    62% |████████████████████            | 1.2MB 10.2MB/s eta 0:00:01\r\u001b[K    63% |████████████████████▎           | 1.2MB 10.2MB/s eta 0:00:01\r\u001b[K    63% |████████████████████▍           | 1.3MB 10.3MB/s eta 0:00:01\r\u001b[K    64% |████████████████████▋           | 1.3MB 9.5MB/s eta 0:00:01\r\u001b[K    64% |████████████████████▊           | 1.3MB 11.3MB/s eta 0:00:01\r\u001b[K    65% |█████████████████████           | 1.3MB 10.3MB/s eta 0:00:01\r\u001b[K    65% |█████████████████████           | 1.3MB 8.6MB/s eta 0:00:01\r\u001b[K    66% |█████████████████████▎          | 1.3MB 10.1MB/s eta 0:00:01\r\u001b[K    66% |█████████████████████▍          | 1.3MB 10.6MB/s eta 0:00:01\r\u001b[K    67% |█████████████████████▌          | 1.3MB 10.5MB/s eta 0:00:01\r\u001b[K    67% |█████████████████████▊          | 1.3MB 11.0MB/s eta 0:00:01\r\u001b[K    68% |█████████████████████▉          | 1.4MB 9.3MB/s eta 0:00:01\r\u001b[K    68% |██████████████████████          | 1.4MB 10.6MB/s eta 0:00:01\r\u001b[K    69% |██████████████████████▏         | 1.4MB 10.3MB/s eta 0:00:01\r\u001b[K    69% |██████████████████████▍         | 1.4MB 9.4MB/s eta 0:00:01\r\u001b[K    70% |██████████████████████▌         | 1.4MB 10.8MB/s eta 0:00:01\r\u001b[K    70% |██████████████████████▊         | 1.4MB 13.2MB/s eta 0:00:01\r\u001b[K    71% |██████████████████████▉         | 1.4MB 13.2MB/s eta 0:00:01\r\u001b[K    72% |███████████████████████         | 1.4MB 12.2MB/s eta 0:00:01\r\u001b[K    72% |███████████████████████▏        | 1.4MB 11.0MB/s eta 0:00:01\r\u001b[K    73% |███████████████████████▍        | 1.4MB 9.5MB/s eta 0:00:01\r\u001b[K    73% |███████████████████████▌        | 1.5MB 11.2MB/s eta 0:00:01\r\u001b[K    74% |███████████████████████▊        | 1.5MB 10.5MB/s eta 0:00:01\r\u001b[K    74% |███████████████████████▉        | 1.5MB 10.3MB/s eta 0:00:01\r\u001b[K    75% |████████████████████████        | 1.5MB 11.4MB/s eta 0:00:01\r\u001b[K    75% |████████████████████████▏       | 1.5MB 10.7MB/s eta 0:00:01\r\u001b[K    76% |████████████████████████▍       | 1.5MB 11.0MB/s eta 0:00:01\r\u001b[K    76% |████████████████████████▌       | 1.5MB 9.9MB/s eta 0:00:01\r\u001b[K    77% |████████████████████████▊       | 1.5MB 10.2MB/s eta 0:00:01\r\u001b[K    77% |████████████████████████▉       | 1.5MB 11.3MB/s eta 0:00:01\r\u001b[K    78% |█████████████████████████       | 1.5MB 12.8MB/s eta 0:00:01\r\u001b[K    78% |█████████████████████████▏      | 1.6MB 11.6MB/s eta 0:00:01\r\u001b[K    79% |█████████████████████████▍      | 1.6MB 12.4MB/s eta 0:00:01\r\u001b[K    79% |█████████████████████████▌      | 1.6MB 12.2MB/s eta 0:00:01\r\u001b[K    80% |█████████████████████████▊      | 1.6MB 12.3MB/s eta 0:00:01\r\u001b[K    80% |█████████████████████████▉      | 1.6MB 11.8MB/s eta 0:00:01\r\u001b[K    81% |██████████████████████████      | 1.6MB 10.5MB/s eta 0:00:01\r\u001b[K    81% |██████████████████████████▏     | 1.6MB 11.8MB/s eta 0:00:01\r\u001b[K    82% |██████████████████████████▍     | 1.6MB 13.7MB/s eta 0:00:01\r\u001b[K    82% |██████████████████████████▌     | 1.6MB 10.5MB/s eta 0:00:01\r\u001b[K    83% |██████████████████████████▊     | 1.6MB 11.3MB/s eta 0:00:01\r\u001b[K    83% |██████████████████████████▉     | 1.7MB 11.5MB/s eta 0:00:01\r\u001b[K    84% |███████████████████████████     | 1.7MB 11.5MB/s eta 0:00:01\r\u001b[K    84% |███████████████████████████▏    | 1.7MB 12.6MB/s eta 0:00:01\r\u001b[K    85% |███████████████████████████▍    | 1.7MB 11.5MB/s eta 0:00:01\r\u001b[K    86% |███████████████████████████▌    | 1.7MB 13.2MB/s eta 0:00:01\r\u001b[K    86% |███████████████████████████▊    | 1.7MB 14.3MB/s eta 0:00:01\r\u001b[K    87% |███████████████████████████▉    | 1.7MB 10.5MB/s eta 0:00:01\r\u001b[K    87% |████████████████████████████    | 1.7MB 10.7MB/s eta 0:00:01\r\u001b[K    88% |████████████████████████████▏   | 1.7MB 13.7MB/s eta 0:00:01\r\u001b[K    88% |████████████████████████████▍   | 1.8MB 13.7MB/s eta 0:00:01\r\u001b[K    89% |████████████████████████████▌   | 1.8MB 14.2MB/s eta 0:00:01\r\u001b[K    89% |████████████████████████████▊   | 1.8MB 12.0MB/s eta 0:00:01\r\u001b[K    90% |████████████████████████████▉   | 1.8MB 12.8MB/s eta 0:00:01\r\u001b[K    90% |█████████████████████████████   | 1.8MB 12.8MB/s eta 0:00:01\r\u001b[K    91% |█████████████████████████████▏  | 1.8MB 11.6MB/s eta 0:00:01\r\u001b[K    91% |█████████████████████████████▍  | 1.8MB 12.4MB/s eta 0:00:01\r\u001b[K    92% |█████████████████████████████▌  | 1.8MB 16.2MB/s eta 0:00:01\r\u001b[K    92% |█████████████████████████████▊  | 1.8MB 27kB/s eta 0:00:06\r\u001b[K    93% |█████████████████████████████▉  | 1.8MB 27kB/s eta 0:00:05\r\u001b[K    93% |██████████████████████████████  | 1.9MB 27kB/s eta 0:00:05\r\u001b[K    94% |██████████████████████████████▏ | 1.9MB 27kB/s eta 0:00:05\r\u001b[K    94% |██████████████████████████████▍ | 1.9MB 27kB/s eta 0:00:04\r\u001b[K    95% |██████████████████████████████▌ | 1.9MB 27kB/s eta 0:00:04\r\u001b[K    95% |██████████████████████████████▊ | 1.9MB 27kB/s eta 0:00:04\r\u001b[K    96% |██████████████████████████████▉ | 1.9MB 27kB/s eta 0:00:03\r\u001b[K    96% |███████████████████████████████ | 1.9MB 27kB/s eta 0:00:03\r\u001b[K    97% |███████████████████████████████▏| 1.9MB 27kB/s eta 0:00:02\r\u001b[K    97% |███████████████████████████████▍| 1.9MB 42.9MB/s eta 0:00:01\r\u001b[K    98% |███████████████████████████████▌| 1.9MB 52.7MB/s eta 0:00:01\r\u001b[K    98% |███████████████████████████████▊| 2.0MB 55.3MB/s eta 0:00:01\r\u001b[K    99% |███████████████████████████████▉| 2.0MB 55.2MB/s eta 0:00:01\r\u001b[K    99% |████████████████████████████████| 2.0MB 56.2MB/s eta 0:00:01\r\u001b[K    100% |████████████████████████████████| 2.0MB 58.7MB/s \n",
            "\u001b[?25hInstalling collected packages: pillow\n",
            "Successfully installed pillow-5.3.0\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "mAQOmU1g1hIm",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# based on https://github.com/zeruniverse/neural-colorization\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import argparse\n",
        "from torch.autograd import Variable\n",
        "import torchvision.models as models\n",
        "import torch.nn.functional as F\n",
        "import os\n",
        "from torch.utils import data\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "from skimage.color import rgb2yuv,yuv2rgb\n",
        "from scipy.ndimage import zoom\n",
        "import cv2\n",
        "from functools import reduce\n",
        "\n",
        "import math"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "sw99cqTF1m65",
        "colab_type": "code",
        "outputId": "b161521a-7f6c-4597-866d-6e7b97b2f982",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "cell_type": "code",
      "source": [
        "SIZE = 224\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "device"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "device(type='cuda')"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 3
        }
      ]
    },
    {
      "metadata": {
        "id": "tHAH0DS71nGz",
        "colab_type": "code",
        "outputId": "f9f47af6-1bd4-4201-d21a-8d17c3eadc15",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "cell_type": "code",
      "source": [
        "import os\n",
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')\n",
        "base_dir = '/content/gdrive/My Drive/ColorMaskGAN'"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Drive already mounted at /content/gdrive; to attempt to forcibly remount, call drive.mount(\"/content/gdrive\", force_remount=True).\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "DlzkeUMD1nKC",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "class ColorGenerator(nn.Module):\n",
        "    '''Generator'''\n",
        "    def __init__(self):\n",
        "        super(ColorGenerator, self).__init__()\n",
        "\n",
        "        self.conv0_1 = nn.Conv2d(1,16, 3, stride=1, padding=1, bias=False)\n",
        "        self.conv0_2 = nn.Conv2d(1,16, 3, stride=1, padding=1, bias=False)\n",
        "        self.bn0 = nn.BatchNorm2d(16)\n",
        "        self.relu0 = nn.LeakyReLU(0.1)\n",
        "        \n",
        "# 1 -> 32\n",
        "        self.conv1 = nn.Conv2d(32, 64, 3, stride=2, padding=1, bias=False)\n",
        "        self.bn1 = nn.BatchNorm2d(64)\n",
        "        self.relu1 = nn.LeakyReLU(0.1)\n",
        "\n",
        "        self.conv2 = nn.Conv2d(64, 128, 3, stride=2, padding=1, bias=False)\n",
        "        self.bn2 = nn.BatchNorm2d(128)\n",
        "        self.relu2 = nn.LeakyReLU(0.1)\n",
        "\n",
        "        self.conv3 = nn.Conv2d(128, 256, 3, stride=2, padding=1, bias=False)\n",
        "        self.bn3 = nn.BatchNorm2d(256)\n",
        "        self.relu3 = nn.LeakyReLU(0.1)\n",
        "\n",
        "        self.conv4 = nn.Conv2d(256, 512, 3, stride=2, padding=1, bias=False)\n",
        "        self.bn4 = nn.BatchNorm2d(512)\n",
        "        self.relu4 = nn.LeakyReLU(0.1)\n",
        "\n",
        "        self.conv5 = nn.Conv2d(512, 512, 3, stride=2, padding=1, bias=False)\n",
        "        self.bn5 = nn.BatchNorm2d(512)\n",
        "        self.relu5 = nn.LeakyReLU(0.1)\n",
        "\n",
        "        self.deconv6 = nn.ConvTranspose2d(512, 512, 3, stride=2, padding=1, output_padding=1, bias=False)\n",
        "        self.bn6 = nn.BatchNorm2d(512)\n",
        "        self.relu6 = nn.ReLU()\n",
        "\n",
        "        self.deconv7 = nn.ConvTranspose2d(512, 256, 3, stride=2, padding=1, output_padding=1, bias=False)\n",
        "        self.bn7 = nn.BatchNorm2d(256)\n",
        "        self.relu7 = nn.ReLU()\n",
        "\n",
        "        self.deconv8 = nn.ConvTranspose2d(256, 128, 3, stride=2, padding=1, output_padding=1, bias=False)\n",
        "        self.bn8 = nn.BatchNorm2d(128)\n",
        "        self.relu8 = nn.ReLU()\n",
        "\n",
        "        self.deconv9 = nn.ConvTranspose2d(128, 64, 3, stride=2, padding=1, output_padding=1, bias=False)\n",
        "        self.bn9 = nn.BatchNorm2d(64)\n",
        "        self.relu9 = nn.ReLU()\n",
        "#Modifiquei para gerar só UV\n",
        "        self.deconv10 = nn.ConvTranspose2d(64, 2, 3, stride=2, padding=1, output_padding=1, bias=False)\n",
        "        self.bn10 = nn.BatchNorm2d(2)\n",
        "        self.relu10 = nn.ReLU()\n",
        "\n",
        "        self._initialize_weights()\n",
        "\n",
        "    def forward(self, x, y):\n",
        "      \n",
        "        x = self.conv0_1(x)\n",
        "        x = self.bn0(x)\n",
        "        x = self.relu0(x)\n",
        "        \n",
        "        y = self.conv0_2(y)\n",
        "        y = self.bn0(y)\n",
        "        y = self.relu0(y)\n",
        "\n",
        "#        print(x.shape, y.shape)\n",
        "        x = torch.cat((x, y), 1)\n",
        "      \n",
        "        h = x\n",
        "        h = self.conv1(h)\n",
        "        h = self.bn1(h)\n",
        "        h = self.relu1(h) # 64,112,112 (if input is 224x224)\n",
        "        pool1 = h\n",
        "\n",
        "        h = self.conv2(h)\n",
        "        h = self.bn2(h)\n",
        "        h = self.relu2(h) # 128,56,56\n",
        "        pool2 = h\n",
        "\n",
        "        h = self.conv3(h) # 256,28,28\n",
        "        h = self.bn3(h)\n",
        "        h = self.relu3(h)\n",
        "        pool3 =h\n",
        "\n",
        "        h = self.conv4(h) # 512,14,14\n",
        "        h = self.bn4(h)\n",
        "        h = self.relu4(h)\n",
        "        pool4 = h\n",
        "\n",
        "        h = self.conv5(h) # 512,7,7\n",
        "        h = self.bn5(h)\n",
        "        h = self.relu5(h)\n",
        "\n",
        "        h = self.deconv6(h)\n",
        "        h = self.bn6(h)\n",
        "        h = self.relu6(h) # 512,14,14\n",
        "        h += pool4\n",
        "\n",
        "        h = self.deconv7(h)\n",
        "        h = self.bn7(h)\n",
        "        h = self.relu7(h) # 256,28,28\n",
        "        h += pool3\n",
        "\n",
        "        h = self.deconv8(h)\n",
        "        h = self.bn8(h)\n",
        "        h = self.relu8(h) # 128,56,56\n",
        "        h += pool2\n",
        "\n",
        "        h = self.deconv9(h)\n",
        "        h = self.bn9(h)\n",
        "        h = self.relu9(h) # 64,112,112\n",
        "        h += pool1\n",
        "\n",
        "        h = self.deconv10(h)\n",
        "        h = F.tanh(h) # 3,224,224\n",
        "\n",
        "        return h\n",
        "\n",
        "    def _initialize_weights(self):\n",
        "        for m in self.modules():\n",
        "            if isinstance(m, nn.Conv2d):\n",
        "                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n",
        "                m.weight.data.normal_(0, math.sqrt(2. / n))\n",
        "            if isinstance(m, nn.ConvTranspose2d):\n",
        "                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n",
        "                m.weight.data.normal_(0, math.sqrt(2. / n))\n",
        "\n",
        "class ColorDiscriminator(nn.Module):\n",
        "    '''Discriminator'''\n",
        "    def __init__(self, large=False):\n",
        "        super(ColorDiscriminator, self).__init__()\n",
        "        \n",
        "#         self.conv0_1 = nn.Conv2d(1,16, 3, stride=1, padding=1, bias=False)\n",
        "#         self.conv0_2 = nn.Conv2d(1,16, 3, stride=1, padding=1, bias=False)\n",
        "#         self.bn0 = nn.BatchNorm2d(16)\n",
        "#         self.relu0 = nn.LeakyReLU(0.1)\n",
        "#CHANGE CHANNELS        \n",
        "        self.conv1 = nn.Conv2d(4, 64, 3, stride=2, padding=1, bias=False)\n",
        "        self.bn1 = nn.BatchNorm2d(64)\n",
        "        self.relu1 = nn.LeakyReLU(0.1)\n",
        "\n",
        "        self.conv2 = nn.Conv2d(64, 128, 3, stride=2, padding=1, bias=False)\n",
        "        self.bn2 = nn.BatchNorm2d(128)\n",
        "        self.relu2 = nn.LeakyReLU(0.1)\n",
        "\n",
        "        self.conv3 = nn.Conv2d(128, 256, 3, stride=2, padding=1, bias=False)\n",
        "        self.bn3 = nn.BatchNorm2d(256)\n",
        "        self.relu3 = nn.LeakyReLU(0.1)\n",
        "\n",
        "        self.conv4 = nn.Conv2d(256, 512, 3, stride=2, padding=1, bias=False)\n",
        "        self.bn4 = nn.BatchNorm2d(512)\n",
        "        self.relu4 = nn.LeakyReLU(0.1)\n",
        "\n",
        "        self.conv5 = nn.Conv2d(512, 512, 3, stride=2, padding=1, bias=False)\n",
        "        self.bn5 = nn.BatchNorm2d(512)\n",
        "        self.relu5 = nn.LeakyReLU(0.1)\n",
        "\n",
        "        if large:\n",
        "            self.conv6 = nn.Conv2d(512, 512, 15, stride=1, padding=0, bias=False)\n",
        "        else:\n",
        "            self.conv6 = nn.Conv2d(512, 512, 7, stride=1, padding=0, bias=False)\n",
        "        self.bn6 = nn.BatchNorm2d(512)\n",
        "        self.relu6 = nn.LeakyReLU(0.1)\n",
        "\n",
        "        self.conv7 = nn.Conv2d(512, 1, 1, stride=1, padding=0, bias=False)\n",
        "\n",
        "        self._initialize_weights()\n",
        "\n",
        "    def forward(self, x, y):\n",
        "      \n",
        "#         x = self.conv0_1(x)\n",
        "#         x = self.bn0(x)\n",
        "#         x = self.relu0(x)\n",
        "        \n",
        "#         y = self.conv0_2(y)\n",
        "#         y = self.bn0(y)\n",
        "#         y = self.relu0(y)\n",
        "        \n",
        "#         x = torch.cat((x, y), 1)\n",
        "        \n",
        "#        h = x\n",
        "        h = torch.cat((x, y), 1)\n",
        "        h = self.conv1(h)\n",
        "        h = self.bn1(h)\n",
        "        h = self.relu1(h) # 64,112,112 (if input is 224x224)\n",
        "\n",
        "        h = self.conv2(h)\n",
        "        h = self.bn2(h)\n",
        "        h = self.relu2(h) # 128,56,56\n",
        "\n",
        "        h = self.conv3(h) # 256,28,28\n",
        "        h = self.bn3(h)\n",
        "        h = self.relu3(h)\n",
        "\n",
        "        h = self.conv4(h) # 512,14,14\n",
        "        h = self.bn4(h)\n",
        "        h = self.relu4(h)\n",
        "\n",
        "        h = self.conv5(h) # 512,7,7\n",
        "        h = self.bn5(h)\n",
        "        h = self.relu5(h)\n",
        "\n",
        "        h = self.conv6(h)\n",
        "        h = self.bn6(h)\n",
        "        h = self.relu6(h) # 512,1,1\n",
        "\n",
        "        h = self.conv7(h)\n",
        "        h = F.sigmoid(h)\n",
        "\n",
        "        return h\n",
        "\n",
        "    def _initialize_weights(self):\n",
        "        for m in self.modules():\n",
        "            if isinstance(m, nn.Conv2d):\n",
        "                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n",
        "                m.weight.data.normal_(0, math.sqrt(2. / n))\n",
        "            if isinstance(m, nn.ConvTranspose2d):\n",
        "                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n",
        "                m.weight.data.normal_(0, math.sqrt(2. / n))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "yUCr7Mov1nMs",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "class dotdict(dict):\n",
        "    \"\"\"dot.notation access to dictionary attributes\"\"\"\n",
        "    __getattr__ = dict.get\n",
        "    __setattr__ = dict.__setitem__\n",
        "    __delattr__ = dict.__delitem__\n",
        "\n",
        "args = {'training_dir':os.path.join(base_dir, \"train\"),\n",
        "        'checkpoint_location': os.path.join(base_dir, \"checkpoints\"),\n",
        "        'epoch': 120,\n",
        "        'gpu': 0,\n",
        "        'batch_size':12,\n",
        "        'num_workers':6,\n",
        "        'pixel_loss_weights':1000,\n",
        "        \"g_every\":1,\n",
        "        \"g_lr\":1e-4,\n",
        "        \"d_lr\":1e-4,\n",
        "        \"checkpoint_every\":1000,\n",
        "        \"d_init\":os.path.join(base_dir, \"D0.pth\"),\n",
        "        \"g_init\":os.path.join(base_dir, \"G0.pth\"),\n",
        "        \"model\": os.path.join(base_dir, \"G0.pth\"),\n",
        "        \"test_image\": os.path.join(base_dir, \"test.jpg\")\n",
        "       }\n",
        "\n",
        "args = dotdict(args)\n",
        "\n",
        "# define data generator\n",
        "class img_data(data.Dataset):\n",
        "    def __init__(self, path):\n",
        "        files = os.listdir(path)\n",
        "        self.files = [os.path.join(path,x) for x in files]\n",
        "    def __len__(self):\n",
        "        return len(self.files)\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        img = Image.open(self.files[index])\n",
        "        yuv = rgb2yuv(img)\n",
        "        y = (yuv[...,0])-0.5\n",
        "        u_t = yuv[...,1] / 0.43601035\n",
        "        v_t = yuv[...,2] / 0.61497538\n",
        "        return torch.Tensor(np.expand_dims(y,axis=0)),torch.Tensor(np.stack([u_t,v_t],axis=0))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "shyBm0RN1nPV",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# mask = Image.open(os.path.join(base_dir, 'mask.jpg'))\n",
        "# mask = Variable(torch.Tensor(rgb2yuv(mask)[...,0]).unsqueeze(0).unsqueeze(0)).to(device)\n",
        "mask_loss_weight = 100000"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "GhtRF5If1vaO",
        "colab_type": "code",
        "outputId": "60ec6998-5772-446f-b941-dc9a3cb348bd",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1193
        }
      },
      "cell_type": "code",
      "source": [
        "if not os.path.exists(os.path.join(args.checkpoint_location,'weights')):\n",
        "    os.makedirs(os.path.join(args.checkpoint_location,'weights'))\n",
        "\n",
        "# Define G, same as torch version\n",
        "#G = generator().to(device)\n",
        "# define D\n",
        "#D = models.resnet18(pretrained=False,num_classes=2)\n",
        "#D.fc = nn.Sequential(nn.Linear(2048, 1), nn.Sigmoid())\n",
        "#D = D.to(device)\n",
        "\n",
        "G = ColorGenerator().to(device)\n",
        "D = ColorDiscriminator().to(device)\n",
        "\n",
        "trainset = img_data(args.training_dir)\n",
        "maskset = img_data(os.path.join(base_dir, \"masks\"))\n",
        "params = {'batch_size': args.batch_size,\n",
        "          'shuffle': True,\n",
        "          'num_workers': args.num_workers}\n",
        "training_generator = data.DataLoader(trainset, **params)\n",
        "masks_generator = data.DataLoader(maskset, batch_size=1, shuffle=True)\n",
        "#mask = mask.expand(args.batch_size, -1, -1, -1)\n",
        "\n",
        "if args.test_image is not None:\n",
        "    test_img = Image.open(args.test_image).convert('RGB').resize((SIZE,SIZE))\n",
        "    test_yuv = rgb2yuv(test_img)\n",
        "    test_inf = test_yuv[...,0].reshape(1,1,SIZE,SIZE)\n",
        "    test_var = Variable(torch.Tensor(test_inf-0.5)).to(device)\n",
        "if args.d_init is not None:\n",
        "    D.load_state_dict(torch.load(args.d_init, map_location='cuda:0'))\n",
        "if args.g_init is not None:\n",
        "    G.load_state_dict(torch.load(args.g_init, map_location='cuda:0'))\n",
        "\n",
        "# save test image for beginning\n",
        "\n",
        "i=0\n",
        "adversarial_loss = torch.nn.BCELoss()\n",
        "optimizer_G = torch.optim.Adam(G.parameters(), lr=args.g_lr, betas=(0.5, 0.999))\n",
        "optimizer_D = torch.optim.Adam(D.parameters(), lr=args.d_lr, betas=(0.5, 0.999))\n",
        "for epoch in range(0, args.epoch):\n",
        "    for mask, _ in masks_generator:\n",
        "      mask = mask + 0.5\n",
        "      mask = mask.to(device)\n",
        "      for y, uv in training_generator:\n",
        "          # Adversarial ground truths\n",
        "          valid = Variable(torch.Tensor(y.size(0), 1).fill_(1.0), requires_grad=False).to(device)\n",
        "          fake = Variable(torch.Tensor(y.size(0), 1).fill_(0.0), requires_grad=False).to(device)\n",
        "\n",
        "          yvar = Variable(y).to(device)\n",
        "          uvvar = Variable(uv).to(device)\n",
        "          real_imgs = torch.cat([yvar,uvvar],dim=1)\n",
        "\n",
        "\n",
        "          mask20 = mask.expand(yvar.shape[0], -1, -1, -1)\n",
        "          optimizer_G.zero_grad()\n",
        "          uvgen = G(yvar, mask20) # MODIFIED HERE\n",
        "          # Generate a batch of images\n",
        "          gen_imgs = torch.cat([yvar.detach(),uvgen],dim=1)\n",
        "\n",
        "          # Loss measures generator's ability to fool the discriminator\n",
        "  #        print(gen_imgs.shape, mask20.shape, valid.shape)\n",
        "#           HOP = .view(-1, 1)\n",
        "          g_loss_gan = adversarial_loss(D(gen_imgs, mask20), valid)\n",
        "          g_loss = (g_loss_gan \n",
        "                    + args.pixel_loss_weights * torch.mean((uvvar*mask20-uvgen)**2) \n",
        "#                     + mask_loss_weight * torch.mean(torch.mul(uvvar, mask20)**2)\n",
        "                   )\n",
        "          if i%args.g_every==0:\n",
        "              g_loss.backward()\n",
        "              optimizer_G.step()\n",
        "\n",
        "          optimizer_D.zero_grad()\n",
        "\n",
        "          # Measure discriminator's ability to classify real from generated samples\n",
        "          real_loss = adversarial_loss(D(real_imgs, mask20), valid)\n",
        "          fake_loss = adversarial_loss(D(gen_imgs.detach(), mask20.detach()), fake)\n",
        "          d_loss = (real_loss + fake_loss) / 2\n",
        "          d_loss.backward()\n",
        "          optimizer_D.step()\n",
        "          i+=1\n",
        "          if i%args.checkpoint_every==0:\n",
        "              print (\"Epoch: %d: [D loss: %f] [G total loss: %f] [G GAN Loss: %f]\" % (epoch, d_loss.item(), g_loss.item(), g_loss_gan.item()))\n",
        "\n",
        "              torch.save(D.state_dict(), os.path.join(args.checkpoint_location,'weights','D'+str(epoch)+'.pth'))\n",
        "              torch.save(G.state_dict(), os.path.join(args.checkpoint_location,'weights','G'+str(epoch)+'.pth'))\n",
        "              if args.test_image is not None:\n",
        "                  test_res = G(test_var, mask)\n",
        "                  uv=test_res.cpu().detach().numpy()\n",
        "                  uv[:,0,:,:] *= 0.436\n",
        "                  uv[:,1,:,:] *= 0.615\n",
        "                  test_yuv = np.concatenate([test_inf,uv],axis=1).reshape(3,SIZE,SIZE)\n",
        "                  test_rgb = yuv2rgb(test_yuv.transpose(1,2,0))\n",
        "                  cv2.imwrite(os.path.join(args.checkpoint_location,'test_epoch_'+str(epoch)+'.jpg'),(test_rgb.clip(min=0,max=1)*SIZE)[:,:,[2,1,0]])\n",
        "torch.save(D.state_dict(), os.path.join(args.checkpoint_location,'D_final.pth'))\n",
        "torch.save(G.state_dict(), os.path.join(args.checkpoint_location,'G_final.pth'))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py:995: UserWarning: nn.functional.tanh is deprecated. Use torch.tanh instead.\n",
            "  warnings.warn(\"nn.functional.tanh is deprecated. Use torch.tanh instead.\")\n",
            "/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py:1006: UserWarning: nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.\n",
            "  warnings.warn(\"nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.\")\n",
            "/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py:1594: UserWarning: Using a target size (torch.Size([12, 1])) that is different to the input size (torch.Size([12, 1, 1, 1])) is deprecated. Please ensure they have the same size.\n",
            "  \"Please ensure they have the same size.\".format(target.size(), input.size()))\n",
            "/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py:1594: UserWarning: Using a target size (torch.Size([4, 1])) that is different to the input size (torch.Size([4, 1, 1, 1])) is deprecated. Please ensure they have the same size.\n",
            "  \"Please ensure they have the same size.\".format(target.size(), input.size()))\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Epoch: 0: [D loss: 0.790424] [G total loss: 18.751749] [G GAN Loss: 1.127248]\n",
            "Epoch: 1: [D loss: 0.389419] [G total loss: 17.068918] [G GAN Loss: 1.922459]\n",
            "Epoch: 2: [D loss: 0.084100] [G total loss: 14.632328] [G GAN Loss: 2.465448]\n",
            "Epoch: 3: [D loss: 0.276594] [G total loss: 14.659258] [G GAN Loss: 3.302816]\n",
            "Epoch: 4: [D loss: 0.006551] [G total loss: 17.098385] [G GAN Loss: 5.177050]\n",
            "Epoch: 5: [D loss: 0.044293] [G total loss: 19.839882] [G GAN Loss: 3.810984]\n",
            "Epoch: 6: [D loss: 0.012644] [G total loss: 20.432430] [G GAN Loss: 4.500637]\n",
            "Epoch: 7: [D loss: 0.039813] [G total loss: 17.535648] [G GAN Loss: 3.933465]\n",
            "Epoch: 7: [D loss: 0.004385] [G total loss: 17.779980] [G GAN Loss: 5.553062]\n",
            "Epoch: 8: [D loss: 0.003467] [G total loss: 23.714794] [G GAN Loss: 6.177885]\n",
            "Epoch: 9: [D loss: 0.093320] [G total loss: 25.188532] [G GAN Loss: 5.624502]\n",
            "Epoch: 10: [D loss: 0.012877] [G total loss: 20.543989] [G GAN Loss: 4.873987]\n",
            "Epoch: 11: [D loss: 0.069864] [G total loss: 20.757734] [G GAN Loss: 5.499758]\n",
            "Epoch: 12: [D loss: 0.016566] [G total loss: 13.626388] [G GAN Loss: 4.254122]\n",
            "Epoch: 13: [D loss: 0.006809] [G total loss: 18.601549] [G GAN Loss: 5.029605]\n",
            "Epoch: 14: [D loss: 0.015382] [G total loss: 16.561878] [G GAN Loss: 4.708138]\n",
            "Epoch: 14: [D loss: 0.015917] [G total loss: 20.505608] [G GAN Loss: 5.637476]\n",
            "Epoch: 15: [D loss: 0.004291] [G total loss: 20.971260] [G GAN Loss: 5.679390]\n",
            "Epoch: 16: [D loss: 0.009555] [G total loss: 17.346901] [G GAN Loss: 5.878819]\n",
            "Epoch: 17: [D loss: 0.004520] [G total loss: 18.145716] [G GAN Loss: 5.984663]\n",
            "Epoch: 18: [D loss: 0.006010] [G total loss: 15.470160] [G GAN Loss: 5.053078]\n",
            "Epoch: 19: [D loss: 0.000518] [G total loss: 15.648239] [G GAN Loss: 7.580986]\n",
            "Epoch: 20: [D loss: 0.001130] [G total loss: 15.112674] [G GAN Loss: 7.154475]\n",
            "Epoch: 21: [D loss: 0.002595] [G total loss: 13.372908] [G GAN Loss: 7.012530]\n",
            "Epoch: 21: [D loss: 0.002638] [G total loss: 12.463200] [G GAN Loss: 5.436820]\n",
            "Epoch: 22: [D loss: 0.005699] [G total loss: 13.036189] [G GAN Loss: 6.126225]\n",
            "Epoch: 23: [D loss: 0.007180] [G total loss: 12.496202] [G GAN Loss: 6.891474]\n",
            "Epoch: 24: [D loss: 0.000605] [G total loss: 13.132366] [G GAN Loss: 7.354624]\n",
            "Epoch: 25: [D loss: 0.001194] [G total loss: 12.566143] [G GAN Loss: 7.298181]\n",
            "Epoch: 26: [D loss: 0.002717] [G total loss: 10.825577] [G GAN Loss: 6.321920]\n",
            "Epoch: 27: [D loss: 0.000268] [G total loss: 12.954484] [G GAN Loss: 7.835577]\n",
            "Epoch: 28: [D loss: 0.005269] [G total loss: 11.898948] [G GAN Loss: 7.948683]\n",
            "Epoch: 28: [D loss: 0.000635] [G total loss: 13.104978] [G GAN Loss: 7.295773]\n",
            "Epoch: 29: [D loss: 0.002500] [G total loss: 10.347769] [G GAN Loss: 6.568535]\n",
            "Epoch: 30: [D loss: 0.000486] [G total loss: 11.614655] [G GAN Loss: 7.871401]\n",
            "Epoch: 31: [D loss: 0.000453] [G total loss: 13.787498] [G GAN Loss: 8.036648]\n",
            "Epoch: 32: [D loss: 0.000229] [G total loss: 10.953821] [G GAN Loss: 8.238957]\n",
            "Epoch: 33: [D loss: 0.003701] [G total loss: 11.447935] [G GAN Loss: 7.820925]\n",
            "Epoch: 34: [D loss: 0.000747] [G total loss: 11.795040] [G GAN Loss: 7.659470]\n",
            "Epoch: 35: [D loss: 0.000773] [G total loss: 11.571221] [G GAN Loss: 8.031262]\n",
            "Epoch: 35: [D loss: 0.000736] [G total loss: 11.543938] [G GAN Loss: 8.120340]\n",
            "Epoch: 36: [D loss: 0.001422] [G total loss: 10.170094] [G GAN Loss: 7.110273]\n",
            "Epoch: 37: [D loss: 0.002204] [G total loss: 10.786551] [G GAN Loss: 7.444535]\n",
            "Epoch: 38: [D loss: 0.000651] [G total loss: 11.876086] [G GAN Loss: 8.964947]\n",
            "Epoch: 39: [D loss: 0.000978] [G total loss: 10.075964] [G GAN Loss: 7.400121]\n",
            "Epoch: 40: [D loss: 0.000742] [G total loss: 9.643085] [G GAN Loss: 7.371357]\n",
            "Epoch: 41: [D loss: 0.001695] [G total loss: 10.163831] [G GAN Loss: 7.232315]\n",
            "Epoch: 42: [D loss: 0.000377] [G total loss: 11.459622] [G GAN Loss: 7.754205]\n",
            "Epoch: 42: [D loss: 0.000409] [G total loss: 11.255841] [G GAN Loss: 7.766150]\n",
            "Epoch: 43: [D loss: 0.001062] [G total loss: 11.029801] [G GAN Loss: 7.894429]\n",
            "Epoch: 44: [D loss: 0.000634] [G total loss: 10.849670] [G GAN Loss: 8.075971]\n",
            "Epoch: 45: [D loss: 0.000569] [G total loss: 12.321198] [G GAN Loss: 8.989406]\n",
            "Epoch: 46: [D loss: 0.000292] [G total loss: 11.066622] [G GAN Loss: 8.458964]\n",
            "Epoch: 47: [D loss: 0.000428] [G total loss: 10.722610] [G GAN Loss: 8.479347]\n",
            "Epoch: 48: [D loss: 0.000617] [G total loss: 10.211001] [G GAN Loss: 7.902008]\n",
            "Epoch: 49: [D loss: 0.000420] [G total loss: 9.972698] [G GAN Loss: 7.673792]\n",
            "Epoch: 49: [D loss: 0.003537] [G total loss: 8.896623] [G GAN Loss: 5.143877]\n",
            "Epoch: 50: [D loss: 0.000429] [G total loss: 10.558670] [G GAN Loss: 7.975339]\n",
            "Epoch: 51: [D loss: 0.000480] [G total loss: 11.638545] [G GAN Loss: 9.312828]\n",
            "Epoch: 52: [D loss: 0.001653] [G total loss: 11.193885] [G GAN Loss: 8.598993]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "k98dDFla1vjv",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "args.input = os.path.join(base_dir, \"images/sunflower.jpg\")\n",
        "args.output = os.path.join(base_dir, \"color_sunflower.jpg\")\n",
        "\n",
        "#G = generator().to(device)\n",
        "G = ColorGenerator().to(device)\n",
        "\n",
        "if torch.cuda.is_available():    \n",
        "    G.load_state_dict(torch.load(args.model, map_location='cuda:0'))\n",
        "else:\n",
        "    G.load_state_dict(torch.load(args.model,map_location='cpu'))\n",
        "\n",
        "def inference(G,in_path,out_path):\n",
        "    p=Image.open(in_path).convert('RGB')\n",
        "    img_yuv = rgb2yuv(p)\n",
        "    H,W,_ = img_yuv.shape\n",
        "    infimg = np.expand_dims(np.expand_dims(img_yuv[...,0], axis=0), axis=0)\n",
        "    img_variable = Variable(torch.Tensor(infimg-0.5))\n",
        "    if args.gpu>=0:\n",
        "        img_variable=img_variable.to(device)\n",
        "    res = G(img_variable, mask)\n",
        "    uv=res.cpu().detach().numpy()\n",
        "    uv[:,0,:,:] *= 0.436\n",
        "    uv[:,1,:,:] *= 0.615\n",
        "    (_,_,H1,W1) = uv.shape\n",
        "    uv = zoom(uv,(1,1,H/H1,W/W1))\n",
        "    yuv = np.concatenate([infimg,uv],axis=1)[0]\n",
        "    rgb=yuv2rgb(yuv.transpose(1,2,0))\n",
        "    cv2.imwrite(out_path,(rgb.clip(min=0,max=1)*256)[:,:,[2,1,0]])\n",
        "\n",
        "\n",
        "if not os.path.isdir(args.input):\n",
        "    inference(G,args.input,args.output)\n",
        "else:\n",
        "    if not os.path.exists(args.output):\n",
        "        os.makedirs(args.output)\n",
        "    for f in os.listdir(args.input):\n",
        "        inference(G,os.path.join(args.input,f),os.path.join(args.output,f))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "m25N8jbH9KyQ",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "if args.test_image is not None:\n",
        "    test_res = G(test_var, mask)\n",
        "    uv=test_res.cpu().detach().numpy()\n",
        "    uv[:,0,:,:] *= 0.436\n",
        "    uv[:,1,:,:] *= 0.615\n",
        "    #print(test_var.shape, mask.shape, test_res.shape, uv.shape, test_inf.shape)\n",
        "    test_yuv = np.concatenate([test_inf,uv],axis=1).reshape(3,SIZE,SIZE)\n",
        "    test_rgb = yuv2rgb(test_yuv.transpose(1,2,0))\n",
        "    cv2.imwrite(os.path.join(args.checkpoint_location,'test_init.jpg'),(test_rgb.clip(min=0,max=1)*SIZE)[:,:,[2,1,0]])\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}