image_to_paint.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "image-to-paint",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/cedro3/stylized-neural-painting/blob/main/image_to_paint.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LFLeFyHW7AfS"
      },
      "source": [
        "# Githubからコードをコピー\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "h72s1Mk26j4V"
      },
      "source": [
        "# githubのコードをコピー\n",
        "!git clone https://github.com/cedro3/stylized-neural-painting.git\n",
        "%cd stylized-neural-painting"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 学習済みモデルのダウンロード"
      ],
      "metadata": {
        "id": "FDboP9pR76ma"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# 学習済みパラメータのダウンロード\n",
        "! pip install --upgrade gdown\n",
        "import gdown\n",
        "gdown.download('https://drive.google.com/uc?id=1sqWhgBKqaBJggl2A8sD1bLSq2_B1ScMG', './checkpoints_G_oilpaintbrush.zip', quiet=False)\n",
        "! unzip checkpoints_G_oilpaintbrush.zip"
      ],
      "metadata": {
        "id": "3EWAl4aq73JT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RwEdnkcI5fFE"
      },
      "source": [
        "# コード本体"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eD8de5Il7JdO"
      },
      "source": [
        "import argparse\n",
        "import torch\n",
        "torch.cuda.current_device()\n",
        "import torch.optim as optim\n",
        "from painter import *\n",
        "\n",
        "# Decide which device we want to run on\n",
        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "# settings\n",
        "parser = argparse.ArgumentParser(description='STYLIZED NEURAL PAINTING')\n",
        "args = parser.parse_args(args=[])\n",
        "args.img_path = './test_images/kasumi.png' # path to input photo\n",
        "args.renderer = 'oilpaintbrush' # [watercolor, markerpen, oilpaintbrush, rectangle]\n",
        "args.canvas_color = 'black' # [black, white]\n",
        "args.canvas_size = 512 # size of the canvas for stroke rendering'\n",
        "args.max_m_strokes = 500 # max number of strokes\n",
        "args.max_divide = 5 # divide an image up-to max_divide x max_divide patches\n",
        "args.beta_L1 = 1.0 # weight for L1 loss\n",
        "args.with_ot_loss = False # set True for imporving the convergence by using optimal transportation loss, but will slow-down the speed\n",
        "args.beta_ot = 0.1 # weight for optimal transportation loss\n",
        "args.net_G = 'zou-fusion-net' # renderer architecture\n",
        "args.renderer_checkpoint_dir = './checkpoints_G_oilpaintbrush' # dir to load the pretrained neu-renderer\n",
        "args.lr = 0.005 # learning rate for stroke searching\n",
        "args.output_dir = './output' # dir to save painting results\n",
        "\n",
        "\n",
        "def _drawing_step_states(pt):\n",
        "    acc = pt._compute_acc().item()\n",
        "    print('iteration step %d, G_loss: %.5f, step_acc: %.5f, grid_scale: %d / %d, strokes: %d / %d'\n",
        "          % (pt.step_id, pt.G_loss.item(), acc,\n",
        "              pt.m_grid, pt.max_divide,\n",
        "              pt.anchor_id, pt.m_strokes_per_block))\n",
        "    vis2 = utils.patches2img(pt.G_final_pred_canvas, pt.m_grid).clip(min=0, max=1)\n",
        "\n",
        "\n",
        "def optimize_x(pt):\n",
        "    pt._load_checkpoint()\n",
        "    pt.net_G.eval()\n",
        "    print('begin drawing...')\n",
        "\n",
        "    PARAMS = np.zeros([1, 0, pt.rderr.d], np.float32)\n",
        "\n",
        "    if pt.rderr.canvas_color == 'white':\n",
        "        CANVAS_tmp = torch.ones([1, 3, 128, 128]).to(device)\n",
        "    else:\n",
        "        CANVAS_tmp = torch.zeros([1, 3, 128, 128]).to(device)\n",
        "\n",
        "    for pt.m_grid in range(1, pt.max_divide + 1):\n",
        "\n",
        "        pt.img_batch = utils.img2patches(pt.img_, pt.m_grid).to(device)\n",
        "        pt.G_final_pred_canvas = CANVAS_tmp\n",
        "\n",
        "        pt.initialize_params()\n",
        "        pt.x_ctt.requires_grad = True\n",
        "        pt.x_color.requires_grad = True\n",
        "        pt.x_alpha.requires_grad = True\n",
        "        utils.set_requires_grad(pt.net_G, False)\n",
        "\n",
        "        pt.optimizer_x = optim.RMSprop([pt.x_ctt, pt.x_color, pt.x_alpha], lr=pt.lr, centered=True)\n",
        "\n",
        "        pt.step_id = 0\n",
        "        for pt.anchor_id in range(0, pt.m_strokes_per_block):\n",
        "            pt.stroke_sampler(pt.anchor_id)\n",
        "            iters_per_stroke = 80\n",
        "            for i in range(iters_per_stroke):\n",
        "                pt.G_pred_canvas = CANVAS_tmp\n",
        "\n",
        "                # update x\n",
        "                pt.optimizer_x.zero_grad()\n",
        "\n",
        "                pt.x_ctt.data = torch.clamp(pt.x_ctt.data, 0.1, 1 - 0.1)\n",
        "                pt.x_color.data = torch.clamp(pt.x_color.data, 0, 1)\n",
        "                pt.x_alpha.data = torch.clamp(pt.x_alpha.data, 0, 1)\n",
        "\n",
        "                pt._forward_pass()\n",
        "                _drawing_step_states(pt)\n",
        "                pt._backward_x()\n",
        "\n",
        "                pt.x_ctt.data = torch.clamp(pt.x_ctt.data, 0.1, 1 - 0.1)\n",
        "                pt.x_color.data = torch.clamp(pt.x_color.data, 0, 1)\n",
        "                pt.x_alpha.data = torch.clamp(pt.x_alpha.data, 0, 1)\n",
        "\n",
        "                pt.optimizer_x.step()\n",
        "                pt.step_id += 1\n",
        "\n",
        "        v = pt._normalize_strokes(pt.x)\n",
        "        PARAMS = np.concatenate([PARAMS, np.reshape(v, [1, -1, pt.rderr.d])], axis=1)\n",
        "        CANVAS_tmp = pt._render(PARAMS)[-1]\n",
        "        CANVAS_tmp = utils.img2patches(CANVAS_tmp, pt.m_grid + 1, to_tensor=True).to(device)\n",
        "\n",
        "    pt._save_stroke_params(PARAMS)\n",
        "    pt.final_rendered_images = pt._render(PARAMS)\n",
        "    pt._save_rendered_images()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b_QxLKdc-7nr"
      },
      "source": [
        "# レンダリング"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7yGUv65G-UCm"
      },
      "source": [
        "  pt = ProgressivePainter(args=args)\n",
        "  optimize_x(pt)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-p4XHLUh_ppN"
      },
      "source": [
        "# 画像表示"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lmYpxdcW_Azv"
      },
      "source": [
        "# show picture\n",
        "fig = plt.figure(figsize=(8,4))\n",
        "plt.subplot(1,2,1)\n",
        "plt.imshow(pt.img_), plt.title('input')\n",
        "plt.subplot(1,2,2)\n",
        "plt.imshow(pt.final_rendered_images[-1]), plt.title('generated')\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5LAIizo7_rxz"
      },
      "source": [
        "# make animation\n",
        "import matplotlib.animation as animation\n",
        "from IPython.display import HTML\n",
        "\n",
        "fig = plt.figure(figsize=(8,8))\n",
        "plt.axis('off')\n",
        "ims = [[plt.imshow(img, animated=True)] for img in pt.final_rendered_images[::10]]\n",
        "ani = animation.ArtistAnimation(fig, ims, interval=100)\n",
        "\n",
        "HTML(ani.to_jshtml())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AlDC5MmeYuYL"
      },
      "source": [
        "# save animation\n",
        "ani.save('anime.mp4', writer='ffmpeg')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}