Council_gan_basic_inference.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Council_gan_basic_inference.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "Ln29iemGiF8e",
        "D67aq1FQjkdl",
        "xD4r65toihTF",
        "WUTk5ZhGGmWN",
        "6MB32ZWck43o"
      ],
      "authorship_tag": "ABX9TyNcAFgU2nU/709EPhmFEi8J",
      "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/Onr/Council-GAN/blob/master/Council_gan_basic_inference.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QzugSnW0irGG",
        "colab_type": "text"
      },
      "source": [
        "# installes\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "t3NM-RTJiqd3",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 170
        },
        "outputId": "e53b8fd3-2519-48c8-c331-f02291e0dd2c"
      },
      "source": [
        "!pip install torchfile "
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Collecting torchfile\n",
            "  Downloading https://files.pythonhosted.org/packages/91/af/5b305f86f2d218091af657ddb53f984ecbd9518ca9fe8ef4103a007252c9/torchfile-0.1.0.tar.gz\n",
            "Building wheels for collected packages: torchfile\n",
            "  Building wheel for torchfile (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for torchfile: filename=torchfile-0.1.0-cp36-none-any.whl size=5712 sha256=a8b73ed4eff5850c52cc9b0a6b591e6c948c8764df6970bed026e26ade659a12\n",
            "  Stored in directory: /root/.cache/pip/wheels/b1/c3/d6/9a1cc8f3a99a0fc1124cae20153f36af59a6e683daca0a0814\n",
            "Successfully built torchfile\n",
            "Installing collected packages: torchfile\n",
            "Successfully installed torchfile-0.1.0\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ln29iemGiF8e",
        "colab_type": "text"
      },
      "source": [
        "#data.py\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZE5YWsW-h85Z",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\"\"\"\n",
        "Copyright (C) 2018 NVIDIA Corporation.  All rights reserved.\n",
        "Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n",
        "\"\"\"\n",
        "import torch.utils.data as data\n",
        "import os.path\n",
        "import torch\n",
        "import numpy as np\n",
        "\n",
        "def default_loader(path):\n",
        "    return Image.open(path).convert('RGB')\n",
        "\n",
        "\n",
        "def default_flist_reader(flist):\n",
        "    \"\"\"\n",
        "    flist format: impath label\\nimpath label\\n ...(same to caffe's filelist)\n",
        "    \"\"\"\n",
        "    imlist = []\n",
        "    with open(flist, 'r') as rf:\n",
        "        for line in rf.readlines():\n",
        "            impath = line.strip()\n",
        "            imlist.append(impath)\n",
        "\n",
        "    return imlist\n",
        "\n",
        "\n",
        "class ImageFilelist(data.Dataset):\n",
        "    def __init__(self, root, flist, transform=None,\n",
        "                 flist_reader=default_flist_reader, loader=default_loader):\n",
        "        self.root = root\n",
        "        self.imlist = flist_reader(flist)\n",
        "        self.transform = transform\n",
        "        self.loader = loader\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        impath = self.imlist[index]\n",
        "        img = self.loader(os.path.join(self.root, impath))\n",
        "        if self.transform is not None:\n",
        "            img = self.transform(img)\n",
        "\n",
        "        return img\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.imlist)\n",
        "\n",
        "\n",
        "class ImageLabelFilelist(data.Dataset):\n",
        "    def __init__(self, root, flist, transform=None,\n",
        "                 flist_reader=default_flist_reader, loader=default_loader):\n",
        "        self.root = root\n",
        "        self.imlist = flist_reader(os.path.join(self.root, flist))\n",
        "        self.transform = transform\n",
        "        self.loader = loader\n",
        "        self.classes = sorted(list(set([path.split('/')[0] for path in self.imlist])))\n",
        "        self.class_to_idx = {self.classes[i]: i for i in range(len(self.classes))}\n",
        "        self.imgs = [(impath, self.class_to_idx[impath.split('/')[0]]) for impath in self.imlist]\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        impath, label = self.imgs[index]\n",
        "        img = self.loader(os.path.join(self.root, impath))\n",
        "        if self.transform is not None:\n",
        "            img = self.transform(img)\n",
        "        return img, label\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.imgs)\n",
        "\n",
        "###############################################################################\n",
        "# Code from\n",
        "# https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py\n",
        "# Modified the original code so that it also loads images from the current\n",
        "# directory as well as the subdirectories\n",
        "###############################################################################\n",
        "\n",
        "import torch.utils.data as data\n",
        "\n",
        "from PIL import Image\n",
        "import os\n",
        "import os.path\n",
        "\n",
        "IMG_EXTENSIONS = [\n",
        "    '.jpg', '.JPG', '.jpeg', '.JPEG',\n",
        "    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',\n",
        "]\n",
        "\n",
        "\n",
        "def is_image_file(filename):\n",
        "    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)\n",
        "\n",
        "\n",
        "def make_dataset(dir):\n",
        "    images = []\n",
        "    assert os.path.isdir(dir), '%s is not a valid directory' % dir\n",
        "\n",
        "    for root, _, fnames in sorted(os.walk(dir)):\n",
        "        for fname in fnames:\n",
        "            if is_image_file(fname):\n",
        "                path = os.path.join(root, fname)\n",
        "                images.append(path)\n",
        "            elif fname.endswith('.npy'):\n",
        "                path = os.path.join(root, fname)\n",
        "                images.append(path)\n",
        "\n",
        "    return images\n",
        "\n",
        "\n",
        "class ImageFolder(data.Dataset):\n",
        "\n",
        "    def __init__(self, root, transform=None, return_paths=False,\n",
        "                 loader=default_loader):\n",
        "        imgs = sorted(make_dataset(root))\n",
        "        if len(imgs) == 0:\n",
        "            raise(RuntimeError(\"Found 0 images in: \" + root + \"\\n\"\n",
        "                               \"Supported image extensions are: \" +\n",
        "                               \",\".join(IMG_EXTENSIONS)))\n",
        "\n",
        "        self.root = root\n",
        "        self.imgs = imgs\n",
        "        self.transform = transform\n",
        "        self.return_paths = return_paths\n",
        "        self.loader = loader\n",
        "\n",
        "        # TODO tmp\n",
        "        self.mean = 0\n",
        "        self.std = None\n",
        "        self.mean_pow2 = 0\n",
        "\n",
        "        self.num_of_sample = 0\n",
        "\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        path = self.imgs[index]\n",
        "        if not path.endswith('.npy'):\n",
        "            try:\n",
        "                img = self.loader(path)\n",
        "            except Exception as e:\n",
        "                print(str(e))\n",
        "                del self.imgs[index]\n",
        "                self.__getitem__(self, index)\n",
        "            if self.transform is not None:\n",
        "                img = self.transform(img)\n",
        "        else:  # numpy data input # for brats dataset  # TODO add transforms\n",
        "            img = torch.from_numpy(np.load(path))\n",
        "            img = img.transpose(dim0=2, dim1=0)\n",
        "            img = img.transpose(dim0=1, dim1=2)\n",
        "            img = img.to(dtype=torch.float)\n",
        "            from torchvision import transforms\n",
        "            import torch.nn.functional as F\n",
        "            mean_vec = torch.mean(img, dim=(1, 2))\n",
        "            std_vec = torch.std(img, dim=(1, 2))\n",
        "            img = transforms.Normalize(mean=mean_vec, std=std_vec)(img)\n",
        "            size = [elem for elem in self.transform.transforms if type(elem) == transforms.transforms.Resize][0].size\n",
        "            img = F.interpolate(input=torch.unsqueeze(img,0), size=size)\n",
        "            img = torch.squeeze(img)\n",
        "            # img = img[:-1,:,:]\n",
        "\n",
        "            # dim0 Flair\n",
        "            # dim1 T1\n",
        "            # dim2 T1ce\n",
        "            # dim3 T2\n",
        "            # dim4 Seg\n",
        "\n",
        "        if self.return_paths:\n",
        "            return img, path\n",
        "        else:\n",
        "            return img\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.imgs)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "class ImageFolder_with_subfolders(data.Dataset):\n",
        "# choose the data from the subfolder acording to probability\n",
        "    # TODO complete usege\n",
        "    def __init__(self, root1, root2, ratio_1_to_2, transform=None, return_paths=False,\n",
        "                 loader=default_loader):\n",
        "        self.ratio_1_to_2 = ratio_1_to_2\n",
        "        imgs1 = sorted(make_dataset(root1))\n",
        "        imgs2 = sorted(make_dataset(root2))\n",
        "        if len(imgs1) == 0:\n",
        "            raise(RuntimeError(\"Found 0 images in: \" + root1 + \"\\n\"\n",
        "                               \"Supported image extensions are: \" +\n",
        "                               \",\".join(IMG_EXTENSIONS)))\n",
        "        if len(imgs2) == 0:\n",
        "            raise(RuntimeError(\"Found 0 images in: \" + root2 + \"\\n\"\n",
        "                               \"Supported image extensions are: \" +\n",
        "                               \",\".join(IMG_EXTENSIONS)))\n",
        "\n",
        "        self.root1 = root1\n",
        "        self.root2 = root2\n",
        "        self.imgs1 = imgs1\n",
        "        self.imgs2 = imgs2\n",
        "        self.transform = transform\n",
        "        self.return_paths = return_paths\n",
        "        self.loader = loader\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        rand_res = torch.rand(1)\n",
        "        path = self.imgs1[index % len(self.imgs1)] if rand_res < self.ratio_1_to_2 else self.imgs2[index % len(self.imgs2)]\n",
        "        img = self.loader(path)\n",
        "        if self.transform is not None:\n",
        "            img = self.transform(img)\n",
        "        if self.return_paths:\n",
        "            return img, path\n",
        "        else:\n",
        "            return img\n",
        "\n",
        "    def __len__(self):\n",
        "        return max(len(self.imgs1), len(self.imgs2))\n"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D67aq1FQjkdl",
        "colab_type": "text"
      },
      "source": [
        "# networks.py"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6AvZuc4ZjpMp",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\"\"\"\n",
        "Copyright (C) 2018 NVIDIA Corporation.  All rights reserved.\n",
        "Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n",
        "\"\"\"\n",
        "from torch import nn\n",
        "from torch.autograd import Variable\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "try:\n",
        "    from itertools import izip as zip\n",
        "except ImportError: # will be 3.x series\n",
        "    pass\n",
        "\n",
        "##################################################################################\n",
        "# Discriminator\n",
        "##################################################################################\n",
        "class MsImageDis(nn.Module):\n",
        "    # Multi-scale discriminator architecture\n",
        "    def __init__(self, input_dim, params, cuda_device='cuda:0'):\n",
        "        super(MsImageDis, self).__init__()\n",
        "        self.prev_real_input = None\n",
        "        self.n_layer = params['n_layer']\n",
        "        self.gan_type = params['gan_type']\n",
        "        self.dim = params['dim']\n",
        "        self.norm = params['norm']\n",
        "        self.activ = params['activ']\n",
        "        self.num_scales = params['num_scales']\n",
        "        self.pad_type = params['pad_type']\n",
        "        self.cuda_device = cuda_device\n",
        "\n",
        "        self.input_dim = input_dim\n",
        "        self.downsample = nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False)\n",
        "        self.cnns = nn.ModuleList()\n",
        "        for _ in range(self.num_scales):\n",
        "            self.cnns.append(self._make_net().cuda(self.cuda_device))\n",
        "\n",
        "    def _make_net(self):\n",
        "        dim = self.dim\n",
        "        cnn_x = []\n",
        "        cnn_x += [Conv2dBlock(self.input_dim, dim, 4, 2, 1, norm='none', activation=self.activ, pad_type=self.pad_type)]\n",
        "        for i in range(self.n_layer - 1):\n",
        "            cnn_x += [Conv2dBlock(dim, dim * 2, 4, 2, 1, norm=self.norm, activation=self.activ, pad_type=self.pad_type)]\n",
        "            dim *= 2\n",
        "        cnn_x += [nn.Conv2d(dim, 1, 1, 1, 0)]\n",
        "        cnn_x = nn.Sequential(*cnn_x)\n",
        "        return cnn_x\n",
        "\n",
        "    def forward(self, x):\n",
        "        outputs = []\n",
        "        x = x.cuda(self.cuda_device)\n",
        "        for model in self.cnns:\n",
        "            outputs.append(model(x))\n",
        "            x = self.downsample(x)\n",
        "        return outputs\n",
        "\n",
        "    def calc_dis_loss(self, input_fake, input_real):\n",
        "        # calculate the loss to train D\n",
        "        outs0 = self.forward(input_fake)\n",
        "        outs1 = self.forward(input_real)\n",
        "        loss = 0\n",
        "\n",
        "        for it, (out0, out1) in enumerate(zip(outs0, outs1)):\n",
        "            if self.gan_type == 'lsgan':\n",
        "                loss += torch.mean((out0 - 0)**2) + torch.mean((out1 - 1)**2)\n",
        "            elif self.gan_type == 'nsgan':\n",
        "                all0 = Variable(torch.zeros_like(out0.data).cuda(self.cuda_device), requires_grad=False)\n",
        "                all1 = Variable(torch.ones_like(out1.data).cuda(self.cuda_device), requires_grad=False)\n",
        "                loss += torch.mean(F.binary_cross_entropy(F.sigmoid(out0), all0) +\n",
        "                                   F.binary_cross_entropy(F.sigmoid(out1), all1))\n",
        "            elif self.gan_type == 'RelativisticAverageHingeGAN':\n",
        "                self.prev_real_input = input_real  # save it for the gen train later\n",
        "                # difference between real and fake:\n",
        "                r_f_diff = out1 - torch.mean(out0, dim=0, keepdim=True).repeat(10, 1, 1, 1)\n",
        "                # difference between fake and real samples\n",
        "                f_r_diff = out0 - torch.mean(out1, dim=0, keepdim=True).repeat(10, 1, 1, 1)\n",
        "                # return the loss\n",
        "                loss += (torch.mean(torch.nn.ReLU()(1 - r_f_diff))\n",
        "                         + torch.mean(torch.nn.ReLU()(1 + f_r_diff)))\n",
        "\n",
        "            else:\n",
        "                assert 0, \"Unsupported GAN type: {}\".format(self.gan_type)\n",
        "        return loss\n",
        "\n",
        "    def calc_gen_loss(self, input_fake, input_real=None):\n",
        "        # calculate the loss to train G\n",
        "        outs0 = self.forward(input_fake)\n",
        "        loss = 0\n",
        "        for it, (out0) in enumerate(outs0):\n",
        "            if self.gan_type == 'lsgan':\n",
        "                loss += torch.mean((out0 - 1)**2)  # LSGAN\n",
        "            elif self.gan_type == 'nsgan':\n",
        "                all1 = Variable(torch.ones_like(out0.data).cuda(self.cuda_device), requires_grad=False)\n",
        "                loss += torch.mean(F.binary_cross_entropy(F.sigmoid(out0), all1))\n",
        "            elif self.gan_type == 'RelativisticAverageHingeGAN':\n",
        "                if input_real is not None:\n",
        "                    outs1 = self.forward(input_real)\n",
        "                elif self.prev_real_input is not None:\n",
        "                    outs1 = self.forward(self.prev_real_input)\n",
        "                else:\n",
        "                    assert 0, \"try using cal_gan_loss with RelativisticAverageHingeGAN but did not provid input_real\"\n",
        "                out1 = outs1[it]\n",
        "                # difference between real and fake:\n",
        "                r_f_diff = out1 - torch.mean(out0, dim=0, keepdim=True).repeat(10, 1, 1, 1)\n",
        "                # difference between fake and real samples\n",
        "                f_r_diff = out0 - torch.mean(out1, dim=0, keepdim=True).repeat(10, 1, 1, 1)\n",
        "                loss = torch.mean(torch.nn.ReLU()(1 + r_f_diff)) + torch.mean(torch.nn.ReLU()(1 - f_r_diff))\n",
        "\n",
        "            else:\n",
        "                assert 0, \"Unsupported GAN type: {}\".format(self.gan_type)\n",
        "        return loss\n",
        "\n",
        "\n",
        "##################################################################################\n",
        "# Council Discriminator\n",
        "##################################################################################\n",
        "class MsImageDisCouncil(nn.Module):\n",
        "    # Multi-scale discriminator architecture\n",
        "    def __init__(self, input_dim, params, cuda_device='cuda:0'):\n",
        "        super(MsImageDisCouncil, self).__init__()\n",
        "        self.n_layer = params['n_layer']\n",
        "        self.gan_type = params['gan_type']\n",
        "        self.dim = params['dim']\n",
        "        self.norm = params['norm']\n",
        "        self.activ = params['activ']\n",
        "        self.num_scales = params['num_scales']\n",
        "        self.pad_type = params['pad_type']\n",
        "        self.cuda_device = cuda_device\n",
        "        self.input_dim = input_dim\n",
        "        self.downsample = nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False)\n",
        "        self.cnns = nn.ModuleList()\n",
        "        for _ in range(self.num_scales):\n",
        "            self.cnns.append(self._make_net())\n",
        "\n",
        "    def _make_net(self):\n",
        "        dim = self.dim\n",
        "        cnn_x = []\n",
        "        # cnn_x += [Conv2dBlock(input_dim=2 * self.input_dim, output_dim=dim, kernel_size=4, stride=2, padding=1, norm='none', activation=self.activ, pad_type=self.pad_type)] # original\n",
        "        cnn_x += [Conv2dBlock(input_dim=2 * self.input_dim, output_dim=dim, kernel_size=3, stride=1, padding=1, norm='none', activation=self.activ, pad_type=self.pad_type)] # ON\n",
        "        for i in range(self.n_layer - 1):\n",
        "            cnn_x += [Conv2dBlock(dim, dim * 2, 4, 2, 1, norm=self.norm, activation=self.activ, pad_type=self.pad_type)]\n",
        "            dim *= 2\n",
        "        cnn_x += [nn.Conv2d(dim, dim, 1, 1, 0)]\n",
        "        cnn_x += [nn.Conv2d(dim, 1, 1, 1, 0)]\n",
        "        cnn_x = nn.Sequential(*cnn_x).cuda(self.cuda_device)\n",
        "        return cnn_x\n",
        "\n",
        "    def forward(self, x, x_input):\n",
        "        x = x.cuda(self.cuda_device)\n",
        "        x_input = x_input.cuda(self.cuda_device)\n",
        "        outputs = []\n",
        "        for model in self.cnns:\n",
        "            model_input = torch.cat((x, x_input), 1)\n",
        "            outputs.append(model(model_input))\n",
        "            x = self.downsample(x)\n",
        "            x_input = self.downsample(x_input)\n",
        "        return outputs\n",
        "\n",
        "    def calc_dis_loss(self, input_fake, input_real, input):\n",
        "        # calculate the loss to train D\n",
        "        outs0 = self.forward(input_fake, input)\n",
        "        outs1 = self.forward(input_real, input)\n",
        "        loss = 0\n",
        "\n",
        "        for it, (out0, out1) in enumerate(zip(outs0, outs1)):\n",
        "            if self.gan_type == 'lsgan':\n",
        "                loss += torch.mean((out0 - 0)**2) + torch.mean((out1 - 1)**2)\n",
        "            elif self.gan_type == 'nsgan':\n",
        "                all0 = Variable(torch.zeros_like(out0.data).cuda(self.cuda_device), requires_grad=False)\n",
        "                all1 = Variable(torch.ones_like(out1.data).cuda(self.cuda_device), requires_grad=False)\n",
        "                loss += torch.mean(F.binary_cross_entropy(F.sigmoid(out0), all0) +\n",
        "                                   F.binary_cross_entropy(F.sigmoid(out1), all1))\n",
        "\n",
        "            elif self.gan_type == 'RelativisticAverageHingeGAN':\n",
        "                self.prev_real_input = input_real  # save it for the gen train later\n",
        "                self.prev_input = input  # save it for the gen train later\n",
        "                # difference between real and fake:\n",
        "                r_f_diff = out1 - torch.mean(out0, dim=0, keepdim=True).repeat(10, 1, 1, 1)\n",
        "                # difference between fake and real samples\n",
        "                f_r_diff = out0 - torch.mean(out1, dim=0, keepdim=True).repeat(10, 1, 1, 1)\n",
        "                loss = torch.mean(torch.nn.ReLU()(1 + r_f_diff)) + torch.mean(torch.nn.ReLU()(1 - f_r_diff))\n",
        "                # return the loss\n",
        "                loss += (torch.mean(torch.nn.ReLU()(1 - r_f_diff))\n",
        "                         + torch.mean(torch.nn.ReLU()(1 + f_r_diff)))\n",
        "            else:\n",
        "                assert 0, \"Unsupported GAN type: {}\".format(self.gan_type)\n",
        "        return loss\n",
        "\n",
        "    def calc_gen_loss(self, input_fake, input, input_real=None):\n",
        "        # calculate the loss to train G\n",
        "        outs0 = self.forward(input_fake, input)\n",
        "        loss = 0\n",
        "        for it, (out0) in enumerate(outs0):\n",
        "            if self.gan_type == 'lsgan':\n",
        "                loss += torch.mean((out0 - 1)**2) # LSGAN\n",
        "            elif self.gan_type == 'nsgan':\n",
        "                all1 = Variable(torch.ones_like(out0.data).cuda(self.cuda_device), requires_grad=False)\n",
        "                loss += torch.mean(F.binary_cross_entropy(F.sigmoid(out0), all1))\n",
        "            elif self.gan_type == 'RelativisticAverageHingeGAN':\n",
        "                if input_real is not None:\n",
        "                    outs1 = self.forward(input_real)\n",
        "                elif self.prev_real_input is not None:\n",
        "                    outs1 = self.forward(self.prev_real_input, self.prev_input)\n",
        "                else:\n",
        "                    assert 0, \"try using cal_gan_loss with RelativisticAverageHingeGAN but did not provid input_real\"\n",
        "\n",
        "                out1 = outs1[it]\n",
        "                # difference between real and fake:\n",
        "                r_f_diff = out1 - torch.mean(out0, dim=0, keepdim=True).repeat(10, 1, 1, 1)\n",
        "                # difference between fake and real samples\n",
        "                f_r_diff = out0 - torch.mean(out1, dim=0, keepdim=True).repeat(10, 1, 1, 1)\n",
        "                loss = torch.mean(torch.nn.ReLU()(1 + r_f_diff)) + torch.mean(torch.nn.ReLU()(1 - f_r_diff))\n",
        "                loss = torch.mean(torch.nn.ReLU()(1 + r_f_diff)) + torch.mean(torch.nn.ReLU()(1 - f_r_diff))\n",
        "            else:\n",
        "                assert 0, \"Unsupported GAN type: {}\".format(self.gan_type)\n",
        "        return loss\n",
        "\n",
        "\n",
        "\n",
        "##################################################################################\n",
        "# Generator\n",
        "##################################################################################\n",
        "\n",
        "class AdaINGen(nn.Module):\n",
        "    # AdaIN auto-encoder architecture\n",
        "    def __init__(self, input_dim, params, cuda_device='cuda:0'):\n",
        "        super(AdaINGen, self).__init__()\n",
        "        dim = params['dim']\n",
        "        style_dim = params['style_dim']\n",
        "        self.n_downsample = params['n_downsample']\n",
        "        n_res = params['n_res']\n",
        "        self.activ = params['activ']\n",
        "        pad_type = params['pad_type']\n",
        "        mlp_dim = params['mlp_dim']\n",
        "        self.do_my_style = params['do_my_style']\n",
        "        self.cuda_device = cuda_device\n",
        "\n",
        "        # style encoder\n",
        "        self.enc_style = StyleEncoder(4, input_dim, dim, style_dim, norm='none', activ=self.activ, pad_type=pad_type).cuda(self.cuda_device)\n",
        "\n",
        "        # content encoder\n",
        "        self.enc_content = ContentEncoder(self.n_downsample, n_res, input_dim, dim, 'in', self.activ, pad_type=pad_type).cuda(self.cuda_device)\n",
        "\n",
        "        if self.do_my_style:\n",
        "            self.dec = Decoder_V2_atten(n_upsample=self.n_downsample, n_res=n_res, dim=self.enc_content.output_dim + style_dim, output_dim=input_dim, res_norm='in', activ=self.activ, pad_type=pad_type, num_of_mask_dim_to_add=params['num_of_mask_dim_to_add']).cuda(self.cuda_device)\n",
        "        else:\n",
        "            self.dec = Decoder_V2_atten(self.n_downsample, n_res, self.enc_content.output_dim, input_dim, res_norm='adain', activ=self.activ, pad_type=pad_type, num_of_mask_dim_to_add=params['num_of_mask_dim_to_add']).cuda(self.cuda_device)\n",
        "\n",
        "        # MLP to generate AdaIN parameters or adding sytle my way\n",
        "        if self.do_my_style:\n",
        "            self.mlp = MLP(input_dim=style_dim, output_dim=style_dim, dim=mlp_dim, n_blk=3, norm='none',\n",
        "                           activ=self.activ).cuda(self.cuda_device)\n",
        "        else:\n",
        "            self.mlp = MLP(input_dim=style_dim, output_dim=self.get_num_adain_params(self.dec), dim=mlp_dim, n_blk=3,\n",
        "                           norm='none', activ=self.activ).cuda(self.cuda_device)\n",
        "\n",
        "    def forward(self, images, return_mask=False):\n",
        "        # reconstruct an image\n",
        "        images = images.cuda(self.cuda_device)\n",
        "        content, style_fake = self.encode(images)\n",
        "        if return_mask:\n",
        "            images_recon, mask = self.decode(content=content, style=style_fake, images=images, return_mask=return_mask)\n",
        "            return images_recon, mask\n",
        "        else:\n",
        "            images_recon = self.decode(content=content, style=style_fake, images=images, return_mask=return_mask)\n",
        "            return images_recon\n",
        "\n",
        "    def forward(self, images, style, return_mask=False):\n",
        "        # reconstruct an image\n",
        "        content, _ = self.encode(images)\n",
        "        if return_mask:\n",
        "            images_recon, mask = self.decode(content=content, style=style, images=images, return_mask=return_mask)\n",
        "            return images_recon, mask\n",
        "        else:\n",
        "            images_recon = self.decode(content=content, style=style, images=images, return_mask=return_mask)\n",
        "            return images_recon\n",
        "\n",
        "\n",
        "    def encode(self, images):\n",
        "        images = images.cuda(self.cuda_device)\n",
        "        # encode an image to its content and style codes\n",
        "        style_fake = self.enc_style(images)\n",
        "        content = self.enc_content(images)\n",
        "        return content, style_fake\n",
        "\n",
        "    def decode(self, content, style, images, return_mask=False):\n",
        "        content, style, images = content.cuda(self.cuda_device), style.cuda(self.cuda_device), images.cuda(self.cuda_device)\n",
        "        # decode content and style codes to an image\n",
        "        if self.do_my_style:\n",
        "            style_to_add = self.mlp(style)\n",
        "            style_to_add = style_to_add.repeat(content.shape[2], content.shape[3], 1, 1)\n",
        "            style_to_add = style_to_add.transpose(0, 2).transpose(3, 1)\n",
        "            content = torch.cat((content, style_to_add), 1)\n",
        "        else:\n",
        "            adain_params = self.mlp(style)\n",
        "            self.assign_adain_params(adain_params, self.dec)\n",
        "        if return_mask:\n",
        "            images, mask = self.dec(content, images, return_mask)\n",
        "            return images, mask\n",
        "        else:\n",
        "            images = self.dec(content, images)\n",
        "            return images\n",
        "\n",
        "    def assign_adain_params(self, adain_params, model):\n",
        "        # assign the adain_params to the AdaIN layers in model\n",
        "        for m in model.modules():\n",
        "            if m.__class__.__name__ == \"AdaptiveInstanceNorm2d\":\n",
        "                mean = adain_params[:, :m.num_features]\n",
        "                std = adain_params[:, m.num_features:2*m.num_features]\n",
        "                m.bias = mean.contiguous().view(-1)\n",
        "                m.weight = std.contiguous().view(-1)\n",
        "                if adain_params.size(1) > 2*m.num_features:\n",
        "                    adain_params = adain_params[:, 2*m.num_features:]\n",
        "\n",
        "    def get_num_adain_params(self, model):\n",
        "        # return the number of AdaIN parameters needed by the model\n",
        "        num_adain_params = 0\n",
        "        for m in model.modules():\n",
        "            if m.__class__.__name__ == \"AdaptiveInstanceNorm2d\":\n",
        "                num_adain_params += 2*m.num_features\n",
        "        return num_adain_params\n",
        "\n",
        "\n",
        "##################################################################################\n",
        "# Encoder and Decoders\n",
        "##################################################################################\n",
        "\n",
        "class StyleEncoder(nn.Module):\n",
        "    def __init__(self, n_downsample, input_dim, dim, style_dim, norm, activ, pad_type):\n",
        "        super(StyleEncoder, self).__init__()\n",
        "        self.model = []\n",
        "        self.model += [Conv2dBlock(input_dim, dim, 7, 1, 3, norm=norm, activation=activ, pad_type=pad_type)]\n",
        "        for i in range(2):\n",
        "            self.model += [Conv2dBlock(dim, 2 * dim, 4, 2, 1, norm=norm, activation=activ, pad_type=pad_type)]\n",
        "            dim *= 2\n",
        "        for i in range(n_downsample - 2):\n",
        "            self.model += [Conv2dBlock(dim, dim, 4, 2, 1, norm=norm, activation=activ, pad_type=pad_type)]\n",
        "        self.model += [nn.AdaptiveAvgPool2d(1)] # global average pooling\n",
        "        self.model += [nn.Conv2d(dim, style_dim, 1, 1, 0)]\n",
        "        self.model = nn.Sequential(*self.model)\n",
        "        self.output_dim = dim\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.model(x)\n",
        "\n",
        "class ContentEncoder(nn.Module):\n",
        "    def __init__(self, n_downsample, n_res, input_dim, dim, norm, activ, pad_type):\n",
        "        super(ContentEncoder, self).__init__()\n",
        "        self.model = []\n",
        "        self.model += [Conv2dBlock(input_dim, dim, 7, 1, 3, norm=norm, activation=activ, pad_type=pad_type)] # ORIGINAL\n",
        "        for i in range(n_downsample):\n",
        "            self.model += [Conv2dBlock(dim, 2 * dim, 4, 2, 1, norm=norm, activation=activ, pad_type=pad_type)]\n",
        "            dim *= 2\n",
        "        # residual blocks\n",
        "        self.model += [ResBlocks(n_res, dim, norm=norm, activation=activ, pad_type=pad_type)]\n",
        "        self.model = nn.Sequential(*self.model)\n",
        "        self.output_dim = dim\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.model(x)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "class Decoder_V2_atten(nn.Module):\n",
        "    def __init__(self, n_upsample, n_res, dim, output_dim, res_norm='adain', activ='relu', pad_type='zero', num_of_mask_dim_to_add=1):\n",
        "        super(Decoder_V2_atten, self).__init__()\n",
        "        self.num_of_mask_dim_to_add = num_of_mask_dim_to_add # 3  # 2\n",
        "        self.model = []\n",
        "        self.output_dim = output_dim\n",
        "        self.mask_s = []\n",
        "        # AdaIN residual blocks\n",
        "        self.model += [ResBlocks(n_res, dim, res_norm, activ, pad_type=pad_type)]\n",
        "        # upsampling blocks\n",
        "        for i in range(n_upsample):\n",
        "            self.model += [nn.Upsample(scale_factor=2)]\n",
        "            self.model += [Conv2dBlock(input_dim=dim, output_dim=dim // 2, kernel_size=3, stride=1, padding=1, norm='adain',\n",
        "                                       activation=activ, pad_type=pad_type)]\n",
        "            dim //= 2\n",
        "            self.model += [Conv2dBlock(input_dim=dim, output_dim=dim, kernel_size=3, stride=1, padding=1, norm='adain',\n",
        "                                       activation=activ, pad_type=pad_type)]\n",
        "\n",
        "\n",
        "        self.model += [Conv2dBlock(input_dim=dim, output_dim=dim, kernel_size=1, stride=1, padding=0, norm='none', activation=activ, pad_type=pad_type)]\n",
        "        self.model += [Conv2dBlock(input_dim=dim, output_dim=dim, kernel_size=1, stride=1, padding=0, norm='none', activation=activ, pad_type=pad_type)]\n",
        "        self.model += [Conv2dBlock(input_dim=dim, output_dim=(output_dim*self.num_of_mask_dim_to_add+self.num_of_mask_dim_to_add), kernel_size=1, stride=1, padding=0, norm='none', activation='tanh', pad_type=pad_type)]\n",
        "        self.model = nn.Sequential(*self.model)\n",
        "\n",
        "    def forward(self, x, im_in, return_mask=False):\n",
        "        new_x = self.model(x)\n",
        "        self.mask_s = ((torch.tanh(10 * new_x[:, (-1 * self.num_of_mask_dim_to_add):, :, :]) + 1) / 2)\n",
        "        new_im = im_in\n",
        "        curr_ind = 0\n",
        "        for k in range(self.num_of_mask_dim_to_add):\n",
        "            new_im_o = new_x[:, curr_ind:self.output_dim * (k + 1), :, :] ##\n",
        "            curr_ind = self.output_dim * (k + 1)  ##\n",
        "            mask = self.mask_s[:, k, :, :].unsqueeze(1).repeat(1, new_im.shape[1], 1, 1)\n",
        "            new_im = (1 - mask) * new_im + mask * new_im_o\n",
        "\n",
        "        if return_mask:\n",
        "            if self.mask_s.shape[1] != 3:\n",
        "                mask_s_prePixTot = torch.sum(self.mask_s, 1).unsqueeze(1).repeat(1, 3, 1, 1) / self.mask_s.shape[1]\n",
        "                self.mask_s = mask_s_prePixTot\n",
        "            return new_im, self.mask_s\n",
        "\n",
        "        return new_im\n",
        "\n",
        "\n",
        "##################################################################################\n",
        "# Sequential Models\n",
        "##################################################################################\n",
        "class ResBlocks(nn.Module):\n",
        "    def __init__(self, num_blocks, dim, norm='in', activation='relu', pad_type='zero'):\n",
        "        super(ResBlocks, self).__init__()\n",
        "        self.model = []\n",
        "        for i in range(num_blocks):\n",
        "            self.model += [ResBlock(dim, norm=norm, activation=activation, pad_type=pad_type)]\n",
        "        self.model = nn.Sequential(*self.model)\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.model(x)\n",
        "\n",
        "class MLP(nn.Module):\n",
        "    def __init__(self, input_dim, output_dim, dim, n_blk, norm='none', activ='relu'):\n",
        "        super(MLP, self).__init__()\n",
        "        self.model = []\n",
        "        self.model += [LinearBlock(input_dim, dim, norm=norm, activation=activ)]\n",
        "        for i in range(n_blk - 2):\n",
        "            self.model += [LinearBlock(dim, dim, norm=norm, activation=activ)]\n",
        "        self.model += [LinearBlock(dim, output_dim, norm='none', activation='none')]  # no output activations\n",
        "        self.model = nn.Sequential(*self.model)\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.model(x.view(x.size(0), -1))\n",
        "\n",
        "##################################################################################\n",
        "# Basic Blocks\n",
        "##################################################################################\n",
        "class ResBlock(nn.Module):\n",
        "    def __init__(self, dim, norm='in', activation='relu', pad_type='zero'):\n",
        "        super(ResBlock, self).__init__()\n",
        "\n",
        "        model = []\n",
        "        model += [Conv2dBlock(dim ,dim, 3, 1, 1, norm=norm, activation=activation, pad_type=pad_type)]\n",
        "        model += [Conv2dBlock(dim ,dim, 3, 1, 1, norm=norm, activation='none', pad_type=pad_type)]\n",
        "        self.model = nn.Sequential(*model)\n",
        "\n",
        "    def forward(self, x):\n",
        "        residual = x\n",
        "        out = self.model(x)\n",
        "        out += residual\n",
        "        return out\n",
        "\n",
        "class Conv2dBlock(nn.Module):\n",
        "    def __init__(self, input_dim, output_dim, kernel_size, stride,\n",
        "                 padding=0, norm='none', activation='relu', pad_type='zero'):\n",
        "        super(Conv2dBlock, self).__init__()\n",
        "        self.use_bias = True\n",
        "        # initialize padding\n",
        "        if pad_type == 'reflect':\n",
        "            self.pad = nn.ReflectionPad2d(padding)\n",
        "        elif pad_type == 'replicate':\n",
        "            self.pad = nn.ReplicationPad2d(padding)\n",
        "        elif pad_type == 'zero':\n",
        "            self.pad = nn.ZeroPad2d(padding)\n",
        "        else:\n",
        "            assert 0, \"Unsupported padding type: {}\".format(pad_type)\n",
        "\n",
        "        # initialize normalization\n",
        "        norm_dim = output_dim\n",
        "        if norm == 'bn':\n",
        "            self.norm = nn.BatchNorm2d(norm_dim)\n",
        "        elif norm == 'in':\n",
        "            self.norm = nn.InstanceNorm2d(norm_dim)\n",
        "        elif norm == 'ln':\n",
        "            self.norm = LayerNorm(norm_dim)\n",
        "        elif norm == 'adain':\n",
        "            self.norm = AdaptiveInstanceNorm2d(norm_dim)\n",
        "        elif norm == 'none' or norm == 'sn':\n",
        "            self.norm = None\n",
        "        else:\n",
        "            assert 0, \"Unsupported normalization: {}\".format(norm)\n",
        "\n",
        "        # initialize activation\n",
        "        if activation == 'relu':\n",
        "            self.activation = nn.ReLU(inplace=True)\n",
        "        elif activation == 'lrelu':\n",
        "            self.activation = nn.LeakyReLU(0.2, inplace=True)\n",
        "        elif activation == 'prelu':\n",
        "            self.activation = nn.PReLU()\n",
        "        elif activation == 'selu':\n",
        "            self.activation = nn.SELU(inplace=True)\n",
        "        elif activation == 'tanh':\n",
        "            self.activation = nn.Tanh()\n",
        "        elif activation == 'none':\n",
        "            self.activation = None\n",
        "        else:\n",
        "            assert 0, \"Unsupported activation: {}\".format(activation)\n",
        "\n",
        "        # initialize convolution\n",
        "        if norm == 'sn':\n",
        "            self.conv = SpectralNorm(nn.Conv2d(input_dim, output_dim, kernel_size, stride, bias=self.use_bias))\n",
        "        else:\n",
        "            self.conv = nn.Conv2d(input_dim, output_dim, kernel_size, stride, bias=self.use_bias)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.conv(self.pad(x))\n",
        "        if self.norm:\n",
        "            x = self.norm(x)\n",
        "        if self.activation:\n",
        "            x = self.activation(x)\n",
        "        return x\n",
        "\n",
        "class LinearBlock(nn.Module):\n",
        "    def __init__(self, input_dim, output_dim, norm='none', activation='relu'):\n",
        "        super(LinearBlock, self).__init__()\n",
        "        use_bias = True\n",
        "        # initialize fully connected layer\n",
        "        if norm == 'sn':\n",
        "            self.fc = SpectralNorm(nn.Linear(input_dim, output_dim, bias=use_bias))\n",
        "        else:\n",
        "            self.fc = nn.Linear(input_dim, output_dim, bias=use_bias)\n",
        "\n",
        "        # initialize normalization\n",
        "        norm_dim = output_dim\n",
        "        if norm == 'bn':\n",
        "            self.norm = nn.BatchNorm1d(norm_dim)\n",
        "        elif norm == 'in':\n",
        "            self.norm = nn.InstanceNorm1d(norm_dim)\n",
        "        elif norm == 'ln':\n",
        "            self.norm = LayerNorm(norm_dim)\n",
        "        elif norm == 'none' or norm == 'sn':\n",
        "            self.norm = None\n",
        "        else:\n",
        "            assert 0, \"Unsupported normalization: {}\".format(norm)\n",
        "\n",
        "        # initialize activation\n",
        "        if activation == 'relu':\n",
        "            self.activation = nn.ReLU(inplace=True)\n",
        "        elif activation == 'lrelu':\n",
        "            self.activation = nn.LeakyReLU(0.2, inplace=True)\n",
        "        elif activation == 'prelu':\n",
        "            self.activation = nn.PReLU()\n",
        "        elif activation == 'selu':\n",
        "            self.activation = nn.SELU(inplace=True)\n",
        "        elif activation == 'tanh':\n",
        "            self.activation = nn.Tanh()\n",
        "        elif activation == 'none':\n",
        "            self.activation = None\n",
        "        else:\n",
        "            assert 0, \"Unsupported activation: {}\".format(activation)\n",
        "\n",
        "    def forward(self, x):\n",
        "        out = self.fc(x)\n",
        "        if self.norm:\n",
        "            out = self.norm(out)\n",
        "        if self.activation:\n",
        "            out = self.activation(out)\n",
        "        return out\n",
        "\n",
        "##################################################################################\n",
        "# VGG network definition\n",
        "##################################################################################\n",
        "class Vgg16(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(Vgg16, self).__init__()\n",
        "        self.conv1_1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)\n",
        "        self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)\n",
        "\n",
        "        self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)\n",
        "        self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)\n",
        "\n",
        "        self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1)\n",
        "        self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)\n",
        "        self.conv3_3 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)\n",
        "\n",
        "        self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1)\n",
        "        self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)\n",
        "        self.conv4_3 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)\n",
        "\n",
        "        self.conv5_1 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)\n",
        "        self.conv5_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)\n",
        "        self.conv5_3 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)\n",
        "\n",
        "    def forward(self, X):\n",
        "        h = F.relu(self.conv1_1(X), inplace=True)\n",
        "        h = F.relu(self.conv1_2(h), inplace=True)\n",
        "        # relu1_2 = h\n",
        "        h = F.max_pool2d(h, kernel_size=2, stride=2)\n",
        "\n",
        "        h = F.relu(self.conv2_1(h), inplace=True)\n",
        "        h = F.relu(self.conv2_2(h), inplace=True)\n",
        "        # relu2_2 = h\n",
        "        h = F.max_pool2d(h, kernel_size=2, stride=2)\n",
        "\n",
        "        h = F.relu(self.conv3_1(h), inplace=True)\n",
        "        h = F.relu(self.conv3_2(h), inplace=True)\n",
        "        h = F.relu(self.conv3_3(h), inplace=True)\n",
        "        # relu3_3 = h\n",
        "        h = F.max_pool2d(h, kernel_size=2, stride=2)\n",
        "\n",
        "        h = F.relu(self.conv4_1(h), inplace=True)\n",
        "        h = F.relu(self.conv4_2(h), inplace=True)\n",
        "        h = F.relu(self.conv4_3(h), inplace=True)\n",
        "        # relu4_3 = h\n",
        "\n",
        "        h = F.relu(self.conv5_1(h), inplace=True)\n",
        "        h = F.relu(self.conv5_2(h), inplace=True)\n",
        "        h = F.relu(self.conv5_3(h), inplace=True)\n",
        "        relu5_3 = h\n",
        "\n",
        "        return relu5_3\n",
        "        # return [relu1_2, relu2_2, relu3_3, relu4_3]\n",
        "\n",
        "##################################################################################\n",
        "# Normalization layers\n",
        "##################################################################################\n",
        "class AdaptiveInstanceNorm2d(nn.Module):\n",
        "    def __init__(self, num_features, eps=1e-5, momentum=0.1):\n",
        "        super(AdaptiveInstanceNorm2d, self).__init__()\n",
        "        self.num_features = num_features\n",
        "        self.eps = eps\n",
        "        self.momentum = momentum\n",
        "        # weight and bias are dynamically assigned\n",
        "        self.weight = None\n",
        "        self.bias = None\n",
        "        # just dummy buffers, not used\n",
        "        self.register_buffer('running_mean', torch.zeros(num_features))\n",
        "        self.register_buffer('running_var', torch.ones(num_features))\n",
        "\n",
        "    def forward(self, x):\n",
        "        assert self.weight is not None and self.bias is not None, \"Please assign weight and bias before calling AdaIN!\"\n",
        "        b, c = x.size(0), x.size(1)\n",
        "        running_mean = self.running_mean.repeat(b)\n",
        "        running_var = self.running_var.repeat(b)\n",
        "\n",
        "        # Apply instance norm\n",
        "        x_reshaped = x.contiguous().view(1, b * c, *x.size()[2:])\n",
        "\n",
        "        out = F.batch_norm(\n",
        "            x_reshaped, running_mean, running_var, self.weight, self.bias,\n",
        "            True, self.momentum, self.eps)\n",
        "\n",
        "        return out.view(b, c, *x.size()[2:])\n",
        "\n",
        "    def __repr__(self):\n",
        "        return self.__class__.__name__ + '(' + str(self.num_features) + ')'\n",
        "\n",
        "\n",
        "class LayerNorm(nn.Module):\n",
        "    def __init__(self, num_features, eps=1e-5, affine=True):\n",
        "        super(LayerNorm, self).__init__()\n",
        "        self.num_features = num_features\n",
        "        self.affine = affine\n",
        "        self.eps = eps\n",
        "\n",
        "        if self.affine:\n",
        "            self.gamma = nn.Parameter(torch.Tensor(num_features).uniform_())\n",
        "            self.beta = nn.Parameter(torch.zeros(num_features))\n",
        "\n",
        "    def forward(self, x):\n",
        "        shape = [-1] + [1] * (x.dim() - 1)\n",
        "        # print(x.size())\n",
        "        if x.size(0) == 1:\n",
        "            # These two lines run much faster in pytorch 0.4 than the two lines listed below.\n",
        "            mean = x.view(-1).mean().view(*shape)\n",
        "            std = x.view(-1).std().view(*shape)\n",
        "        else:\n",
        "            mean = x.view(x.size(0), -1).mean(1).view(*shape)\n",
        "            std = x.view(x.size(0), -1).std(1).view(*shape)\n",
        "\n",
        "        x = (x - mean) / (std + self.eps)\n",
        "\n",
        "        if self.affine:\n",
        "            shape = [1, -1] + [1] * (x.dim() - 2)\n",
        "            x = x * self.gamma.view(*shape) + self.beta.view(*shape)\n",
        "        return x\n",
        "\n",
        "def l2normalize(v, eps=1e-12):\n",
        "    return v / (v.norm() + eps)\n",
        "\n",
        "\n",
        "class SpectralNorm(nn.Module):\n",
        "    \"\"\"\n",
        "    Based on the paper \"Spectral Normalization for Generative Adversarial Networks\" by Takeru Miyato, Toshiki Kataoka, Masanori Koyama, Yuichi Yoshida\n",
        "    and the Pytorch implementation https://github.com/christiancosgrove/pytorch-spectral-normalization-gan\n",
        "    \"\"\"\n",
        "    def __init__(self, module, name='weight', power_iterations=1):\n",
        "        super(SpectralNorm, self).__init__()\n",
        "        self.module = module\n",
        "        self.name = name\n",
        "        self.power_iterations = power_iterations\n",
        "        if not self._made_params():\n",
        "            self._make_params()\n",
        "\n",
        "    def _update_u_v(self):\n",
        "        u = getattr(self.module, self.name + \"_u\")\n",
        "        v = getattr(self.module, self.name + \"_v\")\n",
        "        w = getattr(self.module, self.name + \"_bar\")\n",
        "\n",
        "        height = w.data.shape[0]\n",
        "        for _ in range(self.power_iterations):\n",
        "            v.data = l2normalize(torch.mv(torch.t(w.view(height,-1).data), u.data))\n",
        "            u.data = l2normalize(torch.mv(w.view(height,-1).data, v.data))\n",
        "        sigma = u.dot(w.view(height, -1).mv(v))\n",
        "        setattr(self.module, self.name, w / sigma.expand_as(w))\n",
        "\n",
        "    def _made_params(self):\n",
        "        try:\n",
        "            u = getattr(self.module, self.name + \"_u\")\n",
        "            v = getattr(self.module, self.name + \"_v\")\n",
        "            w = getattr(self.module, self.name + \"_bar\")\n",
        "            return True\n",
        "        except AttributeError:\n",
        "            return False\n",
        "\n",
        "\n",
        "    def _make_params(self):\n",
        "        w = getattr(self.module, self.name)\n",
        "\n",
        "        height = w.data.shape[0]\n",
        "        width = w.view(height, -1).data.shape[1]\n",
        "\n",
        "        u = nn.Parameter(w.data.new(height).normal_(0, 1), requires_grad=False)\n",
        "        v = nn.Parameter(w.data.new(width).normal_(0, 1), requires_grad=False)\n",
        "        u.data = l2normalize(u.data)\n",
        "        v.data = l2normalize(v.data)\n",
        "        w_bar = nn.Parameter(w.data)\n",
        "\n",
        "        del self.module._parameters[self.name]\n",
        "\n",
        "        self.module.register_parameter(self.name + \"_u\", u)\n",
        "        self.module.register_parameter(self.name + \"_v\", v)\n",
        "        self.module.register_parameter(self.name + \"_bar\", w_bar)\n",
        "\n",
        "\n",
        "    def forward(self, *args):\n",
        "        self._update_u_v()\n",
        "        return self.module.forward(*args)\n"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xD4r65toihTF",
        "colab_type": "text"
      },
      "source": [
        "# utils.py\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "603txbXdig7X",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\"\"\"\n",
        "Copyright (C) 2018 NVIDIA Corporation.  All rights reserved.\n",
        "Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n",
        "\"\"\"\n",
        "# from torch.utils.serialization import load_lua\n",
        "import torchfile\n",
        "from torch.utils.data import DataLoader\n",
        "from torchvision.models import inception_v3\n",
        "# from networks import Vgg16\n",
        "from torch.autograd import Variable\n",
        "from torch.optim import lr_scheduler\n",
        "from torchvision import transforms\n",
        "# from data import ImageFilelist, ImageFolder, ImageFolder_with_subfolders\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import os\n",
        "import math\n",
        "import torchvision.utils as vutils\n",
        "import yaml\n",
        "import numpy as np\n",
        "import torch.nn.init as init\n",
        "import time\n",
        "# Methods\n",
        "# get_all_data_loaders      : primary data loader interface (load trainA, testA, trainB, testB)\n",
        "# get_data_loader_list      : list-based data loader\n",
        "# get_data_loader_folder    : folder-based data loader\n",
        "# get_config                : load yaml file\n",
        "# eformat                   :\n",
        "# write_2images             : save output image\n",
        "# prepare_sub_folder        : create checkpoints and images folders for saving outputs\n",
        "# write_one_row_html        : write one row of the html file for output images\n",
        "# write_html                : create the html file.\n",
        "# write_loss\n",
        "# slerp\n",
        "# get_slerp_interp\n",
        "# get_model_list\n",
        "# load_vgg16\n",
        "# load_inception\n",
        "# vgg_preprocess\n",
        "# get_scheduler\n",
        "# weights_init\n",
        "\n",
        "def get_all_data_loaders(conf):\n",
        "    batch_size = conf['batch_size']\n",
        "    num_workers = conf['num_workers']\n",
        "    if 'new_size' in conf:\n",
        "        new_size_a = new_size_b = conf['new_size']\n",
        "    else:\n",
        "        new_size_a = conf['new_size_a']\n",
        "        new_size_b = conf['new_size_b']\n",
        "    height = conf['crop_image_height']\n",
        "    width = conf['crop_image_width']\n",
        "\n",
        "    train_loader_a, train_loader_b, test_loader_a, test_loader_b = [], [], [], []\n",
        "    if 'data_root' in conf:\n",
        "        if 'inbalenceDataSets' in conf:\n",
        "            if (conf['inbalenceDataSets']['imbalance_sub_dataset']):\n",
        "\n",
        "                train_loader_a.append(get_data_loader_folder(os.path.join(conf['data_root'],  'trainA'), batch_size, True,\n",
        "                                                             new_size_a, height, width, num_workers, True, config=conf,\n",
        "                                                             is_data_A=True, ratio_1_to_2=conf['inbalenceDataSets']['ratio_A_1_to_2']))\n",
        "                test_loader_a.append(get_data_loader_folder(os.path.join(conf['data_root'], 'testA'), batch_size, False,\n",
        "                                                            new_size_a, new_size_a, new_size_a, num_workers, True,\n",
        "                                                            config=conf, is_data_A=True, ratio_1_to_2=conf['inbalenceDataSets']['ratio_A_1_to_2']))\n",
        "                train_loader_b.append(get_data_loader_folder(os.path.join(conf['data_root'], 'trainB'), batch_size, True,\n",
        "                                                             new_size_b, height, width, num_workers, True, config=conf,\n",
        "                                                             is_data_A=False, ratio_1_to_2=conf['inbalenceDataSets']['ratio_B_1_to_2']))\n",
        "                test_loader_b.append(get_data_loader_folder(os.path.join(conf['data_root'], 'testB'), batch_size, False,\n",
        "                                                            new_size_b, new_size_b, new_size_b, num_workers, True,\n",
        "                                                            config=conf, is_data_A=False, ratio_1_to_2=conf['inbalenceDataSets']['ratio_B_1_to_2']))\n",
        "            else:\n",
        "                train_loader_a.append(get_data_loader_folder(os.path.join(conf['data_root'], 'trainA'), batch_size, True,\n",
        "                                                        new_size_a, height, width, num_workers, True, config=conf,\n",
        "                                                        is_data_A=True))\n",
        "                test_loader_a.append(get_data_loader_folder(os.path.join(conf['data_root'], 'testA'), batch_size, False,\n",
        "                                                       new_size_a, new_size_a, new_size_a, num_workers, True,\n",
        "                                                       config=conf, is_data_A=True))\n",
        "                train_loader_b.append(get_data_loader_folder(os.path.join(conf['data_root'], 'trainB'), batch_size, True,\n",
        "                                                        new_size_b, height, width, num_workers, True, config=conf,\n",
        "                                                        is_data_A=False))\n",
        "                test_loader_b.append(get_data_loader_folder(os.path.join(conf['data_root'], 'testB'), batch_size, False,\n",
        "                                                       new_size_b, new_size_b, new_size_b, num_workers, True,\n",
        "                                                       config=conf, is_data_A=False))\n",
        "        else:\n",
        "            train_loader_a.append(get_data_loader_folder(os.path.join(conf['data_root'], 'trainA'), batch_size, True,\n",
        "                                                  new_size_a, height, width, num_workers, True, config=conf, is_data_A=True))\n",
        "            test_loader_a.append(get_data_loader_folder(os.path.join(conf['data_root'], 'testA'), batch_size, False,\n",
        "                                                 new_size_a, new_size_a, new_size_a, num_workers, True, config=conf, is_data_A=True))\n",
        "            train_loader_b.append(get_data_loader_folder(os.path.join(conf['data_root'], 'trainB'), batch_size, True,\n",
        "                                                  new_size_b, height, width, num_workers, True, config=conf, is_data_A=False))\n",
        "            test_loader_b.append(get_data_loader_folder(os.path.join(conf['data_root'], 'testB'), batch_size, False,\n",
        "                                                 new_size_b, new_size_b, new_size_b, num_workers, True, config=conf, is_data_A=False))\n",
        "    else:\n",
        "        train_loader_a.append(get_data_loader_list(conf['data_folder_train_a'], conf['data_list_train_a'], batch_size, True,\n",
        "                                                new_size_a, height, width, num_workers, True, is_data_A=True))\n",
        "        test_loader_a.append(get_data_loader_list(conf['data_folder_test_a'], conf['data_list_test_a'], batch_size, False,\n",
        "                                                new_size_a, new_size_a, new_size_a, num_workers, True, is_data_A=True))\n",
        "        train_loader_b.append(get_data_loader_list(conf['data_folder_train_b'], conf['data_list_train_b'], batch_size, True,\n",
        "                                                new_size_b, height, width, num_workers, True, is_data_A=False))\n",
        "        test_loader_b.append(get_data_loader_list(conf['data_folder_test_b'], conf['data_list_test_b'], batch_size, False,\n",
        "                                                new_size_b, new_size_b, new_size_b, num_workers, True, is_data_A=False))\n",
        "    return train_loader_a, train_loader_b, test_loader_a, test_loader_b\n",
        "\n",
        "def get_data_loader_list(root, file_list, batch_size, train, new_size=None,\n",
        "                           height=256, width=256, num_workers=4, crop=True):\n",
        "    transform_list = [transforms.ToTensor(),\n",
        "                      transforms.Normalize((0.5, 0.5, 0.5),\n",
        "                                           (0.5, 0.5, 0.5))]\n",
        "    transform_list = [transforms.RandomCrop((height, width))] + transform_list if crop else transform_list\n",
        "    transform_list = [transforms.Resize(new_size)] + transform_list if new_size is not None else transform_list\n",
        "    transform_list = [transforms.RandomHorizontalFlip()] + transform_list if train else transform_list\n",
        "    transform = transforms.Compose(transform_list)\n",
        "    dataset = ImageFilelist(root, file_list, transform=transform)\n",
        "    loader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=train, drop_last=True, num_workers=num_workers)\n",
        "    return loader\n",
        "\n",
        "def dim3to1(x):\n",
        "    x = torch.sum(x, 0).unsqueeze(0)\n",
        "    return x\n",
        "\n",
        "def get_data_loader_folder(input_folder, batch_size, train, new_size=None,\n",
        "                           height=256, width=256, num_workers=4, crop=True, config=None, is_data_A=None, ratio_1_to_2=None):\n",
        "    transform_list = [transforms.ToTensor(),\n",
        "                      transforms.Normalize((0.5, 0.5, 0.5),\n",
        "                                           (0.5, 0.5, 0.5))]\n",
        "    transform_list = [transforms.CenterCrop((new_size))] + transform_list if not train else transform_list\n",
        "    transform_list = [transforms.RandomCrop((height, width))] + transform_list if crop else transform_list\n",
        "    transform_list = [transforms.Resize(new_size)] + transform_list if new_size is not None else transform_list\n",
        "    transform_list = (transform_list + [dim3to1]) if config['input_dim_a'] == 1 and is_data_A else transform_list\n",
        "    transform_list = (transform_list + [dim3to1]) if config['input_dim_b'] == 1 and not is_data_A else transform_list\n",
        "    if config is not None:\n",
        "        if config['do_HorizontalFlip']:\n",
        "            transform_list = [transforms.RandomHorizontalFlip()] + transform_list if train else transform_list\n",
        "    if config is not None:\n",
        "        if config['do_RandomResizedCrop']:\n",
        "            ratio_max = eval(config['RandomResizedCrop_ratio_max'])\n",
        "            ratio_min = eval(config['RandomResizedCrop_ratio_min'])\n",
        "            transform_list = [transforms.RandomResizedCrop(size=new_size, scale=(config['RandomResizedCrop_scale_min'], config['RandomResizedCrop_scale_max']),\n",
        "                                                           ratio=(ratio_min, ratio_max))] + transform_list if train else transform_list\n",
        "    if config is not None:\n",
        "        if config['do_VerticalFlip']:\n",
        "            transform_list = [transforms.RandomVerticalFlip(p=0.35)] + transform_list if train else transform_list\n",
        "\n",
        "    if config is not None:\n",
        "        if is_data_A is not None:\n",
        "            if config['do_ColorJitter_A'] and is_data_A:\n",
        "                transform_list = [transforms.ColorJitter(brightness=config['ColorJitter_brightness'], contrast=config['ColorJitter_contrast'], saturation=config['ColorJitter_saturation'], hue=config['ColorJitter_hue'])] + transform_list if train else transform_list\n",
        "            elif config['do_ColorJitter_B'] and not is_data_A:\n",
        "                transform_list = [transforms.ColorJitter(brightness=config['ColorJitter_brightness'], contrast=config['ColorJitter_contrast'], saturation=config['ColorJitter_saturation'], hue=config['ColorJitter_hue'])] + transform_list if train else transform_list\n",
        "        else:\n",
        "            if config['do_ColorJitter']:\n",
        "                transform_list = [transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=config['ColorJitter_hue'])] + transform_list if train else transform_list\n",
        "\n",
        "    if config is not None:\n",
        "        if config['do_RandomGrayscale']:\n",
        "            transform_list = [transforms.RandomGrayscale(p=config['RandomGrayscale_P'])] + transform_list if train else transform_list\n",
        "\n",
        "    if config is not None:\n",
        "        if config['do_RandomRotation']:\n",
        "            transform_list = [transforms.RandomRotation(degrees=config['RandomRotation_degree'], expand=False)] + transform_list if train else transform_list\n",
        "\n",
        "    if config is not None:\n",
        "        if config['do_RandomAffine']:\n",
        "            config['RandomAffine_translate_w']\n",
        "            transform_list = [transforms.RandomAffine(translate=(config['RandomAffine_translate_h'], config['RandomAffine_translate_w']), degrees=0, scale=(0.9, 1.1))] + transform_list if train else transform_list\n",
        "\n",
        "    if config is not None:\n",
        "        if config['do_RandomPerspective']:\n",
        "            transform_list = [transforms.RandomPerspective(distortion_scale=0.35, p=0.5)] + transform_list if train else transform_list\n",
        "\n",
        "    transform = transforms.Compose(transform_list)\n",
        "    if ratio_1_to_2 is None:\n",
        "        dataset = ImageFolder(input_folder, transform=transform)\n",
        "    else:\n",
        "        input_folder_1 = os.path.join(input_folder, '1')\n",
        "        input_folder_2 = os.path.join(input_folder, '2')\n",
        "        dataset = ImageFolder_with_subfolders(root1=input_folder_1, root2=input_folder_2, ratio_1_to_2=ratio_1_to_2, transform=transform)\n",
        "\n",
        "    loader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=train, drop_last=True, num_workers=num_workers)\n",
        "    return loader\n",
        "\n",
        "def get_config(config):\n",
        "    with open(config, 'r') as stream:\n",
        "        return yaml.safe_load(stream)\n",
        "\n",
        "def eformat(f, prec):\n",
        "    s = \"%.*e\"%(prec, f)\n",
        "    mantissa, exp = s.split('e')\n",
        "    # add 1 to digits as 1 is taken by sign +/-\n",
        "    return \"%se%d\"%(mantissa, int(exp))\n",
        "\n",
        "def __write_images(image_outputs, display_image_num, file_name):\n",
        "    if image_outputs[0].shape[1]>3:\n",
        "        rand_disp_index_s = range(image_outputs[0].shape[1])\n",
        "        rand_disp_index_s = np.random.permutation(rand_disp_index_s)\n",
        "        rand_disp_index_s = rand_disp_index_s[:3]\n",
        "        rand_disp_index_s = np.sort(rand_disp_index_s)\n",
        "        image_outputs = [images[:, rand_disp_index_s, :, :] if images.shape[1] > 3 else images for images in image_outputs] # TODO tmp\n",
        "    else:\n",
        "        image_outputs = [images.expand(-1, 3, -1, -1) for images in image_outputs] # expand gray-scale images to 3 channels\n",
        "    image_tensor = torch.cat([images[:display_image_num] for images in image_outputs], 0)\n",
        "    image_grid = vutils.make_grid(image_tensor.data, nrow=display_image_num, padding=0, normalize=True)\n",
        "\n",
        "    vutils.save_image(image_grid, file_name, nrow=1)\n",
        "\n",
        "    # from PIL import Image\n",
        "    grid = vutils.make_grid(image_grid, nrow=8, padding=2, pad_value=0,\n",
        "                     normalize=False, range=None, scale_each=False)\n",
        "\n",
        "\n",
        "    return grid.to('cpu', torch.uint8).numpy()\n",
        "\n",
        "def write_2images(image_outputs, display_image_num, image_directory, postfix, do_a2b=True, do_b2a=True):\n",
        "    n = len(image_outputs)\n",
        "    gen_a2b_im, gen_b2a_im = [], []\n",
        "    if do_a2b:\n",
        "        gen_a2b_im = __write_images(image_outputs[0:n//2], display_image_num, '%s/gen_a2b_%s.jpg' % (image_directory, postfix)) if do_a2b else None\n",
        "    if do_b2a:\n",
        "        gen_b2a_im = __write_images(image_outputs[n//2:n], display_image_num, '%s/gen_b2a_%s.jpg' % (image_directory, postfix)) if do_b2a else None\n",
        "    return gen_a2b_im, gen_b2a_im\n",
        "\n",
        "def prepare_sub_folder(output_directory):\n",
        "    image_directory = os.path.join(output_directory, 'images')\n",
        "    if not os.path.exists(image_directory):\n",
        "        print(\"Creating directory: {}\".format(image_directory))\n",
        "        os.makedirs(image_directory)\n",
        "    checkpoint_directory = os.path.join(output_directory, 'checkpoints')\n",
        "    if not os.path.exists(checkpoint_directory):\n",
        "        print(\"Creating directory: {}\".format(checkpoint_directory))\n",
        "        os.makedirs(checkpoint_directory)\n",
        "    checkpoint_log = os.path.join(output_directory, 'log')\n",
        "    if not os.path.exists(checkpoint_log):\n",
        "        print(\"Creating directory: {}\".format(checkpoint_log))\n",
        "        os.makedirs(checkpoint_log)\n",
        "    return checkpoint_directory, image_directory, checkpoint_log\n",
        "\n",
        "def write_one_row_html(html_file, iterations, img_filename, all_size):\n",
        "    html_file.write(\"<h3>iteration [%d] (%s)</h3>\" % (iterations,img_filename.split('/')[-1]))\n",
        "    html_file.write(\"\"\"\n",
        "        <p><a href=\"%s\">\n",
        "          <img src=\"%s\" style=\"width:%dpx\">\n",
        "        </a><br>\n",
        "        <p>\n",
        "        \"\"\" % (img_filename, img_filename, all_size))\n",
        "    return\n",
        "\n",
        "def write_html(filename, iterations, image_save_iterations, image_directory, all_size=1536, do_a2b=True, do_b2a=True):\n",
        "    html_file = open(filename, \"w\")\n",
        "    html_file.write('''\n",
        "    <!DOCTYPE html>\n",
        "    <html>\n",
        "    <head>\n",
        "      <title>Experiment name = %s</title>\n",
        "      <meta http-equiv=\"refresh\" content=\"30\">\n",
        "    </head>\n",
        "    <body>\n",
        "    ''' % os.path.basename(filename))\n",
        "    html_file.write(\"<h3>current</h3>\")\n",
        "    if do_a2b:\n",
        "        write_one_row_html(html_file, iterations, '%s/gen_a2b_train_current.jpg' % (image_directory), all_size)\n",
        "    if do_b2a:\n",
        "        write_one_row_html(html_file, iterations, '%s/gen_b2a_train_current.jpg' % (image_directory), all_size)\n",
        "    for j in range(iterations, image_save_iterations-1, -1):\n",
        "        if j % image_save_iterations == 0:\n",
        "            if do_a2b:\n",
        "                write_one_row_html(html_file, j, '%s/gen_a2b_test_%08d.jpg' % (image_directory, j), all_size)\n",
        "            if do_b2a:\n",
        "                write_one_row_html(html_file, j, '%s/gen_b2a_test_%08d.jpg' % (image_directory, j), all_size)\n",
        "            if do_a2b:\n",
        "                write_one_row_html(html_file, j, '%s/gen_a2b_train_%08d.jpg' % (image_directory, j), all_size)\n",
        "            if do_b2a:\n",
        "                write_one_row_html(html_file, j, '%s/gen_b2a_train_%08d.jpg' % (image_directory, j), all_size)\n",
        "    html_file.write(\"</body></html>\")\n",
        "    html_file.close()\n",
        "\n",
        "def write_loss(iterations, trainer, train_writer):\n",
        "    members = [attr for attr in dir(trainer)\n",
        "               if not callable(getattr(trainer, attr)) and not attr.startswith(\"__\") and ('loss' in attr or 'grad' in attr or 'conf' in attr or 'nwd' in attr or 'do' in attr)]\n",
        "    for m in members:\n",
        "        tmpatter = getattr(trainer, m)\n",
        "        if m.endswith('_conf'):\n",
        "            m = 'configs/' + m\n",
        "        if '_a_' in m or '_b_' in m or '_ab_' in m or '_ba_' in m:\n",
        "            ind = m.find('_a_')\n",
        "            ind = ind if ind != -1 else m.find('_b_')\n",
        "            ind = ind if ind != -1 else m.find('_ab_')\n",
        "            ind = ind if ind != -1 else m.find('_ba_')\n",
        "\n",
        "            m = m[0:ind] + '/' + m\n",
        "        if type(tmpatter) is bool:\n",
        "            tmpatter = 1 if tmpatter else 0\n",
        "            m = 'configs/' + m\n",
        "\n",
        "        if type(tmpatter) is list:\n",
        "            tmpList = tmpatter\n",
        "            tmpScal = {}\n",
        "            for i, listItem in enumerate(tmpList):\n",
        "                if type(listItem) is torch.Tensor:\n",
        "                    tmpScal.update({str(i): listItem.data.cpu().numpy()})\n",
        "                else:\n",
        "                    tmpScal.update({m + '_' + str(i): listItem})\n",
        "            train_writer.add_scalars(m, tmpScal, iterations + 1)\n",
        "        else:\n",
        "            train_writer.add_scalar(m, tmpatter, iterations + 1)\n",
        "\n",
        "\n",
        "def slerp(val, low, high):\n",
        "    \"\"\"\n",
        "    original: Animating Rotation with Quaternion Curves, Ken Shoemake\n",
        "    https://arxiv.org/abs/1609.04468\n",
        "    Code: https://github.com/soumith/dcgan.torch/issues/14, Tom White\n",
        "    \"\"\"\n",
        "    omega = np.arccos(np.dot(low / np.linalg.norm(low), high / np.linalg.norm(high)))\n",
        "    so = np.sin(omega)\n",
        "    return np.sin((1.0 - val) * omega) / so * low + np.sin(val * omega) / so * high\n",
        "\n",
        "def get_slerp_interp(nb_latents, nb_interp, z_dim):\n",
        "    \"\"\"\n",
        "    modified from: PyTorch inference for \"Progressive Growing of GANs\" with CelebA snapshot\n",
        "    https://github.com/ptrblck/prog_gans_pytorch_inference\n",
        "    \"\"\"\n",
        "\n",
        "    latent_interps = np.empty(shape=(0, z_dim), dtype=np.float32)\n",
        "    for _ in range(nb_latents):\n",
        "        low = np.random.randn(z_dim)\n",
        "        high = np.random.randn(z_dim)  # low + np.random.randn(512) * 0.7\n",
        "        interp_vals = np.linspace(0, 1, num=nb_interp)\n",
        "        latent_interp = np.array([slerp(v, low, high) for v in interp_vals],\n",
        "                                 dtype=np.float32)\n",
        "        latent_interps = np.vstack((latent_interps, latent_interp))\n",
        "\n",
        "    return latent_interps[:, :, np.newaxis, np.newaxis]\n",
        "\n",
        "# Get model list for resume\n",
        "def get_model_list(dirname, key):\n",
        "    if os.path.exists(dirname) is False:\n",
        "        return None\n",
        "    gen_models = [os.path.join(dirname, f) for f in os.listdir(dirname) if\n",
        "                  os.path.isfile(os.path.join(dirname, f)) and key in f and \".pt\" in f]\n",
        "    if gen_models is None:\n",
        "        return None\n",
        "    gen_models.sort()\n",
        "    if len(gen_models)>0:\n",
        "        last_model_name = gen_models[-1]\n",
        "    else:\n",
        "        last_model_name = None\n",
        "    return last_model_name\n",
        "\n",
        "def load_vgg16(model_dir):\n",
        "    \"\"\" Use the model from https://github.com/abhiskk/fast-neural-style/blob/master/neural_style/utils.py \"\"\"\n",
        "    if not os.path.exists(model_dir):\n",
        "        os.mkdir(model_dir)\n",
        "    if not os.path.exists(os.path.join(model_dir, 'vgg16.weight')):\n",
        "        if not os.path.exists(os.path.join(model_dir, 'vgg16.t7')):\n",
        "            os.system('wget https://www.dropbox.com/s/76l3rt4kyi3s8x7/vgg16.t7?dl=1 -O ' + os.path.join(model_dir, 'vgg16.t7'))\n",
        "        # vgglua = load_lua(os.path.join(model_dir, 'vgg16.t7'))\n",
        "        vgglua = torchfile.load(os.path.join(model_dir, 'vgg16.t7'))\n",
        "\n",
        "        vgg = Vgg16()\n",
        "        for (src, dst) in zip(vgglua.parameters()[0], vgg.parameters()):\n",
        "            dst.data[:] = src\n",
        "        torch.save(vgg.state_dict(), os.path.join(model_dir, 'vgg16.weight'))\n",
        "    vgg = Vgg16()\n",
        "    vgg.load_state_dict(torch.load(os.path.join(model_dir, 'vgg16.weight')))\n",
        "    return vgg\n",
        "\n",
        "def load_inception(model_path, pretrained=False):\n",
        "    if model_path is not None:\n",
        "        state_dict = torch.load(model_path)\n",
        "    model = inception_v3(pretrained=pretrained, transform_input=True)\n",
        "    model.aux_logits = False\n",
        "    num_ftrs = model.fc.in_features\n",
        "    model.fc = nn.Linear(num_ftrs, state_dict['fc.weight'].size(0))\n",
        "    model.load_state_dict(state_dict)\n",
        "    for param in model.parameters():\n",
        "        param.requires_grad = False\n",
        "    return model\n",
        "\n",
        "def vgg_preprocess(batch):\n",
        "    tensortype = type(batch.data)\n",
        "    (r, g, b) = torch.chunk(batch, 3, dim = 1)\n",
        "    batch = torch.cat((b, g, r), dim = 1) # convert RGB to BGR\n",
        "    batch = (batch + 1) * 255 * 0.5 # [-1, 1] -> [0, 255]\n",
        "    mean = tensortype(batch.data.size()).cuda()\n",
        "    mean[:, 0, :, :] = 103.939\n",
        "    mean[:, 1, :, :] = 116.779\n",
        "    mean[:, 2, :, :] = 123.680\n",
        "    batch = batch.sub(Variable(mean)) # subtract mean\n",
        "    return batch\n",
        "\n",
        "def get_scheduler(optimizer, hyperparameters, iterations=-1):\n",
        "    if 'lr_policy' not in hyperparameters or hyperparameters['lr_policy'] == 'constant':\n",
        "        scheduler = None # constant scheduler\n",
        "    elif hyperparameters['lr_policy'] == 'step':\n",
        "        scheduler = lr_scheduler.StepLR(optimizer, step_size=hyperparameters['step_size'],\n",
        "                                        gamma=hyperparameters['gamma'], last_epoch=iterations)\n",
        "    else:\n",
        "        return NotImplementedError('learning rate policy [%s] is not implemented', hyperparameters['lr_policy'])\n",
        "    return scheduler\n",
        "\n",
        "def weights_init(init_type='gaussian'):\n",
        "    def init_fun(m):\n",
        "        classname = m.__class__.__name__\n",
        "        if (classname.find('Conv') == 0 or classname.find('Linear') == 0) and hasattr(m, 'weight'):\n",
        "            # print m.__class__.__name__\n",
        "            if init_type == 'gaussian':\n",
        "                init.normal_(m.weight.data, 0.0, 0.02)\n",
        "            elif init_type == 'xavier':\n",
        "                init.xavier_normal_(m.weight.data, gain=math.sqrt(2))\n",
        "            elif init_type == 'kaiming':\n",
        "                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')\n",
        "            elif init_type == 'orthogonal':\n",
        "                init.orthogonal_(m.weight.data, gain=math.sqrt(2))\n",
        "            elif init_type == 'default':\n",
        "                pass\n",
        "            else:\n",
        "                assert 0, \"Unsupported initialization: {}\".format(init_type)\n",
        "            if hasattr(m, 'bias') and m.bias is not None:\n",
        "                init.constant_(m.bias.data, 0.0)\n",
        "\n",
        "    return init_fun\n",
        "\n",
        "class Timer:\n",
        "    def __init__(self, msg):\n",
        "        self.msg = msg\n",
        "        self.start_time = None\n",
        "\n",
        "    def __enter__(self):\n",
        "        self.start_time = time.time()\n",
        "\n",
        "    def __exit__(self, exc_type, exc_value, exc_tb):\n",
        "        print(self.msg % (time.time() - self.start_time))\n",
        "\n",
        "def pytorch03_to_pytorch04(state_dict_base, trainer_name):\n",
        "    def __conversion_core(state_dict_base, trainer_name):\n",
        "        state_dict = state_dict_base.copy()\n",
        "        if trainer_name == 'MUNIT':\n",
        "            for key, value in state_dict_base.items():\n",
        "                if key.endswith(('enc_content.model.0.norm.running_mean',\n",
        "                                 'enc_content.model.0.norm.running_var',\n",
        "                                 'enc_content.model.1.norm.running_mean',\n",
        "                                 'enc_content.model.1.norm.running_var',\n",
        "                                 'enc_content.model.2.norm.running_mean',\n",
        "                                 'enc_content.model.2.norm.running_var',\n",
        "                                 'enc_content.model.3.model.0.model.1.norm.running_mean',\n",
        "                                 'enc_content.model.3.model.0.model.1.norm.running_var',\n",
        "                                 'enc_content.model.3.model.0.model.0.norm.running_mean',\n",
        "                                 'enc_content.model.3.model.0.model.0.norm.running_var',\n",
        "                                 'enc_content.model.3.model.1.model.1.norm.running_mean',\n",
        "                                 'enc_content.model.3.model.1.model.1.norm.running_var',\n",
        "                                 'enc_content.model.3.model.1.model.0.norm.running_mean',\n",
        "                                 'enc_content.model.3.model.1.model.0.norm.running_var',\n",
        "                                 'enc_content.model.3.model.2.model.1.norm.running_mean',\n",
        "                                 'enc_content.model.3.model.2.model.1.norm.running_var',\n",
        "                                 'enc_content.model.3.model.2.model.0.norm.running_mean',\n",
        "                                 'enc_content.model.3.model.2.model.0.norm.running_var',\n",
        "                                 'enc_content.model.3.model.3.model.1.norm.running_mean',\n",
        "                                 'enc_content.model.3.model.3.model.1.norm.running_var',\n",
        "                                 'enc_content.model.3.model.3.model.0.norm.running_mean',\n",
        "                                 'enc_content.model.3.model.3.model.0.norm.running_var',\n",
        "                                 )):\n",
        "                    del state_dict[key]\n",
        "        else:\n",
        "            def __conversion_core(state_dict_base):\n",
        "                state_dict = state_dict_base.copy()\n",
        "                for key, value in state_dict_base.items():\n",
        "                    if key.endswith(('enc.model.0.norm.running_mean',\n",
        "                                     'enc.model.0.norm.running_var',\n",
        "                                     'enc.model.1.norm.running_mean',\n",
        "                                     'enc.model.1.norm.running_var',\n",
        "                                     'enc.model.2.norm.running_mean',\n",
        "                                     'enc.model.2.norm.running_var',\n",
        "                                     'enc.model.3.model.0.model.1.norm.running_mean',\n",
        "                                     'enc.model.3.model.0.model.1.norm.running_var',\n",
        "                                     'enc.model.3.model.0.model.0.norm.running_mean',\n",
        "                                     'enc.model.3.model.0.model.0.norm.running_var',\n",
        "                                     'enc.model.3.model.1.model.1.norm.running_mean',\n",
        "                                     'enc.model.3.model.1.model.1.norm.running_var',\n",
        "                                     'enc.model.3.model.1.model.0.norm.running_mean',\n",
        "                                     'enc.model.3.model.1.model.0.norm.running_var',\n",
        "                                     'enc.model.3.model.2.model.1.norm.running_mean',\n",
        "                                     'enc.model.3.model.2.model.1.norm.running_var',\n",
        "                                     'enc.model.3.model.2.model.0.norm.running_mean',\n",
        "                                     'enc.model.3.model.2.model.0.norm.running_var',\n",
        "                                     'enc.model.3.model.3.model.1.norm.running_mean',\n",
        "                                     'enc.model.3.model.3.model.1.norm.running_var',\n",
        "                                     'enc.model.3.model.3.model.0.norm.running_mean',\n",
        "                                     'enc.model.3.model.3.model.0.norm.running_var',\n",
        "\n",
        "                                     'dec.model.0.model.0.model.1.norm.running_mean',\n",
        "                                     'dec.model.0.model.0.model.1.norm.running_var',\n",
        "                                     'dec.model.0.model.0.model.0.norm.running_mean',\n",
        "                                     'dec.model.0.model.0.model.0.norm.running_var',\n",
        "                                     'dec.model.0.model.1.model.1.norm.running_mean',\n",
        "                                     'dec.model.0.model.1.model.1.norm.running_var',\n",
        "                                     'dec.model.0.model.1.model.0.norm.running_mean',\n",
        "                                     'dec.model.0.model.1.model.0.norm.running_var',\n",
        "                                     'dec.model.0.model.2.model.1.norm.running_mean',\n",
        "                                     'dec.model.0.model.2.model.1.norm.running_var',\n",
        "                                     'dec.model.0.model.2.model.0.norm.running_mean',\n",
        "                                     'dec.model.0.model.2.model.0.norm.running_var',\n",
        "                                     'dec.model.0.model.3.model.1.norm.running_mean',\n",
        "                                     'dec.model.0.model.3.model.1.norm.running_var',\n",
        "                                     'dec.model.0.model.3.model.0.norm.running_mean',\n",
        "                                     'dec.model.0.model.3.model.0.norm.running_var',\n",
        "                                     )):\n",
        "                        del state_dict[key]\n",
        "        return state_dict\n",
        "\n",
        "    state_dict = dict()\n",
        "    state_dict['a'] = __conversion_core(state_dict_base['a'], trainer_name)\n",
        "    state_dict['b'] = __conversion_core(state_dict_base['b'], trainer_name)\n",
        "    return state_dict\n",
        "\n"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Hh58FgRpk7EM",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\"\"\"\n",
        "Copyright (C) 2017 NVIDIA Corporation.  All rights reserved.\n",
        "Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n",
        "\"\"\"\n",
        "# from networks import AdaINGen, MsImageDis, MsImageDisCouncil\n",
        "# from utils import weights_init, get_model_list, vgg_preprocess, load_vgg16, get_scheduler\n",
        "from torch.autograd import Variable\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import os\n",
        "import random\n",
        "import threading\n",
        "from multiprocessing.pool import ThreadPool\n",
        "import warnings\n",
        "from collections import deque\n",
        "import numpy as np\n",
        "import torchvision.transforms.functional as TF\n",
        "from scipy import ndimage\n",
        "\n",
        "class Council_Trainer(nn.Module):\n",
        "    def __init__(self, hyperparameters, cuda_device='cuda:0'):\n",
        "        super(Council_Trainer, self).__init__()\n",
        "        lr = hyperparameters['lr']\n",
        "        # Initiate the networks\n",
        "        self.council_size = hyperparameters['council']['council_size']\n",
        "        self.council_size_conf = self.council_size\n",
        "        self.gen_a2b_s = []\n",
        "        self.gen_b2a_s = []\n",
        "        self.dis_a2b_s = []\n",
        "        self.dis_b2a_s = []\n",
        "        self.do_dis_council = hyperparameters['council_w'] != 0\n",
        "        self.do_ads_council_loss = hyperparameters['council_abs_w'] != 0\n",
        "        self.numberOfCouncil_dis_relative_iteration_conf = hyperparameters['council']['numberOfCouncil_dis_relative_iteration']\n",
        "        self.discriminetro_less_style_by_conf = hyperparameters['council']['discriminetro_less_style_by']\n",
        "        self.cuda_device = cuda_device\n",
        "\n",
        "        # all varible with '_conf' at the end will be saved and displayed in tensorboard logs\n",
        "        self.recon_x_w_conf = hyperparameters['recon_x_w']\n",
        "        self.recon_c_w_conf = hyperparameters['recon_c_w']\n",
        "        self.recon_s_w_conf = hyperparameters['recon_s_w']\n",
        "        self.recon_x_cyc_w_conf = hyperparameters['recon_x_cyc_w']\n",
        "        self.gan_w_conf = hyperparameters['gan_w']\n",
        "        self.vgg_w_conf = hyperparameters['vgg_w']\n",
        "        self.abs_beginning_end_w_conf = hyperparameters['abs_beginning_end']\n",
        "        self.flipOnOff_On_iteration_conf = hyperparameters['council']['flipOnOff_On_iteration']\n",
        "        self.flipOnOff_Off_iteration_conf = hyperparameters['council']['flipOnOff_Off_iteration']\n",
        "        self.flipOnOff_Off_iteration_conf = hyperparameters['council']['flipOnOff_start_with']\n",
        "        self.council_abs_w_conf = hyperparameters['council_abs_w']\n",
        "        self.council_w_conf = hyperparameters['council_w']\n",
        "        self.council_start_at_iter_conf = hyperparameters['council']['council_start_at_iter']\n",
        "        self.focus_loss_start_at_iter_conf = hyperparameters['focus_loss']['focus_loss_start_at_iter']\n",
        "        self.mask_zero_or_one_w_conf = hyperparameters['mask_zero_or_one_w']\n",
        "        self.mask_zero_or_one_center_conf = hyperparameters['focus_loss']['mask_zero_or_one_center']\n",
        "        self.mask_zero_or_one_epsilon_conf = hyperparameters['focus_loss']['mask_zero_or_one_epsilon']\n",
        "        self.mask_total_w_conf = hyperparameters['mask_total_w']\n",
        "        self.mask_tv_w_conf = hyperparameters['mask_tv_w']\n",
        "        self.batch_size_conf = hyperparameters['batch_size']\n",
        "        self.do_w_loss_matching = hyperparameters['do_w_loss_matching']\n",
        "        self.do_w_loss_matching_focus = hyperparameters['focus_loss']['do_w_loss_matching_focus']\n",
        "        self.los_matching_hist_size_conf = hyperparameters['loss_matching_hist_size']\n",
        "        self.do_a2b_conf = hyperparameters['do_a2b']\n",
        "        self.do_b2a_conf = hyperparameters['do_b2a']\n",
        "        self.w_match_b2a_conf = 1\n",
        "        self.w_match_a2b_conf = 1\n",
        "        self.w_match_focus_a2b_conf = 1\n",
        "        self.w_match_focus_b2a_conf = 1\n",
        "        self.w_match_focus_zero_one_a2b_conf = 1\n",
        "        self.w_match_focus_zero_one_b2a_conf = 1\n",
        "\n",
        "        if self.do_a2b_conf:\n",
        "            self.los_hist_gan_a2b_s = []\n",
        "            self.los_hist_council_a2b_s = []\n",
        "            self.los_hist_focus_a2b_s = []\n",
        "            self.los_hist_focus_zero_one_a2b_s = []\n",
        "        if self.do_b2a_conf:\n",
        "            self.los_hist_gan_b2a_s = []\n",
        "            self.los_hist_council_b2a_s = []\n",
        "            self.los_hist_focus_b2a_s = []\n",
        "            self.los_hist_focus_zero_one_b2a_s = []\n",
        "\n",
        "        for ind in range(self.council_size):\n",
        "            if self.do_a2b_conf:\n",
        "                self.los_hist_gan_a2b_s.append(deque(np.ones(self.los_matching_hist_size_conf)))\n",
        "                self.los_hist_council_a2b_s.append(deque(np.ones(self.los_matching_hist_size_conf)))\n",
        "                self.los_hist_focus_a2b_s.append(deque(np.ones(self.los_matching_hist_size_conf)))\n",
        "                self.los_hist_focus_zero_one_a2b_s.append(deque(np.ones(self.los_matching_hist_size_conf)))\n",
        "\n",
        "            if self.do_b2a_conf:\n",
        "                self.los_hist_gan_b2a_s.append(deque(np.ones(self.los_matching_hist_size_conf)))\n",
        "                self.los_hist_council_b2a_s.append(deque(np.ones(self.los_matching_hist_size_conf)))\n",
        "                self.los_hist_focus_b2a_s.append(deque(np.ones(self.los_matching_hist_size_conf)))\n",
        "                self.los_hist_focus_zero_one_b2a_s.append(deque(np.ones(self.los_matching_hist_size_conf)))\n",
        "\n",
        "        self.do_council_loss = None\n",
        "\n",
        "        if self.do_dis_council:\n",
        "            self.dis_council_a2b_s = []\n",
        "            self.dis_council_b2a_s = []\n",
        "\n",
        "        # defining all the networks\n",
        "        for i in range(self.council_size):\n",
        "            if self.do_a2b_conf:\n",
        "                self.gen_a2b_s.append(\n",
        "                    AdaINGen(hyperparameters['input_dim_a'], hyperparameters['gen'], cuda_device=self.cuda_device))  # auto-encoder for domain a2b\n",
        "                self.dis_a2b_s.append(\n",
        "                    MsImageDis(hyperparameters['input_dim_a'], hyperparameters['dis'], cuda_device=self.cuda_device))  # discriminator for domain a2b\n",
        "                if self.do_dis_council:\n",
        "                    self.dis_council_a2b_s.append(\n",
        "                        MsImageDisCouncil(hyperparameters['input_dim_a'],\n",
        "                                          hyperparameters['dis'], cuda_device=self.cuda_device))  # council discriminator for domain a2b\n",
        "            if self.do_b2a_conf:\n",
        "                self.gen_b2a_s.append(\n",
        "                    AdaINGen(hyperparameters['input_dim_b'], hyperparameters['gen'], cuda_device=self.cuda_device))  # auto-encoder for domain b\n",
        "                self.dis_b2a_s.append(\n",
        "                    MsImageDis(hyperparameters['input_dim_b'], hyperparameters['dis'], cuda_device=self.cuda_device))  # discriminator for domain b\n",
        "                if self.do_dis_council:\n",
        "                    self.dis_council_b2a_s.append(#\n",
        "                        MsImageDisCouncil(hyperparameters['input_dim_b'],\n",
        "                                          hyperparameters['dis'], cuda_device=self.cuda_device))  # discriminator for domain b\n",
        "\n",
        "        self.instancenorm = nn.InstanceNorm2d(512, affine=False)\n",
        "        self.style_dim = hyperparameters['gen']['style_dim']\n",
        "\n",
        "        if self.do_a2b_conf:\n",
        "            self.gen_a2b_s = nn.ModuleList(self.gen_a2b_s)\n",
        "            self.dis_a2b_s = nn.ModuleList(self.dis_a2b_s)\n",
        "            if self.do_dis_council:\n",
        "                self.dis_council_a2b_s = nn.ModuleList(self.dis_council_a2b_s)\n",
        "        if self.do_b2a_conf:\n",
        "            self.gen_b2a_s = nn.ModuleList(self.gen_b2a_s)\n",
        "            self.dis_b2a_s = nn.ModuleList(self.dis_b2a_s)\n",
        "            if self.do_dis_council:\n",
        "                self.dis_council_b2a_s = nn.ModuleList(self.dis_council_b2a_s)\n",
        "\n",
        "        # fix the noise used in sampling\n",
        "        display_size = int(hyperparameters['display_size'])\n",
        "        self.s_a = torch.randn(display_size, self.style_dim, 1, 1).cuda(self.cuda_device)\n",
        "        self.s_b = torch.randn(display_size, self.style_dim, 1, 1).cuda(self.cuda_device)\n",
        "        # Setup the optimizers\n",
        "        beta1 = hyperparameters['beta1']\n",
        "        beta2 = hyperparameters['beta2']\n",
        "        dis_params_s = []\n",
        "        gen_params_s = []\n",
        "        self.dis_opt_s = []\n",
        "        self.gen_opt_s = []\n",
        "        self.dis_scheduler_s = []\n",
        "        self.gen_scheduler_s = []\n",
        "        if self.do_dis_council:\n",
        "            dis_council_params_s = []\n",
        "            self.dis_council_opt_s = []\n",
        "            self.dis_council_scheduler_s = []\n",
        "        for i in range(self.council_size):\n",
        "            dis_parms = []\n",
        "            gen_parms = []\n",
        "            dis_council_parms = []\n",
        "            if self.do_a2b_conf:\n",
        "                dis_parms += list(self.dis_a2b_s[i].parameters())\n",
        "                gen_parms += list(self.gen_a2b_s[i].parameters())\n",
        "                if self.do_dis_council:\n",
        "                    dis_council_parms += list(self.dis_council_a2b_s[i].parameters())\n",
        "            if self.do_b2a_conf:\n",
        "                dis_parms += list(self.dis_b2a_s[i].parameters())\n",
        "                gen_parms += list(self.gen_b2a_s[i].parameters())\n",
        "                if self.do_dis_council:\n",
        "                    dis_council_parms += list(self.dis_council_b2a_s[i].parameters())\n",
        "            dis_params_s.append(dis_parms)\n",
        "            gen_params_s.append(gen_parms)\n",
        "            if self.do_dis_council:\n",
        "                dis_council_params_s.append(dis_council_parms)\n",
        "            self.dis_opt_s.append(torch.optim.Adam([p for p in dis_params_s[i] if p.requires_grad],\n",
        "                                                   lr=lr, betas=(beta1, beta2),\n",
        "                                                   weight_decay=hyperparameters['weight_decay']))\n",
        "            self.gen_opt_s.append(torch.optim.Adam([p for p in gen_params_s[i] if p.requires_grad],\n",
        "                                                   lr=lr, betas=(beta1, beta2),\n",
        "                                                   weight_decay=hyperparameters['weight_decay']))\n",
        "            if self.do_dis_council:\n",
        "                self.dis_council_opt_s.append(torch.optim.Adam([p for p in dis_council_params_s[i] if p.requires_grad],\n",
        "                                                               lr=lr, betas=(beta1, beta2),\n",
        "                                                               weight_decay=hyperparameters['weight_decay']))\n",
        "            self.dis_scheduler_s.append(get_scheduler(self.dis_opt_s[i], hyperparameters))\n",
        "            self.gen_scheduler_s.append(get_scheduler(self.gen_opt_s[i], hyperparameters))\n",
        "            if self.do_dis_council:\n",
        "                self.dis_council_scheduler_s.append(get_scheduler(self.dis_council_opt_s[i], hyperparameters))\n",
        "\n",
        "        # Network weight initialization\n",
        "        self.apply(weights_init(hyperparameters['init']))\n",
        "        for i in range(self.council_size):\n",
        "            if self.do_a2b_conf:\n",
        "                self.gen_a2b_s[i].apply(weights_init(hyperparameters['init']))\n",
        "                self.dis_a2b_s[i].apply(weights_init('gaussian'))\n",
        "                if self.do_dis_council:\n",
        "                    self.dis_council_a2b_s[i].apply(weights_init('gaussian'))\n",
        "            if self.do_b2a_conf:\n",
        "                self.gen_b2a_s[i].apply(weights_init(hyperparameters['init']))\n",
        "                self.dis_b2a_s[i].apply(weights_init('gaussian'))\n",
        "                if self.do_dis_council:\n",
        "                    self.dis_council_b2a_s[i].apply(weights_init('gaussian'))\n",
        "\n",
        "        # Load VGG model if needed\n",
        "        self.vgg = None\n",
        "        if 'vgg_w' in hyperparameters.keys() and hyperparameters['vgg_w'] > 0:\n",
        "            self.vgg = load_vgg16(hyperparameters['vgg_model_path'] + '/models')\n",
        "            self.vgg.eval()\n",
        "            for param in self.vgg.parameters():\n",
        "                param.requires_grad = False\n",
        "\n",
        "    def recon_criterion(self, input, target):\n",
        "        return torch.mean(torch.abs(input - target))\n",
        "\n",
        "    def recon_criterion_v2_color(self, input, target):\n",
        "        loss_1 = torch.mean(torch.abs(input - target))\n",
        "        loss_2 = torch.mean(torch.pow((input - target), 2))\n",
        "        if loss_1 > loss_2:\n",
        "            return loss_1\n",
        "        return loss_2\n",
        "\n",
        "    def recon_criterion_v3_gray_scale(self, input, target):\n",
        "        loss_1 = torch.mean(torch.abs(torch.sum(input, 1) - torch.sum(target, 1)))\n",
        "        loss_2 = torch.mean(torch.pow((torch.sum(input, 1) - torch.sum(target, 1)), 2))\n",
        "        if loss_1 > loss_2:\n",
        "            return loss_1\n",
        "        return loss_2\n",
        "\n",
        "    def council_basic_criterion_gray_scale(self, input, target):\n",
        "        return torch.mean(torch.abs(torch.sum(input, 1) - torch.sum(target, 1)))\n",
        "\n",
        "    def council_basic_criterion_with_color(self, input, target):\n",
        "        return torch.mean(torch.abs(input - target))\n",
        "\n",
        "    def mask_zero_one_criterion(self, mask, center=0.5, epsilon=0.01):\n",
        "        return torch.sum(1 / (torch.abs(mask - center) + epsilon)) / mask.numel()\n",
        "\n",
        "    def mask_small_criterion(self, mask):\n",
        "        assert self.hyperparameters['focus_loss']['mask_small_use_abs'] or self.hyperparameters['focus_loss']['mask_small_use_square'], 'at leas one small mask loss should be true, mask_small_use_abs or mask_small_use_square'\n",
        "        loss = 0\n",
        "        if self.hyperparameters['focus_loss']['mask_small_use_abs']:\n",
        "            loss += self.mask_small_criterion_abs(mask)\n",
        "        if self.hyperparameters['focus_loss']['mask_small_use_square']:\n",
        "            loss += self.mask_small_criterion_square(mask)\n",
        "        return loss\n",
        "\n",
        "    def mask_small_criterion_square(self, mask):\n",
        "        return (torch.sum(mask) / mask.numel()) ** 2\n",
        "\n",
        "    def mask_small_criterion_abs(self, mask):\n",
        "        return torch.abs((torch.sum(mask))) / mask.numel()\n",
        "\n",
        "    def mask_criterion_TV(self, mask):\n",
        "        return (torch.sum(torch.abs(mask[:, :, 1:, :]-mask[:, :, :-1, :])) + \\\n",
        "               torch.sum(torch.abs(mask[:, :, :, 1:] - mask[:, :, :, :-1]))) / mask.numel()\n",
        "\n",
        "    def forward(self, x_a, x_b=None, s_a=None, s_b=None):\n",
        "        self.eval()\n",
        "        if self.do_a2b_conf:\n",
        "            s_b = self.s_b if s_b is None else s_b\n",
        "            s_b = Variable(s_b)\n",
        "            x_ab_s = []\n",
        "        if self.do_b2a_conf:\n",
        "            s_a = self.s_a if s_a is None else s_a\n",
        "            s_a = Variable(s_a)\n",
        "            x_ba_s = []\n",
        "        for i in range(self.council_size):\n",
        "            if self.do_a2b_conf:\n",
        "                c_a, s_a_fake = self.gen_a2b.encode(x_a)\n",
        "                x_ab_s.append(self.gen_a2b.decode(c_a, s_b, x_a))\n",
        "            if self.do_b2a_conf:\n",
        "                c_b, s_b_fake = self.gen_b2a.encode(x_b)\n",
        "                x_ba_s.append(self.gen_b2a.decode(c_b, s_a, x_b))\n",
        "\n",
        "        self.train()\n",
        "        if self.do_a2b_conf and self.do_b2a_conf:\n",
        "            return x_ab_s, x_ba_s\n",
        "        elif self.do_b2a_conf:\n",
        "            return x_ba_s\n",
        "        return x_ab_s\n",
        "\n",
        "    def gen_update(self, x_a, x_b, hyperparameters, iterations=0):\n",
        "        self.hyperparameters = hyperparameters\n",
        "        for gen_opt in self.gen_opt_s:\n",
        "            gen_opt.zero_grad()\n",
        "        s_a = Variable(torch.randn(x_a.size(0), self.style_dim, 1, 1).cuda(self.cuda_device))\n",
        "        s_b = Variable(torch.randn(x_b.size(0), self.style_dim, 1, 1).cuda(self.cuda_device))\n",
        "        c_a_s = []\n",
        "        s_a_prime_s = []\n",
        "        c_b_s = []\n",
        "        s_b_prime_s = []\n",
        "        x_a_recon_s = []\n",
        "        x_b_recon_s = []\n",
        "        x_ba_s = []\n",
        "        x_ab_s = []\n",
        "        c_b_recon_s = []\n",
        "        s_a_recon_s = []\n",
        "        c_a_recon_s = []\n",
        "        s_b_recon_s = []\n",
        "        x_aba_s = []\n",
        "        x_bab_s = []\n",
        "        mask_ba_s = []\n",
        "        mask_ab_s = []\n",
        "        self.loss_gen_mask_zero_one_ba_s = []\n",
        "        self.loss_gen_mask_zero_one_ab_s = []\n",
        "        self.loss_gen_mask_total_ba_s = []\n",
        "        self.loss_gen_mask_total_ab_s = []\n",
        "        self.loss_gen_mask_TV_ab_s = []\n",
        "        self.loss_gen_mask_TV_ba_s = []\n",
        "        self.loss_gen_recon_x_a_s = []\n",
        "        self.loss_gen_recon_x_b_s = []\n",
        "        self.loss_gen_recon_s_a_s = []\n",
        "        self.loss_gen_recon_s_b_s = []\n",
        "        self.loss_gen_recon_c_a_s = []\n",
        "        self.loss_gen_recon_c_b_s = []\n",
        "        self.loss_gen_cycrecon_x_a_s = []\n",
        "        self.loss_gen_cycrecon_x_b_s = []\n",
        "        self.loss_gen_beginning_end_a_ab_s = []\n",
        "        self.loss_gen_beginning_end_b_ba_s = []\n",
        "        self.loss_gen_adv_a2b_s = []\n",
        "        self.loss_gen_adv_b2a_s = []\n",
        "        self.loss_gen_vgg_a_s = []\n",
        "        self.loss_gen_vgg_b_s = []\n",
        "        self.loss_gen_total_s = []\n",
        "        self.council_w_conf = hyperparameters['council_w'] if hyperparameters['iteration'] > hyperparameters['council']['council_start_at_iter'] else 0\n",
        "        self.mask_zero_or_one_w_conf = hyperparameters['mask_zero_or_one_w'] if hyperparameters['iteration'] > hyperparameters['focus_loss']['focus_loss_start_at_iter'] else 0\n",
        "        self.mask_total_w_conf = hyperparameters['mask_total_w'] if hyperparameters['iteration'] > hyperparameters['focus_loss']['focus_loss_start_at_iter'] else 0\n",
        "        self.mask_tv_w_conf = hyperparameters['mask_tv_w'] if hyperparameters['iteration'] > hyperparameters['focus_loss']['focus_loss_start_at_iter'] else 0\n",
        "\n",
        "        for i in range(self.council_size):\n",
        "            # encode\n",
        "            if self.do_a2b_conf:\n",
        "                c_a, s_a_prime = self.gen_a2b_s[i].encode(x_a)\n",
        "                c_a_s.append(c_a)\n",
        "                s_a_prime_s.append(s_a_prime)\n",
        "            if self.do_b2a_conf:\n",
        "                c_b, s_b_prime = self.gen_b2a_s[i].encode(x_b)\n",
        "                c_b_s.append(c_b)\n",
        "                s_b_prime_s.append(s_b_prime)\n",
        "\n",
        "            # decode (within domain)\n",
        "            if hyperparameters['recon_x_w'] != 0:\n",
        "                if not self.do_a2b_conf and not self.do_b2a_conf:\n",
        "                    print('cant do recon_x loss if not both do_a2b and b2a set to true')\n",
        "                else:\n",
        "                    x_a_recon_s.append(self.gen_b2a_s[i].decode(c_a_s[i], s_a_prime_s[i], x_a))\n",
        "                    x_b_recon_s.append(self.gen_a2b_s[i].decode(c_b_s[i], s_b_prime_s[i], x_b))\n",
        "\n",
        "            # decode (cross domain)\n",
        "            if self.do_a2b_conf:\n",
        "                x_ab_s.append(self.gen_a2b_s[i].decode(c_a_s[i], s_b, x_a))\n",
        "            if self.do_b2a_conf:\n",
        "                x_ba_s.append(self.gen_b2a_s[i].decode(c_b_s[i], s_a, x_b))\n",
        "\n",
        "            if hyperparameters['mask_zero_or_one_w'] != 0 or hyperparameters['mask_total_w'] != 0:\n",
        "                if self.do_a2b_conf:\n",
        "                    mask_ab_s.append(self.gen_a2b_s[i].dec.mask_s)\n",
        "                if self.do_b2a_conf:\n",
        "                    mask_ba_s.append(self.gen_b2a_s[i].dec.mask_s)\n",
        "\n",
        "            # encode again\n",
        "            if hyperparameters['recon_s_w'] != 0 or hyperparameters['recon_c_w'] != 0 or hyperparameters['recon_x_cyc_w'] != 0:\n",
        "                if not self.do_a2b_conf and not self.do_b2a_conf:\n",
        "                    print('cant do recon_s and recon_c loss if not both do_a2b and b2a set to true')\n",
        "                else:\n",
        "                    c_b_recon, s_a_recon = self.gen_a2b_s[i].encode(x_ba_s[i])\n",
        "                    c_a_recon, s_b_recon = self.gen_b2a_s[i].encode(x_ab_s[i])\n",
        "                    c_b_recon_s.append(c_b_recon)\n",
        "                    s_a_recon_s.append(s_a_recon)\n",
        "                    c_a_recon_s.append(c_a_recon)\n",
        "                    s_b_recon_s.append(s_b_recon)\n",
        "\n",
        "            # decode again (if needed)\n",
        "            if hyperparameters['recon_x_cyc_w'] != 0:\n",
        "                if not self.do_a2b_conf and not self.do_b2a_conf:\n",
        "                    print('cant do recon_x_cyc loss if not both do_a2b and b2a set to true')\n",
        "                else:\n",
        "                    x_aba_s.append(\n",
        "                        self.gen_b2a.decode(c_a_recon_s[i], s_a_prime_s[i], x_a) if hyperparameters['recon_x_cyc_w'] > 0 else None)\n",
        "                    x_bab_s.append(\n",
        "                        self.gen_a2b.decode(c_b_recon_s[i], s_b_prime_s[i], x_b) if hyperparameters['recon_x_cyc_w'] > 0 else None)\n",
        "\n",
        "            self.loss_gen_total_s.append(0)\n",
        "            if hyperparameters['do_a2b']:\n",
        "                self.loss_gen_mask_TV_ab_s.append(0)\n",
        "                self.loss_gen_mask_total_ab_s.append(0)\n",
        "            if hyperparameters['do_b2a']:\n",
        "                self.loss_gen_mask_TV_ba_s.append(0)\n",
        "                self.loss_gen_mask_total_ba_s.append(0)\n",
        "\n",
        "            # masks should contain ones or zeros\n",
        "            if hyperparameters['iteration'] > hyperparameters['focus_loss']['focus_loss_start_at_iter'] and (hyperparameters['mask_zero_or_one_w'] != 0 or hyperparameters['mask_total_w'] != 0 ):\n",
        "\n",
        "                if hyperparameters['mask_zero_or_one_w'] != 0:\n",
        "                    if hyperparameters['do_a2b']:\n",
        "                        self.loss_gen_mask_zero_one_ab_s.append(self.mask_zero_one_criterion(mask_ab_s[i], center=hyperparameters['focus_loss']['mask_zero_or_one_center'], epsilon=hyperparameters['focus_loss']['mask_zero_or_one_epsilon']))\n",
        "                    if hyperparameters['do_b2a']:\n",
        "                        self.loss_gen_mask_zero_one_ba_s.append(self.mask_zero_one_criterion(mask_ba_s[i], center=hyperparameters['focus_loss']['mask_zero_or_one_center'], epsilon=hyperparameters['focus_loss']['mask_zero_or_one_epsilon']))\n",
        "\n",
        "                    if self.do_w_loss_matching_focus:\n",
        "                        if hyperparameters['do_a2b']:\n",
        "                            self.los_hist_focus_zero_one_a2b_s[i].append(self.loss_gen_mask_zero_one_ab_s[i].detach().cpu().numpy())\n",
        "                            self.los_hist_focus_zero_one_a2b_s[i].popleft()\n",
        "                            self.w_match_focus_zero_one_a2b_conf = np.mean(self.los_hist_gan_a2b_s[i]) / np.mean(self.los_hist_focus_zero_one_a2b_s[i])\n",
        "                            self.loss_gen_mask_zero_one_ab_s[i] *= self.w_match_focus_zero_one_a2b_conf\n",
        "                            self.loss_gen_total_s[i] += hyperparameters['mask_zero_or_one_w'] * self.loss_gen_mask_zero_one_ab_s[i].cuda(self.cuda_device)\n",
        "                        if hyperparameters['do_b2a']:\n",
        "                            self.los_hist_focus_zero_one_b2a_s[i].append(self.loss_gen_mask_zero_one_ba_s[i].detach().cpu().numpy())\n",
        "                            self.los_hist_focus_zero_one_b2a_s[i].popleft()\n",
        "                            self.w_match_focus_zero_one_b2a_conf = np.mean(self.los_hist_gan_b2a_s[i]) / np.mean(self.los_hist_focus_zero_one_b2a_s[i])\n",
        "                            self.loss_gen_mask_zero_one_ba_s[i] *= self.w_match_focus_zero_one_b2a_conf\n",
        "                            self.loss_gen_total_s[i] += hyperparameters['mask_zero_or_one_w'] * self.loss_gen_mask_zero_one_ba_s[i].cuda(self.cuda_device)\n",
        "                    else:\n",
        "                        if hyperparameters['do_a2b']:\n",
        "                            self.loss_gen_total_s[i] += hyperparameters['mask_zero_or_one_w'] * self.loss_gen_mask_zero_one_ab_s[i].cuda(self.cuda_device)\n",
        "                        if hyperparameters['do_b2a']:\n",
        "                            self.loss_gen_total_s[i] += hyperparameters['mask_zero_or_one_w'] * self.loss_gen_mask_zero_one_ba_s[i].cuda(self.cuda_device)\n",
        "\n",
        "                # masks should as small as possible to leave original domain with little changes\n",
        "                if hyperparameters['mask_total_w'] != 0:\n",
        "                    if hyperparameters['do_a2b']:\n",
        "                        self.loss_gen_mask_total_ab_s[i] += self.mask_small_criterion(mask_ab_s[i])\n",
        "                    if hyperparameters['do_b2a']:\n",
        "                        self.loss_gen_mask_total_ba_s[i] += self.mask_small_criterion(mask_ba_s[i])\n",
        "\n",
        "                # TV loss on the mask\n",
        "                if hyperparameters['mask_tv_w'] != 0:\n",
        "                    if hyperparameters['do_a2b']:\n",
        "                        self.loss_gen_mask_TV_ab_s[i] += self.mask_criterion_TV(mask_ab_s[i])\n",
        "                        self.loss_gen_total_s[i] += hyperparameters['mask_tv_w'] * self.loss_gen_mask_TV_ab_s[i]\n",
        "                    if hyperparameters['do_b2a']:\n",
        "                        self.loss_gen_mask_TV_ba_s[i] += self.mask_criterion_TV(mask_ba_s[i])\n",
        "                        self.loss_gen_total_s[i] += hyperparameters['mask_tv_w'] * self.loss_gen_mask_TV_ba_s[i]\n",
        "\n",
        "                if self.do_w_loss_matching_focus:\n",
        "                    if hyperparameters['do_a2b']:\n",
        "                        self.los_hist_focus_a2b_s[i].append(self.loss_gen_mask_total_ab_s[i].detach().cpu().numpy())\n",
        "                        self.los_hist_focus_a2b_s[i].popleft()\n",
        "                        self.w_match_focus_a2b_conf = np.mean(self.los_hist_gan_a2b_s[i]) / np.mean(self.los_hist_focus_a2b_s[i])\n",
        "                        self.loss_gen_mask_total_ab_s[i] *= self.w_match_focus_a2b_conf\n",
        "                        self.loss_gen_total_s[i] += hyperparameters['mask_total_w'] * self.loss_gen_mask_total_ab_s[i].cuda(self.cuda_device)\n",
        "                    if hyperparameters['do_b2a']:\n",
        "                        self.los_hist_focus_b2a_s[i].append(self.loss_gen_mask_total_ab_s[i].detach().cpu().numpy())\n",
        "                        self.los_hist_focus_b2a_s[i].popleft()\n",
        "                        self.w_match_focus_b2a_conf = np.mean(self.los_hist_gan_b2a_s[i]) / np.mean(self.los_hist_focus_b2a_s[i])\n",
        "                        self.loss_gen_mask_total_ba_s[i] *= self.w_match_focus_b2a_conf\n",
        "                        self.loss_gen_total_s[i] += hyperparameters['mask_total_w'] * self.loss_gen_mask_total_ba_s[i].cuda(self.cuda_device)\n",
        "\n",
        "                else:\n",
        "                    if hyperparameters['do_a2b']:\n",
        "                        self.loss_gen_total_s[i] += hyperparameters['mask_total_w'] * self.loss_gen_mask_total_ab_s[i].cuda(self.cuda_device)\n",
        "                    if hyperparameters['do_b2a']:\n",
        "                        self.loss_gen_total_s[i] += hyperparameters['mask_total_w'] * self.loss_gen_mask_total_ba_s[i].cuda(self.cuda_device)\n",
        "\n",
        "\n",
        "            # reconstruction loss\n",
        "            if hyperparameters['recon_x_w'] != 0 and self.do_a2b_conf and self.do_b2a_conf:\n",
        "                self.loss_gen_recon_x_a_s.append(self.recon_criterion(x_a_recon_s[i], x_a))\n",
        "                self.loss_gen_recon_x_b_s.append(self.recon_criterion(x_b_recon_s[i], x_b))\n",
        "                self.loss_gen_total_s[i] += hyperparameters['recon_x_w'] * (\n",
        "                        self.loss_gen_recon_x_a_s[i].cuda(self.cuda_device) + self.loss_gen_recon_x_b_s[i].cuda(self.cuda_device))\n",
        "            if hyperparameters['recon_s_w'] != 0 and self.do_a2b_conf and self.do_b2a_conf:\n",
        "                self.loss_gen_recon_s_a_s.append(self.recon_criterion(s_a_recon_s[i], s_a))\n",
        "                self.loss_gen_recon_s_b_s.append(self.recon_criterion(s_b_recon_s[i], s_b))\n",
        "                self.loss_gen_total_s[i] += hyperparameters['recon_s_w'] * (\n",
        "                        self.loss_gen_recon_s_a_s[i].cuda(self.cuda_device) + self.loss_gen_recon_s_b_s[i].cuda(self.cuda_device))\n",
        "            if hyperparameters['recon_c_w'] != 0 and self.do_a2b_conf and self.do_b2a_conf:\n",
        "                self.loss_gen_recon_c_a_s.append(self.recon_criterion(c_a_recon_s[i], c_a_s[i]))\n",
        "                self.loss_gen_recon_c_b_s.append(self.recon_criterion(c_b_recon_s[i], c_b_s[i]))\n",
        "                self.loss_gen_total_s[i] += hyperparameters['recon_c_w'] * (\n",
        "                        self.loss_gen_recon_c_a_s[i].cuda(self.cuda_device) + self.loss_gen_recon_c_b_s[i].cuda(self.cuda_device))\n",
        "            if hyperparameters['recon_x_cyc_w'] != 0 and self.do_a2b_conf and self.do_b2a_conf:\n",
        "                self.loss_gen_cycrecon_x_a_s.append(\n",
        "                    self.recon_criterion(x_aba_s[i], x_a) if hyperparameters['recon_x_cyc_w'] > 0 else 0)\n",
        "                self.loss_gen_cycrecon_x_b_s.append(\n",
        "                    self.recon_criterion(x_bab_s[i], x_b) if hyperparameters['recon_x_cyc_w'] > 0 else 0)\n",
        "                self.loss_gen_total_s[i] += hyperparameters['recon_x_cyc_w'] * (\n",
        "                        self.loss_gen_cycrecon_x_a_s[i].cuda(self.cuda_device) + self.loss_gen_cycrecon_x_b_s[i].cuda(self.cuda_device))\n",
        "            if hyperparameters['abs_beginning_end'] != 0 and self.abs_beginning_end_w_conf > 0.005:\n",
        "                if hyperparameters['do_a2b']:\n",
        "                    self.loss_gen_beginning_end_a_ab_s.append(\n",
        "                        self.recon_criterion_v2_color(x_ab_s[i], x_a) if hyperparameters['abs_beginning_end'] > 0 or hyperparameters['abs_beginning_end_minimume'] > 0 else 0)\n",
        "                else:\n",
        "                    self.loss_gen_beginning_end_a_ab_s.append(0)\n",
        "                if hyperparameters['do_b2a']:\n",
        "                    self.loss_gen_beginning_end_b_ba_s.append(\n",
        "                        self.recon_criterion_v2_color(x_ba_s[i], x_b) if hyperparameters['abs_beginning_end'] > 0 or hyperparameters['abs_beginning_end_minimume'] > 0 else 0)\n",
        "                else:\n",
        "                    self.loss_gen_beginning_end_b_ba_s.append(0)\n",
        "\n",
        "                self.abs_beginning_end_w_conf = hyperparameters['abs_beginning_end'] * (hyperparameters['abs_beginning_end_less_by'] ** iterations)\n",
        "                self.abs_beginning_end_w_conf = max(self.abs_beginning_end_w_conf, hyperparameters['abs_beginning_end_minimume'])\n",
        "\n",
        "                if hyperparameters['do_a2b']:\n",
        "                    self.loss_gen_total_s[i] += self.abs_beginning_end_w_conf * self.loss_gen_beginning_end_a_ab_s[i].cuda(self.cuda_device)\n",
        "                if hyperparameters['do_b2a']:\n",
        "                    self.loss_gen_total_s[i] += self.abs_beginning_end_w_conf * self.loss_gen_beginning_end_b_ba_s[i].cuda(self.cuda_device)\n",
        "\n",
        "            # GAN loss\n",
        "            if hyperparameters['gan_w'] != 0:\n",
        "                i_dis = i\n",
        "                if hyperparameters['gen']['useRandomDis']:\n",
        "                    i_dis = np.random.randint(self.council_size)\n",
        "\n",
        "                if hyperparameters['do_a2b']:\n",
        "                    x_ab_s_curr = x_ab_s[i] if not hyperparameters['dis']['do_Dis_only_gray'] else torch.sum(x_ab_s[i], 1).unsqueeze(1).repeat(1, hyperparameters['input_dim_b'], 1, 1) / hyperparameters['input_dim_b']\n",
        "                    loss_gen_adv_a2b = self.dis_a2b_s[i_dis].calc_gen_loss(x_ab_s_curr)\n",
        "                else:\n",
        "                    loss_gen_adv_a2b = 0\n",
        "\n",
        "                if hyperparameters['do_b2a']:\n",
        "                    x_ba_s_curr = x_ba_s[i] if not hyperparameters['dis']['do_Dis_only_gray'] else torch.sum(x_ba_s[i], 1).unsqueeze(1).repeat(1, hyperparameters['input_dim_a'], 1, 1) / hyperparameters['input_dim_a']\n",
        "                    loss_gen_adv_b2a = self.dis_b2a_s[i_dis].calc_gen_loss(x_ba_s_curr)\n",
        "                else:\n",
        "                    loss_gen_adv_b2a = 0\n",
        "\n",
        "                self.loss_gen_adv_a2b_s.append(loss_gen_adv_a2b)\n",
        "                self.loss_gen_adv_b2a_s.append(loss_gen_adv_b2a)\n",
        "\n",
        "                if self.do_w_loss_matching:\n",
        "                    if hyperparameters['do_a2b']:\n",
        "                        self.los_hist_gan_a2b_s[i].append(loss_gen_adv_a2b.detach().cpu().numpy())\n",
        "                        self.los_hist_gan_a2b_s[i].popleft()\n",
        "                    if hyperparameters['do_b2a']:\n",
        "                        self.los_hist_gan_b2a_s[i].append(loss_gen_adv_b2a.detach().cpu().numpy())\n",
        "                        self.los_hist_gan_b2a_s[i].popleft()\n",
        "\n",
        "                if hyperparameters['do_a2b']:\n",
        "                    self.loss_gen_total_s[i] += hyperparameters['gan_w'] * self.loss_gen_adv_a2b_s[i].cuda(self.cuda_device)\n",
        "                if hyperparameters['do_b2a']:\n",
        "                    self.loss_gen_total_s[i] += hyperparameters['gan_w'] * self.loss_gen_adv_b2a_s[i].cuda(self.cuda_device)\n",
        "\n",
        "            # domain-invariant perceptual loss\n",
        "            if hyperparameters['vgg_w'] != 0:\n",
        "                self.loss_gen_vgg_a_s.append(\n",
        "                    self.compute_vgg_loss(self.vgg, x_ba_s[i], x_b) if hyperparameters['vgg_w'] > 0 else 0)\n",
        "                self.loss_gen_vgg_b_s.append(\n",
        "                    self.compute_vgg_loss(self.vgg, x_ab_s[i], x_a) if hyperparameters['vgg_w'] > 0 else 0)\n",
        "                self.loss_gen_total_s[i] += hyperparameters['vgg_w'] * (\n",
        "                        self.loss_gen_vgg_a_s[i].cuda(self.cuda_device) + self.loss_gen_vgg_b_s[i].cuda(self.cuda_device))\n",
        "\n",
        "        # Council loss\n",
        "        onOffCycle = hyperparameters['council']['flipOnOff_On_iteration'] + hyperparameters['council'][\n",
        "            'flipOnOff_Off_iteration']\n",
        "        currIterCyc = hyperparameters['iteration'] % onOffCycle\n",
        "        if hyperparameters['council']['flipOnOff_start_with']:\n",
        "            startCyc = hyperparameters['council']['flipOnOff_On_iteration']\n",
        "        else:\n",
        "            startCyc = hyperparameters['council']['flipOnOff_Off_iteration']\n",
        "\n",
        "        self.do_council_loss = hyperparameters['council']['flipOnOff_start_with'] if (currIterCyc < startCyc) \\\n",
        "            else not hyperparameters['council']['flipOnOff_start_with']\n",
        "\n",
        "        if not hyperparameters['council']['flipOnOff']:\n",
        "            self.do_council_loss = True\n",
        "        if hyperparameters['iteration'] < hyperparameters['council']['council_start_at_iter']:\n",
        "            self.do_council_loss = False\n",
        "        self.council_loss_ba_s = []\n",
        "        self.council_loss_ab_s = []\n",
        "        for i in range(self.council_size):\n",
        "            if (hyperparameters['council_w'] != 0 or hyperparameters['council_abs_w'] != 0) and self.do_council_loss and self.council_size > 1:\n",
        "                # if i == 0:\n",
        "                #     print('do council loss: True')\n",
        "                if self.do_a2b_conf:\n",
        "                    self.council_loss_ab_s.append(0)\n",
        "                if self.do_b2a_conf:\n",
        "                    self.council_loss_ba_s.append(0)\n",
        "\n",
        "                if self.do_dis_council:  # do council discriminator\n",
        "                    if hyperparameters['do_a2b']:\n",
        "                        dis_council_loss_ab = self.dis_council_a2b_s[i].calc_gen_loss(x_ab_s[i], x_a)\n",
        "                    else:\n",
        "                        dis_council_loss_ab = 0\n",
        "                    if hyperparameters['do_b2a']:\n",
        "                        dis_council_loss_ba = self.dis_council_b2a_s[i].calc_gen_loss(x_ba_s[i], x_b)\n",
        "                    else:\n",
        "                        dis_council_loss_ba = 0\n",
        "                    if self.do_w_loss_matching:\n",
        "                        if hyperparameters['do_a2b']:\n",
        "                            self.los_hist_council_a2b_s[i].append(dis_council_loss_ab.detach().cpu().numpy())\n",
        "                            self.los_hist_council_a2b_s[i].popleft()\n",
        "                            self.w_match_a2b_conf = np.mean(self.los_hist_gan_a2b_s[i]) / np.mean(self.los_hist_council_a2b_s[i])\n",
        "                            dis_council_loss_ab *= self.w_match_a2b_conf\n",
        "                        if hyperparameters['do_b2a']:\n",
        "                            self.los_hist_council_b2a_s[i].append(dis_council_loss_ba.detach().cpu().numpy())\n",
        "                            self.los_hist_council_b2a_s[i].popleft()\n",
        "                            self.w_match_b2a_conf = np.mean(self.los_hist_gan_b2a_s[i]) / np.mean(self.los_hist_council_b2a_s[i])\n",
        "                            dis_council_loss_ba *= self.w_match_b2a_conf\n",
        "\n",
        "                    if hyperparameters['do_a2b']:\n",
        "                        dis_council_loss_ab *= hyperparameters['council_w']\n",
        "                        self.council_loss_ab_s[i] += dis_council_loss_ab\n",
        "                    if hyperparameters['do_b2a']:\n",
        "                        dis_council_loss_ba *= hyperparameters['council_w']\n",
        "                        self.council_loss_ba_s[i] += dis_council_loss_ba\n",
        "\n",
        "                if hyperparameters['council_abs_w'] != 0 and self.council_size > 1:  # ads loss without discriminetor\n",
        "                    tmp = list(range(0, i)) + list(range(i + 1, self.council_size))\n",
        "                    comper_to_i = random.choice(tmp)\n",
        "                    if hyperparameters['council_abs_gray_scale']:\n",
        "                        if hyperparameters['do_a2b']:\n",
        "                            abs_council_loss_ab = hyperparameters['council_abs_w'] * self.council_basic_criterion_gray_scale(x_ab_s[i], x_ab_s[comper_to_i].detach())\n",
        "                        else:\n",
        "                            abs_council_loss_ab = 0\n",
        "                        if hyperparameters['do_b2a']:\n",
        "                            abs_council_loss_ba = hyperparameters['council_abs_w'] * self.council_basic_criterion_gray_scale(x_ba_s[i], x_ba_s[comper_to_i].detach())\n",
        "                        else:\n",
        "                            abs_council_loss_ba = 0\n",
        "                    else:\n",
        "                        if hyperparameters['do_a2b']:\n",
        "                            abs_council_loss_ab = hyperparameters['council_abs_w'] * self.council_basic_criterion_with_color(x_ab_s[i], x_ab_s[comper_to_i].detach())\n",
        "                        else:\n",
        "                            abs_council_loss_ab = 0\n",
        "                        if hyperparameters['do_b2a']:\n",
        "                            abs_council_loss_ba = hyperparameters['council_abs_w'] * self.council_basic_criterion_with_color(x_ba_s[i], x_ba_s[comper_to_i].detach())\n",
        "                        else:\n",
        "                            abs_council_loss_ba = 0\n",
        "                    if self.do_a2b_conf:\n",
        "                        self.council_loss_ab_s[i] += abs_council_loss_ba.cuda(self.cuda_device)\n",
        "                    if self.do_b2a_conf:\n",
        "                        self.council_loss_ba_s[i] += abs_council_loss_ab.cuda(self.cuda_device)\n",
        "\n",
        "                if hyperparameters['do_a2b']:\n",
        "                    self.loss_gen_total_s[i] += self.council_loss_ab_s[i].cuda(self.cuda_device)\n",
        "                if hyperparameters['do_b2a']:\n",
        "                        self.loss_gen_total_s[i] += self.council_loss_ba_s[i].cuda(self.cuda_device)\n",
        "\n",
        "            else:\n",
        "                if self.do_a2b_conf:\n",
        "                    self.council_loss_ab_s.append(0)\n",
        "                if self.do_b2a_conf:\n",
        "                    self.council_loss_ba_s.append(0)\n",
        "\n",
        "            # backpropogation\n",
        "            self.loss_gen_total_s[i].backward()\n",
        "            self.gen_opt_s[i].step()\n",
        "\n",
        "    def compute_vgg_loss(self, vgg, img, target):\n",
        "        img_vgg = vgg_preprocess(img)\n",
        "        target_vgg = vgg_preprocess(target)\n",
        "        img_fea = vgg(img_vgg)\n",
        "        target_fea = vgg(target_vgg)\n",
        "        return torch.mean((self.instancenorm(img_fea) - self.instancenorm(target_fea)) ** 2)\n",
        "\n",
        "    def sample(self, x_a=None, x_b=None, s_a=None, s_b=None, council_member_to_sample_vec=None, return_mask=True):\n",
        "        self.eval()\n",
        "        if self.do_a2b_conf:\n",
        "            x_a_s = []\n",
        "            s_b = self.s_b if s_b is None else s_b\n",
        "            s_b1 = Variable(s_b)\n",
        "            s_b2 = Variable(torch.randn(x_a.size(0), self.style_dim, 1, 1).cuda(self.cuda_device))\n",
        "            x_a_recon, x_ab1, x_ab2, x_ab1_mask = [], [], [], []\n",
        "        if self.do_b2a_conf:\n",
        "            x_b_s = []\n",
        "            s_a = self.s_a if s_a is None else s_a\n",
        "            s_a1 = Variable(s_a)\n",
        "            s_a2 = Variable(torch.randn(x_b.size(0), self.style_dim, 1, 1).cuda(self.cuda_device))\n",
        "            x_b_recon, x_ba1, x_ba2, x_ba1_mask = [], [], [], []\n",
        "\n",
        "        council_member_to_sample_vec = range(self.council_size) if council_member_to_sample_vec is None else council_member_to_sample_vec\n",
        "        x_size = x_a.size(0) if x_a is not None else x_b.size(0)\n",
        "        for i in range(x_size):\n",
        "            for j in council_member_to_sample_vec:\n",
        "                if self.do_b2a_conf:\n",
        "                    x_b_s.append(x_b[i].unsqueeze(0))\n",
        "                    c_b, s_b_fake = self.gen_b2a_s[j].encode(x_b[i].unsqueeze(0))\n",
        "                    if not return_mask:\n",
        "                        x_b_recon.append(self.gen_b2a_s[j].decode(c_b, s_b_fake, x_b[i].unsqueeze(0)))\n",
        "                        x_ba1.append(self.gen_b2a_s[j].decode(c_b, s_a1[i].unsqueeze(0), x_b[i].unsqueeze(0)))\n",
        "                        x_ba2.append(self.gen_b2a_s[j].decode(c_b, s_a2[i].unsqueeze(0), x_b[i].unsqueeze(0)))\n",
        "                    else:\n",
        "                        x_ba1_tmp, x_ba1_mask_tmp = self.gen_b2a_s[j].decode(c_b, s_a1[i].unsqueeze(0), x_b[i].unsqueeze(0), return_mask=return_mask)\n",
        "                        x_ba1_mask.append(x_ba1_mask_tmp)\n",
        "                        x_ba1.append(x_ba1_tmp)\n",
        "                        x_ba2.append(self.gen_b2a_s[j].decode(c_b, s_a2[i].unsqueeze(0), x_b[i].unsqueeze(0)))\n",
        "                if self.do_a2b_conf:\n",
        "                    x_a_s.append(x_a[i].unsqueeze(0))\n",
        "                    c_a, s_a_fake = self.gen_a2b_s[j].encode(x_a[i].unsqueeze(0))\n",
        "                    if not return_mask:\n",
        "                        x_a_recon.append(self.gen_a2b_s[j].decode(c_a, s_a_fake, x_a[i].unsqueeze(0)))\n",
        "                        x_ab1.append(self.gen_a2b_s[j].decode(c_a, s_b1[i].unsqueeze(0), x_a[i].unsqueeze(0)))\n",
        "                        x_ab2.append(self.gen_a2b_s[j].decode(c_a, s_b2[i].unsqueeze(0), x_a[i].unsqueeze(0)))\n",
        "                    else:\n",
        "                        x_ab1_tmp, x_ab1_mask_tmp = self.gen_a2b_s[j].decode(c_a, s_b1[i].unsqueeze(0), x_a[i].unsqueeze(0), return_mask=return_mask)\n",
        "                        do_double = False\n",
        "                        if do_double:\n",
        "                            c_a_double, s_a_fake = self.gen_a2b_s[j].encode(x_ab1_tmp)\n",
        "                            x_ab1_tmp, x_ab1_mask_tmp = self.gen_a2b_s[j].decode(c_a_double, s_b1[i].unsqueeze(0),\n",
        "                                                                               x_ab1_tmp,\n",
        "                                                                               return_mask=return_mask)\n",
        "\n",
        "                        x_ab1_mask.append(x_ab1_mask_tmp)\n",
        "                        x_ab1.append(x_ab1_tmp)\n",
        "                        x_ab2.append(self.gen_a2b_s[j].decode(c_a, s_b2[i].unsqueeze(0), x_a[i].unsqueeze(0)))\n",
        "\n",
        "        if self.do_b2a_conf:\n",
        "            x_b_s = torch.cat(x_b_s)\n",
        "            x_ba1, x_ba2 = torch.cat(x_ba1), torch.cat(x_ba2)\n",
        "            if not return_mask:\n",
        "                x_b_recon = torch.cat(x_b_recon)\n",
        "            else:\n",
        "                x_ba1_mask = torch.cat(x_ba1_mask)\n",
        "        if self.do_a2b_conf:\n",
        "            x_a_s = torch.cat(x_a_s)\n",
        "            x_ab1, x_ab2 = torch.cat(x_ab1), torch.cat(x_ab2)\n",
        "            if not return_mask:\n",
        "                x_a_recon = torch.cat(x_a_recon)\n",
        "            else:\n",
        "                x_ab1_mask = torch.cat(x_ab1_mask)\n",
        "\n",
        "        self.train()\n",
        "\n",
        "        do_diff = False\n",
        "        if do_diff:\n",
        "            if self.do_a2b_conf:\n",
        "                x_ab1 = x_a_s - x_ab1\n",
        "                x_ab2 = x_a_s - x_ab2\n",
        "            if self.do_b2a_conf:\n",
        "                x_ba1 = x_b_s - x_ba1\n",
        "                x_ba2 = x_b_s - x_ba2\n",
        "\n",
        "        if not return_mask:\n",
        "            if self.do_a2b_conf and self.do_b2a_conf:\n",
        "                return x_a_s, x_a_recon, x_ab1, x_ab2, x_b_s, x_b_recon, x_ba1, x_ba2\n",
        "            if self.do_a2b_conf:\n",
        "                return x_a_s, x_a_recon, x_ab1, x_ab2, None, None, None, None\n",
        "            if self.do_b2a_conf:\n",
        "                return None, None, None, None, x_b_s, x_b_recon, x_ba1, x_ba2\n",
        "        else:\n",
        "            if self.do_a2b_conf and self.do_b2a_conf:\n",
        "                return x_a_s, x_ab1_mask, x_ab1, x_ab2, x_b_s, x_ba1_mask, x_ba1, x_ba2\n",
        "            if self.do_a2b_conf:\n",
        "                return x_a_s, x_ab1_mask, x_ab1, x_ab2, None, None, None, None\n",
        "            if self.do_b2a_conf:\n",
        "                return None, None, None, None, x_b_s, x_ba1_mask, x_ba1, x_ba2\n",
        "\n",
        "    def dis_update(self, x_a=None, x_b=None, hyperparameters=None):\n",
        "        x_a_dis = x_a if not hyperparameters['dis']['do_Dis_only_gray'] else torch.sum(x_a.detach(), 1).unsqueeze(1).repeat(1, hyperparameters['input_dim_a'], 1, 1) / hyperparameters['input_dim_a']\n",
        "        x_b_dis = x_b if not hyperparameters['dis']['do_Dis_only_gray'] else torch.sum(x_b.detach(), 1).unsqueeze(1).repeat(1, hyperparameters['input_dim_b'], 1, 1) / hyperparameters['input_dim_b']\n",
        "        for dis_opt in self.dis_opt_s:\n",
        "            dis_opt.zero_grad()\n",
        "        if self.do_a2b_conf:\n",
        "            s_b = Variable(torch.randn(x_b.size(0), self.style_dim, 1, 1).cuda(self.cuda_device))\n",
        "            self.loss_dis_a2b_s = []\n",
        "        if self.do_b2a_conf:\n",
        "            s_a = Variable(torch.randn(x_a.size(0), self.style_dim, 1, 1).cuda(self.cuda_device))\n",
        "            self.loss_dis_b2a_s = []\n",
        "        self.loss_dis_total_s = []\n",
        "        for i in range(self.council_size):\n",
        "            i_gen = i\n",
        "            if hyperparameters['dis']['useRandomGen']:\n",
        "                i_gen = np.random.randint(self.council_size)\n",
        "\n",
        "            # encode\n",
        "            if hyperparameters['do_a2b']:\n",
        "                c_a, _ = self.gen_a2b_s[i_gen].encode(x_a)\n",
        "            if hyperparameters['do_b2a']:\n",
        "                c_b, _ = self.gen_b2a_s[i_gen].encode(x_b)\n",
        "\n",
        "            # decode (cross domain)\n",
        "            if hyperparameters['do_a2b']:\n",
        "                x_ab = self.gen_a2b_s[i_gen].decode(c_a, s_b, x_a)\n",
        "                x_ab = x_ab if not hyperparameters['dis']['do_Dis_only_gray'] else torch.sum(x_ab.detach(), 1).unsqueeze(1).repeat(1, hyperparameters['input_dim_b'], 1, 1) / hyperparameters['input_dim_b']\n",
        "\n",
        "            if hyperparameters['do_b2a']:\n",
        "                x_ba = self.gen_b2a_s[i_gen].decode(c_b, s_a, x_b)\n",
        "                x_ba = x_ba if not hyperparameters['dis']['do_Dis_only_gray'] else torch.sum(x_ba.detach(), 1).unsqueeze(1).repeat(1, hyperparameters['input_dim_a'], 1, 1) / hyperparameters['input_dim_a']\n",
        "\n",
        "            # D loss\n",
        "            if hyperparameters['do_a2b']:\n",
        "                self.loss_dis_a2b_s.append(self.dis_a2b_s[i].calc_dis_loss(x_ab.detach(), x_b_dis))\n",
        "            if hyperparameters['do_b2a']:\n",
        "                self.loss_dis_b2a_s.append(self.dis_b2a_s[i].calc_dis_loss(x_ba.detach(), x_a_dis))\n",
        "\n",
        "            self.loss_dis_total_s.append(0)\n",
        "            if hyperparameters['do_a2b']:\n",
        "                self.loss_dis_total_s[i] += hyperparameters['gan_w'] * self.loss_dis_a2b_s[i]\n",
        "            if hyperparameters['do_b2a']:\n",
        "                self.loss_dis_total_s[i] += self.loss_dis_b2a_s[i]\n",
        "\n",
        "            self.loss_dis_total_s[i].backward()\n",
        "            self.dis_opt_s[i].step()\n",
        "\n",
        "    def dis_council_update(self, x_a=None, x_b=None, hyperparameters=None):\n",
        "\n",
        "        if self.council_size <= 1 or hyperparameters['council']['numberOfCouncil_dis_relative_iteration'] == 0:\n",
        "            print('no council discriminetor is needed (council size <= 1 or numberOfCouncil_dis_relative_iteration == 0)')\n",
        "            return\n",
        "        onOffCycle = hyperparameters['council']['flipOnOff_On_iteration'] + hyperparameters['council'][\n",
        "            'flipOnOff_Off_iteration']\n",
        "        currIterCyc = hyperparameters['iteration'] % onOffCycle\n",
        "        if hyperparameters['council']['flipOnOff_start_with']:\n",
        "            startCyc = hyperparameters['council']['flipOnOff_On_iteration']\n",
        "        else:\n",
        "            startCyc = hyperparameters['council']['flipOnOff_Off_iteration']\n",
        "\n",
        "        self.do_council_loss = hyperparameters['council']['flipOnOff_start_with'] if (currIterCyc < startCyc) \\\n",
        "            else not hyperparameters['council']['flipOnOff_start_with']\n",
        "        if not hyperparameters['council']['flipOnOff']:\n",
        "            self.do_council_loss = hyperparameters['council']['flipOnOff_start_with']\n",
        "\n",
        "        if not self.do_council_loss or hyperparameters['council_w'] == 0 or hyperparameters['iteration'] < hyperparameters['council']['council_start_at_iter']:\n",
        "            return\n",
        "\n",
        "        for dis_council_opt in self.dis_council_opt_s:\n",
        "            dis_council_opt.zero_grad()\n",
        "\n",
        "        if self.do_b2a_conf:\n",
        "            s_a = Variable(torch.randn(x_a.size(0), self.style_dim, 1, 1).cuda(self.cuda_device))\n",
        "        if self.do_a2b_conf:\n",
        "            s_b = Variable(torch.randn(x_b.size(0), self.style_dim, 1, 1).cuda(self.cuda_device))\n",
        "        if hyperparameters['council']['discriminetro_less_style_by'] != 0:\n",
        "            if self.do_b2a_conf:\n",
        "                s_a_less = s_a * hyperparameters['council']['discriminetro_less_style_by']\n",
        "            if self.do_a2b_conf:\n",
        "                s_b_less = s_b * hyperparameters['council']['discriminetro_less_style_by']\n",
        "\n",
        "        self.loss_dis_council_a2b_s = []\n",
        "        self.loss_dis_council_b2a_s = []\n",
        "        self.loss_dis_council_total_s = []\n",
        "        c_a_s = []\n",
        "        c_b_s = []\n",
        "        x_ba_s = []\n",
        "        x_ab_s = []\n",
        "        x_ba_s_less = []\n",
        "        x_ab_s_less = []\n",
        "\n",
        "        for i in range(self.council_size):\n",
        "            # encode\n",
        "            if hyperparameters['do_a2b']:\n",
        "                c_a, _ = self.gen_a2b_s[i].encode(x_a)\n",
        "                c_a_s.append(c_a)\n",
        "            if hyperparameters['do_b2a']:\n",
        "                c_b, _ = self.gen_b2a_s[i].encode(x_b)\n",
        "                c_b_s.append(c_b)\n",
        "\n",
        "            # decode (cross domain)\n",
        "            if hyperparameters['do_a2b']:\n",
        "                x_ab = self.gen_a2b_s[i].decode(c_a, s_b, x_a)\n",
        "                x_ab_s.append(x_ab)\n",
        "            if hyperparameters['do_b2a']:\n",
        "                x_ba = self.gen_b2a_s[i].decode(c_b, s_a, x_b)\n",
        "                x_ba_s.append(x_ba)\n",
        "\n",
        "            if hyperparameters['council']['discriminetro_less_style_by'] != 0:\n",
        "                # decode (cross domain) less_style_by\n",
        "                if hyperparameters['do_a2b']:\n",
        "                    x_ab_less = self.gen_a2b_s[i].decode(c_a, s_b_less, x_a)\n",
        "                    x_ab_s_less.append(x_ab_less)\n",
        "\n",
        "                if hyperparameters['do_b2a']:\n",
        "                    x_ba_less = self.gen_b2a_s[i].decode(c_b, s_a_less, x_b)\n",
        "                    x_ba_s_less.append(x_ba_less)\n",
        "\n",
        "        if self.do_a2b_conf:\n",
        "            comper_x_ab_s = x_ab_s if hyperparameters['council']['discriminetro_less_style_by'] == 0 else x_ab_s_less\n",
        "        if self.do_b2a_conf:\n",
        "            comper_x_ba_s = x_ba_s if hyperparameters['council']['discriminetro_less_style_by'] == 0 else x_ba_s_less\n",
        "\n",
        "        for i in range(self.council_size):\n",
        "            self.loss_dis_council_a2b_s.append(0)\n",
        "            self.loss_dis_council_b2a_s.append(0)\n",
        "            index_to_chose_from = list(range(0, i)) + list(range(i + 1, self.council_size))\n",
        "            for k in range(hyperparameters['council']['numberOfCouncil_dis_relative_iteration']):\n",
        "                if k == self.council_size:\n",
        "                    break\n",
        "                if len(index_to_chose_from) == 0:\n",
        "                    index_to_chose_from = list(range(0, i)) + list(range(i + 1, self.council_size)) # reinitilize the indexes to chose from if numberOfCouncil_dis_relative_iteration is biger then thr number of council members\n",
        "                index_to_comper = random.choice(index_to_chose_from)\n",
        "                index_to_chose_from.remove(index_to_comper)\n",
        "\n",
        "                # D loss\n",
        "                if hyperparameters['do_a2b']:\n",
        "                    self.loss_dis_council_a2b_s[i] += self.dis_council_a2b_s[i].calc_dis_loss(x_ab_s[i].detach(), comper_x_ab_s[index_to_comper].detach(), x_a)  # original\n",
        "                if hyperparameters['do_b2a']:\n",
        "                    self.loss_dis_council_b2a_s[i] += self.dis_council_b2a_s[i].calc_dis_loss(x_ba_s[i].detach(), comper_x_ba_s[index_to_comper].detach(), x_b)  # original\n",
        "\n",
        "            self.loss_dis_council_total_s.append(0)\n",
        "            if hyperparameters['do_a2b']:\n",
        "                self.loss_dis_council_total_s[i] += hyperparameters['council_w'] * self.loss_dis_council_a2b_s[i] / hyperparameters['council']['numberOfCouncil_dis_relative_iteration']\n",
        "            if hyperparameters['do_b2a']:\n",
        "                self.loss_dis_council_total_s[i] += hyperparameters['council_w'] * self.loss_dis_council_b2a_s[i] / hyperparameters['council']['numberOfCouncil_dis_relative_iteration']\n",
        "\n",
        "            self.loss_dis_council_total_s[i].backward()\n",
        "            self.dis_council_opt_s[i].step()\n",
        "\n",
        "    def update_learning_rate(self):\n",
        "        for dis_scheduler in self.dis_scheduler_s:\n",
        "            if dis_scheduler is not None:\n",
        "                dis_scheduler.step()\n",
        "        for gen_scheduler in self.gen_scheduler_s:\n",
        "            if gen_scheduler is not None:\n",
        "                gen_scheduler.step()\n",
        "        if not self.do_dis_council:\n",
        "            return\n",
        "        for dis_council_scheduler in self.dis_council_scheduler_s:\n",
        "            if dis_council_scheduler is not None:\n",
        "                dis_council_scheduler.step()\n",
        "\n",
        "    def resume(self, checkpoint_dir, hyperparameters):\n",
        "        iterations = 0\n",
        "        # Load generators\n",
        "        for i in range(self.council_size):\n",
        "            last_model_name = get_model_list(checkpoint_dir, \"gen_\" + str(i))\n",
        "            if last_model_name is not None:\n",
        "                last_model_name = last_model_name.replace('a2b_gen_', 'gen_').replace('b2a_gen_', 'gen_')\n",
        "                print('loading: ' + last_model_name)\n",
        "                if self.do_a2b_conf:\n",
        "                    state_dict = torch.load(last_model_name.replace('gen_', 'a2b_gen_'), map_location=torch.device(self.cuda_device))\n",
        "                    self.gen_a2b_s[i].load_state_dict(state_dict['a2b'])\n",
        "                if self.do_b2a_conf:\n",
        "                    state_dict = torch.load(last_model_name.replace('gen_', 'b2a_gen_'), map_location=torch.device(self.cuda_device))\n",
        "                    self.gen_b2a_s[i].load_state_dict(state_dict['b2a'])\n",
        "                iterations = int(last_model_name[-11:-3])\n",
        "            else:\n",
        "                warnings.warn('Failed to find gen checkpoint, did not load model')\n",
        "\n",
        "            # Load discriminators\n",
        "            last_model_name = get_model_list(checkpoint_dir, \"dis_\" + str(i))\n",
        "            if last_model_name is not None:\n",
        "                last_model_name = last_model_name.replace('a2b_dis_', 'dis_').replace('b2a_dis_', 'dis_')\n",
        "                print('loading: ' + last_model_name)\n",
        "                if self.do_a2b_conf:\n",
        "                    state_dict = torch.load(last_model_name.replace('dis_', 'a2b_dis_'), map_location=torch.device(self.cuda_device))\n",
        "                    self.dis_a2b_s[i].load_state_dict(state_dict['a2b'])\n",
        "                if self.do_b2a_conf:\n",
        "                    state_dict = torch.load(last_model_name.replace('dis_', 'b2a_dis_'), map_location=torch.device(self.cuda_device))\n",
        "                    self.dis_b2a_s[i].load_state_dict(state_dict['b2a'])\n",
        "            else:\n",
        "                warnings.warn('Failed to find dis checkpoint, did not load model')\n",
        "            # Load council discriminators\n",
        "            if self.do_dis_council:\n",
        "                try:\n",
        "                    last_model_name = get_model_list(checkpoint_dir, \"dis_council_\" + str(i))\n",
        "                    print('loading: ' + last_model_name)\n",
        "                    if last_model_name is not None:\n",
        "                        last_model_name = last_model_name.replace('a2b_dis_council_', 'dis_council_').replace('b2a_dis_council_', 'dis_council_')\n",
        "\n",
        "                        if self.do_a2b_conf:\n",
        "                            state_dict = torch.load(last_model_name.replace('dis_council_', 'a2b_dis_council_'), map_location=torch.device(self.cuda_device))\n",
        "                            self.dis_council_a2b_s[i].load_state_dict(state_dict['a2b'])\n",
        "                        if self.do_b2a_conf:\n",
        "                            state_dict = torch.load(last_model_name.replace('dis_council_', 'b2a_dis_council_'), map_location=torch.device(self.cuda_device))\n",
        "                            self.dis_council_b2a_s[i].load_state_dict(state_dict['b2a'])\n",
        "                    else:\n",
        "                        warnings.warn('Failed to find dis checkpoint, did not load model')\n",
        "                except:\n",
        "                    warnings.warn('some council discriminetor FAILED to load')\n",
        "\n",
        "            # Load optimizers\n",
        "            try:\n",
        "                state_dict = torch.load(os.path.join(checkpoint_dir, 'optimizer_' + str(i) + '.pt'), map_location=torch.device(self.cuda_device))\n",
        "                self.dis_opt_s[i].load_state_dict(state_dict['dis'])\n",
        "                self.gen_opt_s[i].load_state_dict(state_dict['gen'])\n",
        "                if self.do_dis_council:\n",
        "                    self.dis_council_opt_s[i].load_state_dict(state_dict['dis_council'])\n",
        "\n",
        "                # Reinitilize schedulers\n",
        "                self.dis_scheduler_s[i] = get_scheduler(self.dis_opt_s[i], hyperparameters, iterations)\n",
        "                self.gen_scheduler = get_scheduler(self.gen_opt_s[i], hyperparameters, iterations)\n",
        "                if self.do_dis_council:\n",
        "                    self.dis_council_scheduler_s[i] = get_scheduler(self.dis_council_opt_s[i], hyperparameters, iterations)\n",
        "            except:\n",
        "                warnings.warn('some optimizer FAILED to load ')\n",
        "        if iterations > 0 :\n",
        "            print('Resume from iteration %d' % iterations)\n",
        "        else:\n",
        "            warnings.warn('FAILED TO RESUME STARTED FROM 0')\n",
        "        return iterations\n",
        "\n",
        "    def save(self, snapshot_dir, iterations):\n",
        "        for i in range(self.council_size):\n",
        "\n",
        "            # Save generators, discriminators, and optimizers\n",
        "            gen_name = os.path.join(snapshot_dir, 'gen_' + str(i) + '_%08d.pt' % (iterations + 1))\n",
        "            dis_name = os.path.join(snapshot_dir, 'dis_' + str(i) + '_%08d.pt' % (iterations + 1))\n",
        "            if self.do_dis_council:\n",
        "                dis_council_name = os.path.join(snapshot_dir, 'dis_council_' + str(i) + '_%08d.pt' % (iterations + 1))\n",
        "            opt_name = os.path.join(snapshot_dir, 'optimizer_' + str(i) + '.pt')\n",
        "            if self.do_a2b_conf:\n",
        "                torch.save({'a2b': self.gen_a2b_s[i].state_dict()}, gen_name.replace('gen_', 'a2b_gen_'))\n",
        "                torch.save({'a2b': self.dis_a2b_s[i].state_dict()}, dis_name.replace('dis_', 'a2b_dis_'))\n",
        "            if self.do_b2a_conf:\n",
        "                torch.save({'b2a': self.gen_b2a_s[i].state_dict()}, gen_name.replace('gen_', 'b2a_gen_'))\n",
        "                torch.save({'b2a': self.dis_b2a_s[i].state_dict()}, dis_name.replace('dis_', 'b2a_dis_'))\n",
        "            if self.do_dis_council:\n",
        "                if self.do_a2b_conf:\n",
        "                    torch.save({'a2b': self.dis_council_a2b_s[i].state_dict()}, dis_council_name.replace('dis_council_', 'a2b_dis_council_'))\n",
        "                if self.do_b2a_conf:\n",
        "                    torch.save({'b2a': self.dis_council_b2a_s[i].state_dict()},  dis_council_name.replace('dis_council_', 'b2a_dis_council_'))\n",
        "                torch.save({'gen': self.gen_opt_s[i].state_dict(), 'dis': self.dis_opt_s[i].state_dict(),\n",
        "                            'dis_council': self.dis_council_opt_s[i].state_dict()}, opt_name)\n",
        "            else:\n",
        "                torch.save({'gen': self.gen_opt_s[i].state_dict(), 'dis': self.dis_opt_s[i].state_dict()}, opt_name)\n",
        "\n",
        "#%%\n"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WUTk5ZhGGmWN",
        "colab_type": "text"
      },
      "source": [
        "#download.sh"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "V9B_bbc2Gd3E",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 0
        },
        "outputId": "b479b2a2-379a-4485-b524-d55b3042d3b8"
      },
      "source": [
        "\n",
        "#!/usr/bin/env bash\n",
        "%%shell\n",
        "\n",
        "FILE=\"pretrain_selfie_to_anime\"\n",
        "#pretrain_selfie_to_anime, pretrain_male_to_female, pretrain_glasses_removal\n",
        "\n",
        "#celeba dataset\n",
        "# @inproceedings{liu2015faceattributes,\n",
        "# title = {Deep Learning Face Attributes in the Wild},\n",
        "# author = {Liu, Ziwei and Luo, Ping and Wang, Xiaogang and Tang, Xiaoou},\n",
        "# booktitle = {Proceedings of International Conference on Computer Vision (ICCV)},\n",
        "# month = {December},\n",
        "# year = {2015}\n",
        "#}\n",
        "if [ $FILE == \"celeba_male2female\" ]; then\n",
        "    URL=https://cgm.technion.ac.il/Computer-Graphics-Multimedia/CouncilGAN/DataSet/celeba_male2female.zip?dl=0\n",
        "    ZIP_FILE=./datasets/celeba_male2female.zip\n",
        "    mkdir -p ./datasets/\n",
        "    wget -N $URL -O $ZIP_FILE\n",
        "    unzip $ZIP_FILE -d ./datasets/\n",
        "    rm $ZIP_FILE\n",
        "\n",
        "elif [ $FILE == \"celeba_glasses_removal\" ]; then\n",
        "    URL=https://cgm.technion.ac.il/Computer-Graphics-Multimedia/CouncilGAN/DataSet/celeba_glasses.zip?dl=0\n",
        "    ZIP_FILE=./datasets/celeba_glasses.zip\n",
        "    mkdir -p ./datasets/\n",
        "    wget -N $URL -O $ZIP_FILE\n",
        "    unzip $ZIP_FILE -d ./datasets/\n",
        "    rm $ZIP_FILE\n",
        "\n",
        "#  selfie2anime dataset used in https://github.com/taki0112/UGATIT\n",
        "#  title={U-GAT-IT: Unsupervised Generative Attentional Networks with Adaptive Layer-Instance Normalization for Image-to-Image Translation},\n",
        "#  author={Kim, Junho and Kim, Minjae and Kang, Hyeonwoo and Lee, Kwanghee},\n",
        "#  journal={arXiv preprint arXiv:1907.10830},\n",
        "#  year={2019}\n",
        "elif [ $FILE == \"U_GAT_IT_selfie2anime\" ]; then\n",
        "  URL=https://www.dropbox.com/s/9lz6gwwwyyxpdnn/selfie2anime.zip?dl=0\n",
        "  ZIP_FILE=./datasets/selfie2anime.zip\n",
        "  mkdir -p ./datasets/\n",
        "  wget -N $URL -O $ZIP_FILE\n",
        "  unzip $ZIP_FILE -d ./datasets/\n",
        "  rm $ZIP_FILE\n",
        "\n",
        "elif [ $FILE == \"pretrain_male_to_female\" ]; then\n",
        "    URL=\"https://onedrive.live.com/download?cid=552378614E6BA583&resid=552378614E6BA583%2119549&authkey=ANzlMZsBQy77urU\"\n",
        "    ZIP_FILE=./pretrain/pretrain_m2f.zip\n",
        "\n",
        "    mkdir -p ./pretrain/\n",
        "    wget -N $URL -O $ZIP_FILE\n",
        "    unzip $ZIP_FILE -d ./pretrain/\n",
        "    rm $ZIP_FILE\n",
        "\n",
        "elif [ $FILE == \"pretrain_glasses_removal\" ]; then\n",
        "    URL=\"https://onedrive.live.com/download?cid=552378614E6BA583&resid=552378614E6BA583%2119556&authkey=AAjSgID8ia6_Wt8\"\n",
        "    ZIP_FILE=./pretrain/pretrain_glasses_removal.zip\n",
        "    mkdir -p ./pretrain/\n",
        "    wget -N $URL -O $ZIP_FILE\n",
        "    unzip $ZIP_FILE -d ./pretrain/\n",
        "    rm $ZIP_FILE\n",
        "\n",
        "elif [ $FILE == \"pretrain_selfie_to_anime\" ]; then\n",
        "    URL=\"https://onedrive.live.com/download?cid=552378614E6BA583&resid=552378614E6BA583%2119566&authkey=AIcUuinR_eQPp1c\"\n",
        "    ZIP_FILE=./pretrain/pretrain_anime.zip\n",
        "    mkdir -p ./pretrain/\n",
        "    wget -N $URL -O $ZIP_FILE\n",
        "    unzip $ZIP_FILE -d ./pretrain/\n",
        "    rm $ZIP_FILE\n",
        "\n",
        "\n",
        "else\n",
        "    echo \"Available arguments are celeba_male2female, celeba_glasses_removal, selfie2anime, pretrain_male_to_female, pretrain_glasses_removal, pretrain_selfie_to_anime.\"\n",
        "    exit 1\n",
        "\n",
        "fi\n"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "WARNING: timestamping does nothing in combination with -O. See the manual\n",
            "for details.\n",
            "\n",
            "--2020-08-05 12:37:32--  https://onedrive.live.com/download?cid=552378614E6BA583&resid=552378614E6BA583%2119566&authkey=AIcUuinR_eQPp1c\n",
            "Resolving onedrive.live.com (onedrive.live.com)... 13.107.42.13\n",
            "Connecting to onedrive.live.com (onedrive.live.com)|13.107.42.13|:443... connected.\n",
            "HTTP request sent, awaiting response... 302 Found\n",
            "Location: https://dy0vaw.db.files.1drv.com/y4mWRnf4inWtXM5phjKCP7CSIvRsEMFxtNLajy9WI46vyvnqNK-NAQg8NS4WvCSnxu7f84ialAiiizf8aJk-TTNFNKZhTcA14FgefUqnZ9aqNL9ENgz1_T_i_uFYW_J14VgKhoyAJTzIeY1-XbfQ_BJgCsuWCTdk3jny8ZkB-V1HvTTIQvw0bvQqoQxzn5r8hQVeogTDRoNoakWDOtzyAiy7w/pretrain_anime.zip?download&psid=1 [following]\n",
            "--2020-08-05 12:37:33--  https://dy0vaw.db.files.1drv.com/y4mWRnf4inWtXM5phjKCP7CSIvRsEMFxtNLajy9WI46vyvnqNK-NAQg8NS4WvCSnxu7f84ialAiiizf8aJk-TTNFNKZhTcA14FgefUqnZ9aqNL9ENgz1_T_i_uFYW_J14VgKhoyAJTzIeY1-XbfQ_BJgCsuWCTdk3jny8ZkB-V1HvTTIQvw0bvQqoQxzn5r8hQVeogTDRoNoakWDOtzyAiy7w/pretrain_anime.zip?download&psid=1\n",
            "Resolving dy0vaw.db.files.1drv.com (dy0vaw.db.files.1drv.com)... 13.107.42.12\n",
            "Connecting to dy0vaw.db.files.1drv.com (dy0vaw.db.files.1drv.com)|13.107.42.12|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 359108374 (342M) [application/zip]\n",
            "Saving to: ‘./pretrain/pretrain_anime.zip’\n",
            "\n",
            "./pretrain/pretrain 100%[===================>] 342.47M  29.7MB/s    in 15s     \n",
            "\n",
            "2020-08-05 12:37:49 (23.4 MB/s) - ‘./pretrain/pretrain_anime.zip’ saved [359108374/359108374]\n",
            "\n",
            "Archive:  ./pretrain/pretrain_anime.zip\n",
            "   creating: ./pretrain/anime/\n",
            "   creating: ./pretrain/anime/256/\n",
            "  inflating: ./pretrain/anime/256/anime2face_council_folder.yaml  \n",
            "  inflating: ./pretrain/anime/256/b2a_gen_3_01000000.pt  \n",
            "  inflating: ./pretrain/anime/256/b2a_gen_2_01000000.pt  \n",
            "  inflating: ./pretrain/anime/256/b2a_gen_1_01000000.pt  \n",
            "  inflating: ./pretrain/anime/256/b2a_gen_0_01000000.pt  \n",
            "   creating: ./pretrain/anime/128/\n",
            "  inflating: ./pretrain/anime/128/anime2face_council_folder.yaml  \n",
            "  inflating: ./pretrain/anime/128/b2a_gen_3_01000000.pt  \n",
            "  inflating: ./pretrain/anime/128/b2a_gen_2_01000000.pt  \n",
            "  inflating: ./pretrain/anime/128/b2a_gen_1_01000000.pt  \n",
            "  inflating: ./pretrain/anime/128/b2a_gen_0_01000000.pt  \n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              ""
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6MB32ZWck43o",
        "colab_type": "text"
      },
      "source": [
        "# test_on_folder def\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bvQVvKaamyDt",
        "colab_type": "code",
        "cellView": "both",
        "colab": {}
      },
      "source": [
        "#@title input_configs { run: \"auto\" }\n",
        "config = './pretrain/anime/256/anime2face_council_folder.yaml'  #@param {type: \"string\"}\n",
        "checkpoint = './pretrain/anime/256/b2a_gen_0_01000000.pt'  #@param {type: \"string\"}\n",
        "input_folder = './input_folder'  #@param {type: \"string\"}\n",
        "output_folder = './output_folder'  #@param {type: \"string\"}\n",
        "output_path = 'outputs' #@param {type: \"string\"}\n",
        "a2b = 1 #@param {type: \"integer\"} #1 for a2b 0 for b2a\n",
        "seed = 1 #@param {type: \"integer\"}\n",
        "num_style = 10 #@param {type: \"integer\"}\n",
        "output_only = False #@param{type: \"boolean\"}\n",
        "num_of_images_to_test = 100 #@param {type: \"integer\"}\n",
        "config = get_config(config)"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-kRR3iT9jqfe",
        "colab_type": "text"
      },
      "source": [
        "# UPLOAD IMAGES"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "H7mBvPYDjWXH",
        "colab_type": "code",
        "colab": {
          "resources": {
            "http://localhost:8080/nbextensions/google.colab/files.js": {
              "data": "Ly8gQ29weXJpZ2h0IDIwMTcgR29vZ2xlIExMQwovLwovLyBMaWNlbnNlZCB1bmRlciB0aGUgQXBhY2hlIExpY2Vuc2UsIFZlcnNpb24gMi4wICh0aGUgIkxpY2Vuc2UiKTsKLy8geW91IG1heSBub3QgdXNlIHRoaXMgZmlsZSBleGNlcHQgaW4gY29tcGxpYW5jZSB3aXRoIHRoZSBMaWNlbnNlLgovLyBZb3UgbWF5IG9idGFpbiBhIGNvcHkgb2YgdGhlIExpY2Vuc2UgYXQKLy8KLy8gICAgICBodHRwOi8vd3d3LmFwYWNoZS5vcmcvbGljZW5zZXMvTElDRU5TRS0yLjAKLy8KLy8gVW5sZXNzIHJlcXVpcmVkIGJ5IGFwcGxpY2FibGUgbGF3IG9yIGFncmVlZCB0byBpbiB3cml0aW5nLCBzb2Z0d2FyZQovLyBkaXN0cmlidXRlZCB1bmRlciB0aGUgTGljZW5zZSBpcyBkaXN0cmlidXRlZCBvbiBhbiAiQVMgSVMiIEJBU0lTLAovLyBXSVRIT1VUIFdBUlJBTlRJRVMgT1IgQ09ORElUSU9OUyBPRiBBTlkgS0lORCwgZWl0aGVyIGV4cHJlc3Mgb3IgaW1wbGllZC4KLy8gU2VlIHRoZSBMaWNlbnNlIGZvciB0aGUgc3BlY2lmaWMgbGFuZ3VhZ2UgZ292ZXJuaW5nIHBlcm1pc3Npb25zIGFuZAovLyBsaW1pdGF0aW9ucyB1bmRlciB0aGUgTGljZW5zZS4KCi8qKgogKiBAZmlsZW92ZXJ2aWV3IEhlbHBlcnMgZm9yIGdvb2dsZS5jb2xhYiBQeXRob24gbW9kdWxlLgogKi8KKGZ1bmN0aW9uKHNjb3BlKSB7CmZ1bmN0aW9uIHNwYW4odGV4dCwgc3R5bGVBdHRyaWJ1dGVzID0ge30pIHsKICBjb25zdCBlbGVtZW50ID0gZG9jdW1lbnQuY3JlYXRlRWxlbWVudCgnc3BhbicpOwogIGVsZW1lbnQudGV4dENvbnRlbnQgPSB0ZXh0OwogIGZvciAoY29uc3Qga2V5IG9mIE9iamVjdC5rZXlzKHN0eWxlQXR0cmlidXRlcykpIHsKICAgIGVsZW1lbnQuc3R5bGVba2V5XSA9IHN0eWxlQXR0cmlidXRlc1trZXldOwogIH0KICByZXR1cm4gZWxlbWVudDsKfQoKLy8gTWF4IG51bWJlciBvZiBieXRlcyB3aGljaCB3aWxsIGJlIHVwbG9hZGVkIGF0IGEgdGltZS4KY29uc3QgTUFYX1BBWUxPQURfU0laRSA9IDEwMCAqIDEwMjQ7CgpmdW5jdGlvbiBfdXBsb2FkRmlsZXMoaW5wdXRJZCwgb3V0cHV0SWQpIHsKICBjb25zdCBzdGVwcyA9IHVwbG9hZEZpbGVzU3RlcChpbnB1dElkLCBvdXRwdXRJZCk7CiAgY29uc3Qgb3V0cHV0RWxlbWVudCA9IGRvY3VtZW50LmdldEVsZW1lbnRCeUlkKG91dHB1dElkKTsKICAvLyBDYWNoZSBzdGVwcyBvbiB0aGUgb3V0cHV0RWxlbWVudCB0byBtYWtlIGl0IGF2YWlsYWJsZSBmb3IgdGhlIG5leHQgY2FsbAogIC8vIHRvIHVwbG9hZEZpbGVzQ29udGludWUgZnJvbSBQeXRob24uCiAgb3V0cHV0RWxlbWVudC5zdGVwcyA9IHN0ZXBzOwoKICByZXR1cm4gX3VwbG9hZEZpbGVzQ29udGludWUob3V0cHV0SWQpOwp9CgovLyBUaGlzIGlzIHJvdWdobHkgYW4gYXN5bmMgZ2VuZXJhdG9yIChub3Qgc3VwcG9ydGVkIGluIHRoZSBicm93c2VyIHlldCksCi8vIHdoZXJlIHRoZXJlIGFyZSBtdWx0aXBsZSBhc3luY2hyb25vdXMgc3RlcHMgYW5kIHRoZSBQeXRob24gc2lkZSBpcyBnb2luZwovLyB0byBwb2xsIGZvciBjb21wbGV0aW9uIG9mIGVhY2ggc3RlcC4KLy8gVGhpcyB1c2VzIGEgUHJvbWlzZSB0byBibG9jayB0aGUgcHl0aG9uIHNpZGUgb24gY29tcGxldGlvbiBvZiBlYWNoIHN0ZXAsCi8vIHRoZW4gcGFzc2VzIHRoZSByZXN1bHQgb2YgdGhlIHByZXZpb3VzIHN0ZXAgYXMgdGhlIGlucHV0IHRvIHRoZSBuZXh0IHN0ZXAuCmZ1bmN0aW9uIF91cGxvYWRGaWxlc0NvbnRpbnVlKG91dHB1dElkKSB7CiAgY29uc3Qgb3V0cHV0RWxlbWVudCA9IGRvY3VtZW50LmdldEVsZW1lbnRCeUlkKG91dHB1dElkKTsKICBjb25zdCBzdGVwcyA9IG91dHB1dEVsZW1lbnQuc3RlcHM7CgogIGNvbnN0IG5leHQgPSBzdGVwcy5uZXh0KG91dHB1dEVsZW1lbnQubGFzdFByb21pc2VWYWx1ZSk7CiAgcmV0dXJuIFByb21pc2UucmVzb2x2ZShuZXh0LnZhbHVlLnByb21pc2UpLnRoZW4oKHZhbHVlKSA9PiB7CiAgICAvLyBDYWNoZSB0aGUgbGFzdCBwcm9taXNlIHZhbHVlIHRvIG1ha2UgaXQgYXZhaWxhYmxlIHRvIHRoZSBuZXh0CiAgICAvLyBzdGVwIG9mIHRoZSBnZW5lcmF0b3IuCiAgICBvdXRwdXRFbGVtZW50Lmxhc3RQcm9taXNlVmFsdWUgPSB2YWx1ZTsKICAgIHJldHVybiBuZXh0LnZhbHVlLnJlc3BvbnNlOwogIH0pOwp9CgovKioKICogR2VuZXJhdG9yIGZ1bmN0aW9uIHdoaWNoIGlzIGNhbGxlZCBiZXR3ZWVuIGVhY2ggYXN5bmMgc3RlcCBvZiB0aGUgdXBsb2FkCiAqIHByb2Nlc3MuCiAqIEBwYXJhbSB7c3RyaW5nfSBpbnB1dElkIEVsZW1lbnQgSUQgb2YgdGhlIGlucHV0IGZpbGUgcGlja2VyIGVsZW1lbnQuCiAqIEBwYXJhbSB7c3RyaW5nfSBvdXRwdXRJZCBFbGVtZW50IElEIG9mIHRoZSBvdXRwdXQgZGlzcGxheS4KICogQHJldHVybiB7IUl0ZXJhYmxlPCFPYmplY3Q+fSBJdGVyYWJsZSBvZiBuZXh0IHN0ZXBzLgogKi8KZnVuY3Rpb24qIHVwbG9hZEZpbGVzU3RlcChpbnB1dElkLCBvdXRwdXRJZCkgewogIGNvbnN0IGlucHV0RWxlbWVudCA9IGRvY3VtZW50LmdldEVsZW1lbnRCeUlkKGlucHV0SWQpOwogIGlucHV0RWxlbWVudC5kaXNhYmxlZCA9IGZhbHNlOwoKICBjb25zdCBvdXRwdXRFbGVtZW50ID0gZG9jdW1lbnQuZ2V0RWxlbWVudEJ5SWQob3V0cHV0SWQpOwogIG91dHB1dEVsZW1lbnQuaW5uZXJIVE1MID0gJyc7CgogIGNvbnN0IHBpY2tlZFByb21pc2UgPSBuZXcgUHJvbWlzZSgocmVzb2x2ZSkgPT4gewogICAgaW5wdXRFbGVtZW50LmFkZEV2ZW50TGlzdGVuZXIoJ2NoYW5nZScsIChlKSA9PiB7CiAgICAgIHJlc29sdmUoZS50YXJnZXQuZmlsZXMpOwogICAgfSk7CiAgfSk7CgogIGNvbnN0IGNhbmNlbCA9IGRvY3VtZW50LmNyZWF0ZUVsZW1lbnQoJ2J1dHRvbicpOwogIGlucHV0RWxlbWVudC5wYXJlbnRFbGVtZW50LmFwcGVuZENoaWxkKGNhbmNlbCk7CiAgY2FuY2VsLnRleHRDb250ZW50ID0gJ0NhbmNlbCB1cGxvYWQnOwogIGNvbnN0IGNhbmNlbFByb21pc2UgPSBuZXcgUHJvbWlzZSgocmVzb2x2ZSkgPT4gewogICAgY2FuY2VsLm9uY2xpY2sgPSAoKSA9PiB7CiAgICAgIHJlc29sdmUobnVsbCk7CiAgICB9OwogIH0pOwoKICAvLyBXYWl0IGZvciB0aGUgdXNlciB0byBwaWNrIHRoZSBmaWxlcy4KICBjb25zdCBmaWxlcyA9IHlpZWxkIHsKICAgIHByb21pc2U6IFByb21pc2UucmFjZShbcGlja2VkUHJvbWlzZSwgY2FuY2VsUHJvbWlzZV0pLAogICAgcmVzcG9uc2U6IHsKICAgICAgYWN0aW9uOiAnc3RhcnRpbmcnLAogICAgfQogIH07CgogIGNhbmNlbC5yZW1vdmUoKTsKCiAgLy8gRGlzYWJsZSB0aGUgaW5wdXQgZWxlbWVudCBzaW5jZSBmdXJ0aGVyIHBpY2tzIGFyZSBub3QgYWxsb3dlZC4KICBpbnB1dEVsZW1lbnQuZGlzYWJsZWQgPSB0cnVlOwoKICBpZiAoIWZpbGVzKSB7CiAgICByZXR1cm4gewogICAgICByZXNwb25zZTogewogICAgICAgIGFjdGlvbjogJ2NvbXBsZXRlJywKICAgICAgfQogICAgfTsKICB9CgogIGZvciAoY29uc3QgZmlsZSBvZiBmaWxlcykgewogICAgY29uc3QgbGkgPSBkb2N1bWVudC5jcmVhdGVFbGVtZW50KCdsaScpOwogICAgbGkuYXBwZW5kKHNwYW4oZmlsZS5uYW1lLCB7Zm9udFdlaWdodDogJ2JvbGQnfSkpOwogICAgbGkuYXBwZW5kKHNwYW4oCiAgICAgICAgYCgke2ZpbGUudHlwZSB8fCAnbi9hJ30pIC0gJHtmaWxlLnNpemV9IGJ5dGVzLCBgICsKICAgICAgICBgbGFzdCBtb2RpZmllZDogJHsKICAgICAgICAgICAgZmlsZS5sYXN0TW9kaWZpZWREYXRlID8gZmlsZS5sYXN0TW9kaWZpZWREYXRlLnRvTG9jYWxlRGF0ZVN0cmluZygpIDoKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgJ24vYSd9IC0gYCkpOwogICAgY29uc3QgcGVyY2VudCA9IHNwYW4oJzAlIGRvbmUnKTsKICAgIGxpLmFwcGVuZENoaWxkKHBlcmNlbnQpOwoKICAgIG91dHB1dEVsZW1lbnQuYXBwZW5kQ2hpbGQobGkpOwoKICAgIGNvbnN0IGZpbGVEYXRhUHJvbWlzZSA9IG5ldyBQcm9taXNlKChyZXNvbHZlKSA9PiB7CiAgICAgIGNvbnN0IHJlYWRlciA9IG5ldyBGaWxlUmVhZGVyKCk7CiAgICAgIHJlYWRlci5vbmxvYWQgPSAoZSkgPT4gewogICAgICAgIHJlc29sdmUoZS50YXJnZXQucmVzdWx0KTsKICAgICAgfTsKICAgICAgcmVhZGVyLnJlYWRBc0FycmF5QnVmZmVyKGZpbGUpOwogICAgfSk7CiAgICAvLyBXYWl0IGZvciB0aGUgZGF0YSB0byBiZSByZWFkeS4KICAgIGxldCBmaWxlRGF0YSA9IHlpZWxkIHsKICAgICAgcHJvbWlzZTogZmlsZURhdGFQcm9taXNlLAogICAgICByZXNwb25zZTogewogICAgICAgIGFjdGlvbjogJ2NvbnRpbnVlJywKICAgICAgfQogICAgfTsKCiAgICAvLyBVc2UgYSBjaHVua2VkIHNlbmRpbmcgdG8gYXZvaWQgbWVzc2FnZSBzaXplIGxpbWl0cy4gU2VlIGIvNjIxMTU2NjAuCiAgICBsZXQgcG9zaXRpb24gPSAwOwogICAgd2hpbGUgKHBvc2l0aW9uIDwgZmlsZURhdGEuYnl0ZUxlbmd0aCkgewogICAgICBjb25zdCBsZW5ndGggPSBNYXRoLm1pbihmaWxlRGF0YS5ieXRlTGVuZ3RoIC0gcG9zaXRpb24sIE1BWF9QQVlMT0FEX1NJWkUpOwogICAgICBjb25zdCBjaHVuayA9IG5ldyBVaW50OEFycmF5KGZpbGVEYXRhLCBwb3NpdGlvbiwgbGVuZ3RoKTsKICAgICAgcG9zaXRpb24gKz0gbGVuZ3RoOwoKICAgICAgY29uc3QgYmFzZTY0ID0gYnRvYShTdHJpbmcuZnJvbUNoYXJDb2RlLmFwcGx5KG51bGwsIGNodW5rKSk7CiAgICAgIHlpZWxkIHsKICAgICAgICByZXNwb25zZTogewogICAgICAgICAgYWN0aW9uOiAnYXBwZW5kJywKICAgICAgICAgIGZpbGU6IGZpbGUubmFtZSwKICAgICAgICAgIGRhdGE6IGJhc2U2NCwKICAgICAgICB9LAogICAgICB9OwogICAgICBwZXJjZW50LnRleHRDb250ZW50ID0KICAgICAgICAgIGAke01hdGgucm91bmQoKHBvc2l0aW9uIC8gZmlsZURhdGEuYnl0ZUxlbmd0aCkgKiAxMDApfSUgZG9uZWA7CiAgICB9CiAgfQoKICAvLyBBbGwgZG9uZS4KICB5aWVsZCB7CiAgICByZXNwb25zZTogewogICAgICBhY3Rpb246ICdjb21wbGV0ZScsCiAgICB9CiAgfTsKfQoKc2NvcGUuZ29vZ2xlID0gc2NvcGUuZ29vZ2xlIHx8IHt9OwpzY29wZS5nb29nbGUuY29sYWIgPSBzY29wZS5nb29nbGUuY29sYWIgfHwge307CnNjb3BlLmdvb2dsZS5jb2xhYi5fZmlsZXMgPSB7CiAgX3VwbG9hZEZpbGVzLAogIF91cGxvYWRGaWxlc0NvbnRpbnVlLAp9Owp9KShzZWxmKTsK",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "application/javascript"
                ]
              ],
              "status": 200,
              "status_text": ""
            }
          },
          "base_uri": "https://localhost:8080/",
          "height": 514
        },
        "outputId": "f9a9b755-faeb-4a1f-9a38-46753ba62d84"
      },
      "source": [
        "print(f\"Upload images into {input_folder}\")\n",
        "from google.colab import files\n",
        "import shutil\n",
        "uploaded = files.upload()\n",
        "\n",
        "for fn in uploaded.keys():\n",
        "  print('User uploaded file \"{name}\" with length {length} bytes'.format(\n",
        "      name=fn, length=len(uploaded[fn])))\n",
        "\n",
        "try:\n",
        "  shutil.rmtree(input_folder)\n",
        "except:\n",
        "  pass\n",
        "os.makedirs(input_folder, exist_ok=True)\n",
        "for fn in uploaded.keys():\n",
        "  shutil.move(os.path.join(os.curdir,fn), input_folder)"
      ],
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Upload images into ./input_folder\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "     <input type=\"file\" id=\"files-dddadc2c-c8d1-4d30-abd0-c93303a14e7b\" name=\"files[]\" multiple disabled\n",
              "        style=\"border:none\" />\n",
              "     <output id=\"result-dddadc2c-c8d1-4d30-abd0-c93303a14e7b\">\n",
              "      Upload widget is only available when the cell has been executed in the\n",
              "      current browser session. Please rerun this cell to enable.\n",
              "      </output>\n",
              "      <script src=\"/nbextensions/google.colab/files.js\"></script> "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "Saving female_281.jpg to female_281 (1).jpg\n",
            "Saving female_348.jpg to female_348 (1).jpg\n",
            "Saving female_1409.jpg to female_1409 (1).jpg\n",
            "Saving female_1469.jpg to female_1469 (1).jpg\n",
            "Saving female_1479.jpg to female_1479 (1).jpg\n",
            "Saving female_2316.jpg to female_2316 (1).jpg\n",
            "Saving female_2321.jpg to female_2321 (1).jpg\n",
            "Saving female_2386.jpg to female_2386.jpg\n",
            "Saving female_4056.jpg to female_4056.jpg\n",
            "User uploaded file \"female_281.jpg\" with length 16161 bytes\n",
            "User uploaded file \"female_348.jpg\" with length 18444 bytes\n",
            "User uploaded file \"female_1409.jpg\" with length 26316 bytes\n",
            "User uploaded file \"female_1469.jpg\" with length 24914 bytes\n",
            "User uploaded file \"female_1479.jpg\" with length 18972 bytes\n",
            "User uploaded file \"female_2316.jpg\" with length 18206 bytes\n",
            "User uploaded file \"female_2321.jpg\" with length 19882 bytes\n",
            "User uploaded file \"female_2386.jpg\" with length 24244 bytes\n",
            "User uploaded file \"female_4056.jpg\" with length 17061 bytes\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vYfw_7MdSWTj",
        "colab_type": "text"
      },
      "source": [
        "# test_on_folder.py"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mT-tF_Uvim2q",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 255
        },
        "outputId": "a2127a41-a728-4df6-b6c6-29c37f3a1f28"
      },
      "source": [
        "from __future__ import print_function\n",
        "# from utils import get_config, get_data_loader_folder, pytorch03_to_pytorch04\n",
        "# from trainer_council import Council_Trainer\n",
        "from torch import nn\n",
        "from scipy.stats import entropy\n",
        "import torch.nn.functional as F\n",
        "import argparse\n",
        "from torch.autograd import Variable\n",
        "# from data import ImageFolder\n",
        "import numpy as np\n",
        "import torchvision.utils as vutils\n",
        "try:\n",
        "    from itertools import izip as zip\n",
        "except ImportError: # will be 3.x series\n",
        "    pass\n",
        "import torch\n",
        "import os\n",
        "from tqdm import tqdm\n",
        "\n",
        "data_name = 'out'\n",
        "\n",
        "torch.manual_seed(seed)\n",
        "torch.cuda.manual_seed(seed)\n",
        "\n",
        "# Load experiment setting\n",
        "input_dim = config['input_dim_a'] if a2b else config['input_dim_b']\n",
        "council_size = config['council']['council_size']\n",
        "\n",
        "\n",
        "# Setup model and data loader\n",
        "try:\n",
        "    image_names = ImageFolder(input_folder, transform=None, return_paths=True)\n",
        "except:\n",
        "  print('Put images to be translated in ./input_folder')\n",
        "  input(\"Press Enter when done\")\n",
        "  image_names = ImageFolder(input_folder, transform=None, return_paths=True)\n",
        "\n",
        "if not 'new_size_a' in config.keys():\n",
        "    config['new_size_a'] = config['new_size']\n",
        "is_data_A = a2b\n",
        "data_loader = get_data_loader_folder(input_folder, 1, False,\\\n",
        "                                     new_size=config['new_size_a'] if 'new_size_a' in config.keys() else config['new_size'],\\\n",
        "                                     crop=False, config=config, is_data_A=is_data_A)\n",
        "\n",
        "\n",
        "style_dim = config['gen']['style_dim']\n",
        "trainer = Council_Trainer(config)\n",
        "only_one = False\n",
        "if 'gen_' in checkpoint[-21:]:\n",
        "    state_dict = torch.load(checkpoint, map_location='cuda:0')\n",
        "    try:\n",
        "        if a2b:\n",
        "            trainer.gen_a2b_s[0].load_state_dict(state_dict['a2b'])\n",
        "        else:\n",
        "            trainer.gen_b2a_s[0].load_state_dict(state_dict['b2a'])\n",
        "    except:\n",
        "       print('opts.a2b should be set to ' + str(not a2b) + ' , Or config file could be wrong')\n",
        "       a2b = not a2b\n",
        "       if a2b:\n",
        "           trainer.gen_a2b_s[0].load_state_dict(state_dict['a2b'])\n",
        "       else:\n",
        "           trainer.gen_b2a_s[0].load_state_dict(state_dict['b2a'])\n",
        "            \n",
        "    council_size = 1\n",
        "    only_one = True\n",
        "else:\n",
        "    for i in range(council_size):\n",
        "        try:\n",
        "            if a2b:\n",
        "                tmp_checkpoint = checkpoint[:-8] + 'a2b_gen_' + str(i) + '_' + opts.checkpoint[-8:] + '.pt'\n",
        "                state_dict = torch.load(tmp_checkpoint)\n",
        "                trainer.gen_a2b_s[i].load_state_dict(state_dict['a2b'])\n",
        "            else:\n",
        "                tmp_checkpoint = checkpoint[:-8] + 'b2a_gen_' + str(i) + '_' + checkpoint[-8:] + '.pt'\n",
        "                state_dict = torch.load(tmp_checkpoint)\n",
        "                trainer.gen_b2a_s[i].load_state_dict(state_dict['b2a'])\n",
        "        except:\n",
        "            print('opts.a2b should be set to ' + str(not a2b) + ' , Or config file could be wrong')\n",
        "            \n",
        "            opts.a2b = not a2b\n",
        "            if a2b:\n",
        "                tmp_checkpoint = checkpoint[:-8] + 'a2b_gen_' + str(i) + '_' + checkpoint[-8:] + '.pt'\n",
        "                state_dict = torch.load(tmp_checkpoint)\n",
        "                trainer.gen_a2b_s[i].load_state_dict(state_dict['a2b'])\n",
        "            else:\n",
        "                tmp_checkpoint = checkpoint[:-8] + 'b2a_gen_' + str(i) + '_' + checkpoint[-8:] + '.pt'\n",
        "                state_dict = torch.load(tmp_checkpoint)\n",
        "                trainer.gen_b2a_s[i].load_state_dict(state_dict['b2a'])\n",
        "            \n",
        "\n",
        "\n",
        "trainer.cuda()\n",
        "trainer.eval()\n",
        "\n",
        "encode_s = []\n",
        "decode_s = []\n",
        "if a2b:\n",
        "    for i in range(council_size):\n",
        "        encode_s.append(trainer.gen_a2b_s[i].encode)  # encode function\n",
        "        decode_s.append(trainer.gen_a2b_s[i].decode)  # decode function\n",
        "else:\n",
        "    for i in range(council_size):\n",
        "        encode_s.append(trainer.gen_b2a_s[i].encode)  # encode function\n",
        "        decode_s.append(trainer.gen_b2a_s[i].decode)  # decode function\n",
        "\n",
        "\n",
        "# creat testing images\n",
        "num_of_images_to_test = num_of_images_to_test\n",
        "seed = 1\n",
        "curr_image_num = -1\n",
        "for i, (images, names) in tqdm(enumerate(zip(data_loader, image_names)), total=num_of_images_to_test):\n",
        "    if curr_image_num == num_of_images_to_test:\n",
        "        break\n",
        "    curr_image_num += 1\n",
        "    k = np.random.randint(council_size)\n",
        "    style_fixed = Variable(torch.randn(num_style, style_dim, 1, 1).cuda(), volatile=True)\n",
        "    print(names[1])\n",
        "    images = Variable(images.cuda(), volatile=True)\n",
        "\n",
        "    content, _ = encode_s[k](images)\n",
        "    seed += 1\n",
        "    torch.random.manual_seed(seed)\n",
        "    style = Variable(torch.randn(num_style, style_dim, 1, 1).cuda(), volatile=True)\n",
        "    for j in range(num_style):\n",
        "        s = style[j].unsqueeze(0)\n",
        "        outputs = decode_s[k](content, s, images)\n",
        "        basename = os.path.basename(names[1])\n",
        "        output_folder = os.path.join(output_path, 'test_res')\n",
        "        if only_one:\n",
        "            path = os.path.join(output_folder, checkpoint[-11:-3] + \"_%02d\" % j, data_name + '_out_' + str(curr_image_num) + '_' + str(j) + '.jpg')\n",
        "            path_all_in_one = os.path.join(output_folder, checkpoint[-11:-3] + '_all_in_1', data_name + '_out_' + str(curr_image_num) + '_' + str(j) + '.jpg')\n",
        "\n",
        "        else:\n",
        "            path = os.path.join(output_folder, checkpoint[-8:] + \"_%02d\" % j, data_name + '_out_' + str(curr_image_num) + '_' + str(j) + '.jpg')\n",
        "            path_all_in_one = os.path.join(output_folder,  checkpoint[-8:] + '_all_in_1', data_name + '_out_' + str(curr_image_num) + '_' + str(j) + '.jpg')\n",
        "\n",
        "        if not os.path.exists(os.path.dirname(path)):\n",
        "            os.makedirs(os.path.dirname(path))\n",
        "        vutils.save_image(outputs.data, path, padding=0, normalize=True)\n",
        "        do_all_in_one = True\n",
        "        if do_all_in_one:\n",
        "            if not os.path.exists(os.path.dirname(path_all_in_one)):\n",
        "                os.makedirs(os.path.dirname(path_all_in_one))\n",
        "            vutils.save_image(outputs.data, path_all_in_one, padding=0, normalize=True)\n",
        "    if not output_only:\n",
        "        # also save input images\n",
        "        output_folder = os.path.join(output_folder, 'input')\n",
        "        if not os.path.exists(output_folder):\n",
        "            os.makedirs(output_folder)\n",
        "        vutils.save_image(images.data, os.path.join(output_folder, 'input{:03d}.jpg'.format(i)), padding=0, normalize=True)\n"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "opts.a2b should be set to False , Or config file could be wrong\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "\r  0%|          | 0/100 [00:00<?, ?it/s]/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:114: UserWarning: volatile was removed and now has no effect. Use `with torch.no_grad():` instead.\n",
            "/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:116: UserWarning: volatile was removed and now has no effect. Use `with torch.no_grad():` instead.\n",
            "/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:121: UserWarning: volatile was removed and now has no effect. Use `with torch.no_grad():` instead.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "./input_folder/female_1409.jpg\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "\r  1%|          | 1/100 [00:01<01:44,  1.05s/it]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "./input_folder/female_1469.jpg\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "\r  2%|▏         | 2/100 [00:01<01:36,  1.02it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "./input_folder/female_1479.jpg\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "\r  3%|▎         | 3/100 [00:02<01:29,  1.08it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "./input_folder/female_2316.jpg\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "\r  4%|▍         | 4/100 [00:03<01:24,  1.14it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "./input_folder/female_2321.jpg\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "\r  5%|▌         | 5/100 [00:04<01:20,  1.18it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "./input_folder/female_2386.jpg\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "\r  6%|▌         | 6/100 [00:04<01:17,  1.21it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "./input_folder/female_281.jpg\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "\r  7%|▋         | 7/100 [00:05<01:15,  1.24it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "./input_folder/female_348.jpg\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "\r  8%|▊         | 8/100 [00:06<01:13,  1.25it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "./input_folder/female_4056.jpg\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "  9%|▉         | 9/100 [00:07<01:14,  1.22it/s]\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h-qKOfxEmbCq",
        "colab_type": "text"
      },
      "source": [
        "# Display\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "w9MwOH3Pmh7U",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "2447c90f-90fd-45df-9e3a-a66f2852d465"
      },
      "source": [
        "import matplotlib.pyplot as plt\n",
        "from PIL import Image\n",
        "\n",
        "fig=plt.figure(figsize=(128, 128))\n",
        "columns = num_style + 1\n",
        "rows = len(os.listdir(input_folder))\n",
        "\n",
        "for i in range(1, columns*rows +1):\n",
        "    curr_colum = i % columns\n",
        "    curr_row = int(i / columns)\n",
        "    if curr_colum == 0: # input image\n",
        "      img = Image.open('outputs/test_res/input/input{:03d}.jpg'.format(curr_row-1))\n",
        "    else:\n",
        "      img = Image.open('outputs/test_res/01000000_all_in_1/out_out_{:d}_{:d}.jpg'.format( curr_row, curr_colum-1))\n",
        "\n",
        "\n",
        "    fig.add_subplot(rows, columns, i)\n",
        "    plt.imshow(img)\n",
        "print(f'results are stored in output_folder')\n",
        "plt.show()\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "results are stored in output_folder\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}