DiffAugment_GAN.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/cedro3/data-efficient-gans/blob/master/DiffAugment_GAN.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KZb8pDhKRxcS"
      },
      "source": [
        "# Data-Efficient GANs with DiffAugment\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A4MvOMh1BNef"
      },
      "source": [
        "## セットアップ\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EqqY49pVYDqg"
      },
      "outputs": [],
      "source": [
        "# tensorflow 1.15.0　のインストール\n",
        "!pip uninstall -y tensorflow tensorflow-probability\n",
        "!pip install tensorflow-gpu==1.15.0 tensorflow-datasets==2.1.0\n",
        "\n",
        "# Githubからコードをコピー\n",
        "!git clone https://github.com/cedro3/data-efficient-gans.git\n",
        "%cd data-efficient-gans/DiffAugment-stylegan2\n",
        "\n",
        "# gdownアップデート\n",
        "! pip install --upgrade gdown\n",
        "\n",
        "# numpy 1.19.4のインストール\n",
        "! pip install numpy==1.19.4 "
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Google drive から新垣結衣の画像と学習済みの重みをダウンロード\n",
        "import gdown\n",
        "gdown.download('https://drive.google.com/uc?id=1rlmCCuLH4euzIwzqlYh8LiF5iizh5UWV', 'yui.zip', quiet=False)\n",
        "! unzip yui.zip"
      ],
      "metadata": {
        "id": "VNQqfSnLtaFJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b599aK-wCRBB"
      },
      "source": [
        "# 関数の定義"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Re5R6VX8VNgo"
      },
      "outputs": [],
      "source": [
        "# 関数の定義\n",
        "import tensorflow as tf\n",
        "import os\n",
        "import numpy as np\n",
        "import PIL\n",
        "import IPython\n",
        "from multiprocessing import Pool\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "from dnnlib import tflib, EasyDict\n",
        "from training import misc, dataset_tool\n",
        "from metrics import metric_base\n",
        "from metrics.metric_defaults import metric_defaults\n",
        "\n",
        "def _generate(network_name, num_rows, num_cols, seed, resolution):\n",
        "  if seed is not None:\n",
        "    np.random.seed(seed)\n",
        "  with tf.Session():\n",
        "    _, _, Gs = misc.load_pkl(network_name)\n",
        "    z = np.random.randn(num_rows * num_cols, Gs.input_shape[1])\n",
        "    outputs = Gs.run(z, None, output_transform=dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True))\n",
        "    outputs = np.reshape(outputs, [num_rows, num_cols, *outputs.shape[1:]])\n",
        "    outputs = np.concatenate(outputs, axis=1)\n",
        "    outputs = np.concatenate(outputs, axis=1)\n",
        "    img = PIL.Image.fromarray(outputs)\n",
        "    img = img.resize((resolution * num_cols, resolution * num_rows), PIL.Image.ANTIALIAS)\n",
        "  return img\n",
        "\n",
        "def generate(network_name, num_rows, num_cols, seed=None, resolution=128):\n",
        "  with Pool(1) as pool:\n",
        "    return pool.apply(_generate, (network_name, num_rows, num_cols, seed, resolution))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ObOQUWb9jEu-"
      },
      "source": [
        "# データセットの作成\n",
        "64×64の画像100枚からtfrecords形式のデータセットを作成します。\\"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "emaQzEznPWYt"
      },
      "outputs": [],
      "source": [
        "# 100-shot-gakki を読み込み学習用データセットを作成\n",
        "data_dir = dataset_tool.create_dataset('100-shot-gakki')\n",
        "training_images = []\n",
        "for fname in os.listdir(data_dir):\n",
        "  if fname.endswith('.jpg'):\n",
        "    training_images.append(np.array(PIL.Image.open(os.path.join(data_dir, fname))))\n",
        "imgs = np.reshape(training_images, [5, 20, *training_images[0].shape])\n",
        "imgs = np.concatenate(imgs, axis=1)\n",
        "imgs = np.concatenate(imgs, axis=1)\n",
        "PIL.Image.fromarray(imgs).resize((1000, 250), PIL.Image.ANTIALIAS)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O4YdmLWDyRzP"
      },
      "source": [
        "# 学習の実行\n",
        "学習時間は、割り当てられているGPUによって異なります。下記を参考にして下さい。\\ \n",
        "P100 : 7.3H (kimg=300)\\\n",
        "V100 : 4.1H (kimg=300)\\\n",
        "\\\n",
        "※割り当てがK80であったり、学習に時間を掛けたくない方は、kimg=500で学習した重みがありますので、ここはパスでOKです。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JC9zsb4XvIU8"
      },
      "outputs": [],
      "source": [
        "# GPUの確認\n",
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ywR-TqhzY01I"
      },
      "outputs": [],
      "source": [
        "# 学習の実行\n",
        "!python3 run_few_shot.py --dataset=100-shot-gakki --resolution=64 --total-kimg=300"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PXXzNi5N1ZJ1"
      },
      "source": [
        "# 学習済みの重みを使う\n",
        "実際に学習を行った場合は、resultsフォルダーの1段下に重み( network-snapshot-XXXXXX.pkl)が作成されますので、それをDiffAugment-stylegan2のディレクトリーに移動して下さい。\\\n",
        "そして、**generate() , generate_gif.py** の引数をそのファイル名に変更して下さい。\\\n",
        "\\"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "riefZb82lhTH"
      },
      "outputs": [],
      "source": [
        "# 学習済みの重みを使って画像生成\n",
        "generate('network-snapshot-gakki-000500.pkl', num_rows=2, num_cols=5, seed=3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "S9Q7bwdKvws9"
      },
      "outputs": [],
      "source": [
        "# 学習済みの重みを使って、GIF動画 (interp.gif) を作成\n",
        "!python3 generate_gif.py -r network-snapshot-gakki-000500.pkl -o interp.gif --num-rows=2 --num-cols=3 --seed=1\n",
        "IPython.display.Image(open('interp.gif', 'rb').read())"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "DiffAugment_GAN",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}