Labeller.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Modified from @bryandlee"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install ipympl\n",
    "# !jupyter labextension list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import numpy as np\n",
    "import time\n",
    "import cv2\n",
    "import pickle\n",
    "\n",
    "def tensor2image(tensor):\n",
    "    tensor = tensor.clamp_(-1., 1.).detach().squeeze().permute(1,2,0).cpu().numpy()\n",
    "    return tensor*0.5 + 0.5\n",
    "\n",
    "def imshow(img, size=5, cmap='jet'):\n",
    "    %matplotlib inline\n",
    "    plt.figure(figsize=(size,size))\n",
    "    plt.imshow(img, cmap=cmap)\n",
    "    plt.axis('off')\n",
    "    plt.show()\n",
    "\n",
    "def horizontal_concat(imgs):\n",
    "    return torch.cat([img.unsqueeze(0) for img in imgs], 3) \n",
    "\n",
    "def save_generated_image(tensor):\n",
    "    imgs = tensor.clamp_(-1., 1.).detach().squeeze().permute(0,2,3,1).cpu().numpy()\n",
    "    imgs = imgs*0.5 + 0.5\n",
    "    for i in range(imgs.shape[0]):\n",
    "        filename = f'image_generated_{i}.png'\n",
    "        plt.imsave(filename, imgs[i])\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pre-Trained Generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[StyleGAN2 generator loaded] ./checkpoint/550000.pt\n",
      "\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "CUDA out of memory. Tried to allocate 2.62 GiB (GPU 0; 11.78 GiB total capacity; 2.97 GiB already allocated; 1.78 GiB free; 2.99 GiB reserved in total by PyTorch)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-2-fd82ff45b39f>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     26\u001b[0m                                    \u001b[0minput_is_latent\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     27\u001b[0m                                    \u001b[0mrandomize_noise\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 28\u001b[0;31m                                    concat_features=True)\n\u001b[0m\u001b[1;32m     29\u001b[0m     \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcuda\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mempty_cache\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     30\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    725\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    726\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 727\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    728\u001b[0m         for hook in itertools.chain(\n\u001b[1;32m    729\u001b[0m                 \u001b[0m_global_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/regan/model/stylegan_model.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, styles, return_latents, inject_index, truncation, truncation_latent, input_is_latent, noise, randomize_noise, concat_features)\u001b[0m\n\u001b[1;32m    645\u001b[0m         \u001b[0mimage\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mskip\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    646\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mconcat_features\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 647\u001b[0;31m             \u001b[0mfeatures\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mconcat_f\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfeatures\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    648\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mreturn_latents\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    649\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mimage\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlatent\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/torch/autograd/grad_mode.py\u001b[0m in \u001b[0;36mdecorate_context\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     24\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mdecorate_context\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     25\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__class__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 26\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     27\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mcast\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mF\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdecorate_context\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     28\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/regan/model/stylegan_model.py\u001b[0m in \u001b[0;36mconcat_f\u001b[0;34m(features)\u001b[0m\n\u001b[1;32m    458\u001b[0m     \u001b[0mh\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmax\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mf\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mfeatures\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    459\u001b[0m     \u001b[0mw\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmax\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mf\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mfeatures\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 460\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfunctional\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minterpolate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mh\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mw\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmode\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'nearest'\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mf\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mfeatures\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    461\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    462\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: CUDA out of memory. Tried to allocate 2.62 GiB (GPU 0; 11.78 GiB total capacity; 2.97 GiB already allocated; 1.78 GiB free; 2.99 GiB reserved in total by PyTorch)"
     ]
    }
   ],
   "source": [
    "############ self - config ##############\n",
    "device = 'cuda:0'\n",
    "image_size = 256\n",
    "n_samples = 2   #n-shot segmentation\n",
    "generator_path = './checkpoint/550000.pt'\n",
    "latent_dim = 512\n",
    "truncation = 0.7\n",
    "imshow_size = 3\n",
    "#########################################\n",
    "\n",
    "from model.stylegan_model import Generator\n",
    "\n",
    "generator = Generator(image_size, latent_dim, 8)\n",
    "generator_ckpt = torch.load(generator_path, map_location='cpu')\n",
    "generator.load_state_dict(generator_ckpt[\"g_ema\"], strict=False)\n",
    "generator.eval().to(device)\n",
    "print(f'[StyleGAN2 generator loaded] {generator_path}\\n')\n",
    "\n",
    "\n",
    "with torch.no_grad():\n",
    "    trunc_mean = generator.mean_latent(4096).detach().clone()\n",
    "    latent = generator.get_latent(torch.randn(n_samples, latent_dim, device=device)) #random latent vector z 생성\n",
    "    imgs_gen, features = generator([latent],\n",
    "                                   truncation=truncation,\n",
    "                                   truncation_latent=trunc_mean,\n",
    "                                   input_is_latent=True,\n",
    "                                   randomize_noise=True,\n",
    "                                   concat_features=True)\n",
    "    torch.cuda.empty_cache()\n",
    "\n",
    "print(\"sample images:\")\n",
    "imshow(tensor2image(horizontal_concat(imgs_gen)), size=imshow_size*n_samples)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## save image and vector\n",
    "\n",
    "latent vector,\n",
    "images_gen,\n",
    "features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import pickle\n",
    "# import os\n",
    "# import cv2\n",
    "\n",
    "\n",
    "# ## Save image\n",
    "# save_generated_image(imgs_gen)    \n",
    "# os.makedirs('dataset', exist_ok=True)\n",
    "# data_list = {'image' : imgs_gen, 'features':features,'latent':latent}\n",
    "\n",
    "# ## Save pickle\n",
    "# with open(\"./dataset/dataset.pkl\",\"wb\") as fw:\n",
    "#     pickle.dump(data_list, fw)\n",
    "    \n",
    "# import matplotlib"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Few-Shot Labelling with a Simple Interface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%matplotlib widget\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import widgets\n",
    "from matplotlib.colors import hsv_to_rgb\n",
    "# print(matplotlib.get_backend())\n",
    "\n",
    "import random\n",
    "import copy\n",
    "import cv2\n",
    "\n",
    "class Labeller:\n",
    "    def __init__(self, images, classes):\n",
    "        \n",
    "        self.n_image = len(images)\n",
    "        self.img_idx = 0\n",
    "\n",
    "        self.images = images        \n",
    "        self._reset_label()\n",
    "        \n",
    "        self.classes = classes\n",
    "        self.n_class = len(classes)\n",
    "        self.colors = self._sample_colors(self.n_class)\n",
    "        self.colors[0] = np.array([1., 1., 1.])\n",
    "        self._class = 1\n",
    "\n",
    "        self.fig = plt.figure('Simple Labeller')\n",
    "        self.ax = self.fig.add_subplot()\n",
    "        self.fig.subplots_adjust(left=0.0, bottom=0.0, right=0.80, top=1.0)\n",
    "\n",
    "        self.ax.axis('off')\n",
    "        self.ax_img = self.ax.imshow(self.images[self.img_idx])\n",
    "        \n",
    "        self._add_buttons()        \n",
    "        self.fig.canvas.mpl_connect('key_press_event', self._key_maps)\n",
    "\n",
    "        self.show_overlay = True\n",
    "        self.history = []\n",
    "        self.brush_size = 1\n",
    "        \n",
    "        plt.show()\n",
    "       \n",
    "    def _sample_colors(self, n=1):\n",
    "        h = np.linspace(0.0, 1.0, n)[:,np.newaxis]\n",
    "        s = np.ones((n,1))*0.5\n",
    "        v = np.ones((n,1))*1.0\n",
    "        print(np.concatenate([h,s,v], axis=1))\n",
    "        return hsv_to_rgb(np.concatenate([h,s,v], axis=1))\n",
    "\n",
    "    def _draw(self, image):\n",
    "        self.ax_img.set_data(image)\n",
    "        \n",
    "    def _key_maps(self, event):\n",
    "        key_maps = {\n",
    "            'c': self._lasso,\n",
    "            'v': self._poly,\n",
    "            'z': self._undo,\n",
    "            'right': self._next_class,\n",
    "            'left': self._prev_class,\n",
    "            'o': self._overlay,\n",
    "            'up': self._brush_up,\n",
    "            'down': self._brush_down,\n",
    "        }\n",
    "        key = event.key.lower()\n",
    "        if key in key_maps:\n",
    "            key_maps[key](None)\n",
    "        \n",
    "        \n",
    "    def _add_buttons(self):\n",
    "        axes_coords = [0.84, 0.94, 0.15, 0.05]\n",
    "        interval = 0.08\n",
    "\n",
    "        self.class_box = widgets.Button(plt.axes(axes_coords),\n",
    "                                        self.classes[self._class],\n",
    "                                        color=self.colors[self._class],\n",
    "                                        hovercolor=self.colors[self._class])\n",
    "\n",
    "        axes_coords[1] -= interval\n",
    "        axes_coords_split = copy.deepcopy(axes_coords)\n",
    "        axes_coords_split[2] = 0.06\n",
    "        self.prev_class = widgets.Button(plt.axes(axes_coords_split), '<')\n",
    "        self.cid_prev_class = self.prev_class.on_clicked(self._prev_class)\n",
    "        axes_coords_split[0] = 0.84 + 0.15 - 0.06\n",
    "        self.next_class = widgets.Button(plt.axes(axes_coords_split), '>')\n",
    "        self.cid_next_class = self.next_class.on_clicked(self._next_class)\n",
    "        \n",
    "        axes_coords[1] -= interval\n",
    "        self.lasso = widgets.Button(plt.axes(axes_coords), 'Brush (C)')\n",
    "        self.cid_lasso = self.lasso.on_clicked(self._lasso)\n",
    "\n",
    "        axes_coords[1] -= interval\n",
    "        axes_coords_split = copy.deepcopy(axes_coords)\n",
    "        axes_coords_split[2] = 0.06\n",
    "        self.brush_up = widgets.Button(plt.axes(axes_coords_split), '+')\n",
    "        self.cid_brush_up = self.brush_up.on_clicked(self._brush_up)\n",
    "        axes_coords_split[0] = 0.84 + 0.15 - 0.06\n",
    "        self.brush_down = widgets.Button(plt.axes(axes_coords_split), '-')\n",
    "        self.cid_brush_down = self.brush_down.on_clicked(self._brush_down)\n",
    "        \n",
    "        axes_coords[1] -= interval\n",
    "        self.poly = widgets.Button(plt.axes(axes_coords), 'Polygon (V)')\n",
    "        self.cid_poly = self.poly.on_clicked(self._poly)\n",
    "        \n",
    "        axes_coords[1] -= interval\n",
    "        self.undo = widgets.Button(plt.axes(axes_coords), 'Undo (Z)')\n",
    "        self.cid_undo = self.undo.on_clicked(self._undo)\n",
    "\n",
    "        axes_coords[1] -= interval\n",
    "        self.overlay = widgets.Button(plt.axes(axes_coords), 'Overlay (O)')\n",
    "        self.cid_overlay = self.overlay.on_clicked(self._overlay)\n",
    "        \n",
    "        axes_coords[1] -= interval\n",
    "        self.reset = widgets.Button(plt.axes(axes_coords), 'Reset',\n",
    "                                    color=[1, 0.3, 0.3], hovercolor=[1, 0.5, 0.5])\n",
    "        self.cid_reset = self.reset.on_clicked(self._reset)\n",
    "        \n",
    "        axes_coords[1] -= interval\n",
    "        axes_coords_split = copy.deepcopy(axes_coords)\n",
    "        axes_coords_split[2] = 0.06\n",
    "        self.prev_img = widgets.Button(plt.axes(axes_coords_split), 'Prev')\n",
    "        self.cid_prev_img = self.prev_img.on_clicked(self._prev_img)\n",
    "        axes_coords_split[0] = 0.84 + 0.15 - 0.06\n",
    "        self.next_img = widgets.Button(plt.axes(axes_coords_split), 'Next')\n",
    "        self.cid_next_img = self.next_img.on_clicked(self._next_img)\n",
    "        \n",
    "    def _next_class(self, event):\n",
    "        self._class = (self._class + 1) % self.n_class\n",
    "        self._update_class_box()\n",
    "        \n",
    "    def _prev_class(self, event):\n",
    "        self._class = (self._class - 1) % self.n_class\n",
    "        self._update_class_box()\n",
    "\n",
    "    def _update_class_box(self):\n",
    "        self.class_box.label.set_text(self.classes[self._class])\n",
    "        self.class_box.color = self.colors[self._class]\n",
    "        self.class_box.hovercolor = self.class_box.color\n",
    "        self.fig.canvas.draw()\n",
    "        \n",
    "    def _undo(self, event):\n",
    "        if len(self.history) > 0:\n",
    "            self.history.pop(-1)\n",
    "            self._reset_label(only_current_img=True)\n",
    "            for inputs in self.history:\n",
    "                self._update_label(inputs)\n",
    "            self._draw(self.get_image_label_overlay())\n",
    "\n",
    "    def _overlay(self, event):\n",
    "        self.show_overlay = not self.show_overlay\n",
    "        if self.show_overlay:\n",
    "            self._draw(self.get_image_label_overlay())\n",
    "        else:\n",
    "            self._draw(self.images[self.img_idx])\n",
    "    \n",
    "    def _reset(self, event):\n",
    "        self.history = []\n",
    "        self._reset_label(only_current_img=True)\n",
    "        self._draw(self.images[self.img_idx])\n",
    "\n",
    "    def _next_img(self, event):\n",
    "        self.img_idx = (self.img_idx + 1) % self.n_image\n",
    "        self._on_img_change()\n",
    "\n",
    "    def _prev_img(self, event):\n",
    "        self.img_idx = (self.img_idx - 1) % self.n_image\n",
    "        self._on_img_change()\n",
    "\n",
    "    def _on_img_change(self):\n",
    "        self.history = []\n",
    "        self.show_overlay = True        \n",
    "        self._draw(self.get_image_label_overlay())\n",
    "        \n",
    "    def _poly(self, event):\n",
    "        self._reset_selectors()\n",
    "        self.poly_selector = widgets.PolygonSelector(self.ax, self._process_polygon)\n",
    "    \n",
    "    def _process_polygon(self, vert):\n",
    "        polygon = np.array(vert, np.int32).reshape((-1, 1, 2))\n",
    "        inputs = ('poly', polygon, self._class)\n",
    "        self.history.append(inputs)\n",
    "        self._update_label(inputs)\n",
    "        self._after_new_label()\n",
    "        self._reset_selectors()\n",
    "\n",
    "    def _brush_up(self, dummy):\n",
    "        self.brush_size += 1\n",
    "        \n",
    "    def _brush_down(self, dummy):\n",
    "        self.brush_size = max(self.brush_size-1, 1)\n",
    "        \n",
    "    def _lasso(self, event):\n",
    "        self._reset_selectors()        \n",
    "        self.lasso_selector = widgets.LassoSelector(\n",
    "            self.ax, self._process_lasso, lineprops=dict(linewidth=self.brush_size//2)\n",
    "        )\n",
    "    \n",
    "    def _process_lasso(self, vert):\n",
    "        path = np.array(vert, np.int32).reshape((-1, 1, 2))\n",
    "        path = np.unique(path, axis=1)\n",
    "        inputs = ('lasso', path, self._class, self.brush_size)\n",
    "        self.history.append(inputs)\n",
    "        self._update_label(inputs)\n",
    "        self._after_new_label()\n",
    "        self._reset_selectors()\n",
    "\n",
    "    def _reset_selectors(self):\n",
    "        if hasattr(self, 'lasso_selector'):\n",
    "            self.lasso_selector.set_visible(False)\n",
    "            del(self.lasso_selector)\n",
    "        if hasattr(self, 'poly_selector'):\n",
    "            self.poly_selector.set_visible(False)            \n",
    "            del(self.poly_selector)\n",
    "    \n",
    "    def _after_new_label(self):\n",
    "        self.show_overlay = True\n",
    "        self._draw(self.get_image_label_overlay())\n",
    "            \n",
    "    def _reset_label(self, only_current_img=False):\n",
    "        if only_current_img:\n",
    "            self.labels[self.img_idx] = np.zeros(\n",
    "                (self.images.shape[1], self.images.shape[2]),\n",
    "                np.uint8\n",
    "            )\n",
    "        else:\n",
    "            self.labels = np.zeros(\n",
    "                (self.n_image, self.images.shape[1], self.images.shape[2]),\n",
    "                np.uint8\n",
    "            )\n",
    "        \n",
    "    def _update_label(self, inputs):\n",
    "        if inputs[0] == 'poly':\n",
    "            self.labels[self.img_idx] = cv2.fillPoly(\n",
    "                self.labels[self.img_idx], [inputs[1]], inputs[2], 0\n",
    "            )\n",
    "        elif inputs[0] == 'lasso':\n",
    "            self.labels[self.img_idx] = cv2.polylines(\n",
    "                self.labels[self.img_idx], [inputs[1]], isClosed=False,\n",
    "                color=inputs[2], thickness=inputs[3]\n",
    "            )\n",
    "        \n",
    "    def get_image_label_overlay(self):\n",
    "        overlay = self.images[self.img_idx].copy()\n",
    "        label_image = self.get_visualized_label()\n",
    "        non_zeros = label_image > 0\n",
    "        overlay[non_zeros] = label_image[non_zeros]\n",
    "        return overlay\n",
    "        \n",
    "    def get_visualized_label(self, label=None):\n",
    "        if label is None:\n",
    "            label = self.labels[self.img_idx]\n",
    "        \n",
    "        label_image = np.zeros_like(self.images[self.img_idx])\n",
    "        for c in range(1,self.n_class):\n",
    "            label_image[label == c] = self.colors[c]\n",
    "        return label_image\n",
    "    \n",
    "    def get_labels(self):\n",
    "        return self.labels\n",
    "\n",
    "# set classes here\n",
    "classes = ['background', 'skin', 'hair', 'eye', 'eyebrow', 'nose', 'mouth', 'ear']\n",
    "        \n",
    "labeller = Labeller(imgs_gen.clamp_(-1., 1.).detach().permute(0,2,3,1).cpu().numpy() * 0.5 + 0.5, classes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = labeller.get_labels()\n",
    "\n",
    "\n",
    "imshow(tensor2image(horizontal_concat(imgs_gen)), size=imshow_size*n_samples)\n",
    "imshow(np.concatenate([labeller.get_visualized_label(l) for l in labels], axis=1), size=imshow_size*n_samples)\n",
    "\n",
    "@torch.no_grad()\n",
    "def concat_features(features):\n",
    "    h = max([f.shape[-2] for f in features])\n",
    "    w = max([f.shape[-1] for f in features])\n",
    "    return torch.cat([torch.nn.functional.interpolate(f, (h,w), mode='nearest') for f in features], dim=1)\n",
    "# data = dataset\n",
    "data = dict(\n",
    "    images = imgs_gen,\n",
    "    latents=latent.cpu(),\n",
    "    features=features.cpu(),  #input img 사이즈로 representation vector concatenate\n",
    "    labels=torch.tensor(labels).long(),\n",
    ")\n",
    "print(f'Dataset for {n_samples}-Shot Training is Prepared')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "with open(\"./dataset/dataset.pkl\",\"wb\") as fw:\n",
    "    pickle.dump(data, fw)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "from matplotlib.colors import hsv_to_rgb\n",
    "\n",
    "\n",
    "def _sample_colors(n):\n",
    "        h = np.linspace(0.0, 1.0, n)[:,np.newaxis]\n",
    "        s = np.ones((n,1))*0.5\n",
    "        v = np.ones((n,1))*1.0\n",
    "        return hsv_to_rgb(np.concatenate([h,s,v], axis=1))\n",
    "    \n",
    "def get_visualized_label(label, class_num):\n",
    "    colors = _sample_colors(class_num)\n",
    "    colors[0] = np.array([1., 1., 1.])\n",
    "    print(colors)\n",
    "    label_image = np.zeros_like(3,label[1].shape,label[2].shape)\n",
    "    \n",
    "    for c in range(1,class_num):\n",
    "        label_image[label == c] = colors[c]\n",
    "    return label_image\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   50-th epoch | loss: 0.0840 | time:   37.0sec\n",
      "  100-th epoch | loss: 0.0290 | time:   78.0sec\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "import torch.nn.functional as F\n",
    "import pickle\n",
    "import torch\n",
    "import numpy as np\n",
    "import time\n",
    "import random\n",
    "import copy\n",
    "import cv2\n",
    "\n",
    "from model.segmentation_model import FewShotCNN\n",
    "n_samples = 2\n",
    "PATH = './checkpoint/'\n",
    "device = 'cuda'\n",
    "\n",
    "with open(\"./dataset/dataset.pkl\",\"rb\") as fw:\n",
    "    data = pickle.load(fw)\n",
    "    \n",
    "classes = ['background', 'skin', 'hair', 'eye', 'eyebrow', 'nose', 'mouth', 'ear']\n",
    "\n",
    "net = FewShotCNN(data['features'].shape[1], len(classes), size='S')\n",
    "\n",
    "optimizer = torch.optim.Adam(net.parameters(), lr=0.001, weight_decay=0.001)\n",
    "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.1)\n",
    "labels = torch.tensor(data['labels']).long()\n",
    "net.train().to(device)\n",
    "start_time = time.time()\n",
    "\n",
    "for epoch in range(1, 100+1):\n",
    "    sample_order = list(range(n_samples))\n",
    "    random.shuffle(sample_order)\n",
    "\n",
    "    for idx in sample_order:\n",
    "        \n",
    "        sample = data['features'][idx].unsqueeze(0).to(device)\n",
    "        label = labels[idx].unsqueeze(0).to(device)\n",
    "        \n",
    "        out = net(sample)\n",
    "        \n",
    "        loss = F.cross_entropy(out, label, reduction='mean')\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    if epoch % 50 == 0:\n",
    "        print(f'{epoch:5}-th epoch | loss: {loss.item():6.4f} | time: {time.time()-start_time:6.1f}sec')\n",
    "\n",
    "    scheduler.step()\n",
    "    torch.save(net, PATH + 'FewShotCNN.pt')  # 전체 모델 저장\n",
    "print('Done!')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_save_path = './dataset/images'\n",
    "label_save_path = './dataset/labels'\n",
    "\n",
    "def save_img_and_label(imgs_gen, predictions):\n",
    "    imgs_gen = imgs_gen.clamp_(-1., 1.).detach().squeeze().permute(0,2,3,1).cpu().numpy()\n",
    "    imgs_gen = imgs_gen*0.5 + 0.5\n",
    "    for i in range(imgs_gen.shape[0]):\n",
    "        img_filename = f'generated_data_{i}.png'\n",
    "        label_filename = f'generated_label_{i}.png'\n",
    "        plt.imsave(os.path.join(img_save_path,img_filename), imgs_gen[i])\n",
    "        plt.imsave(os.path.join(label_save_path,label_filename), predictions[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([3, 3, 256, 256])\n",
      "(3, 256, 256)\n",
      "torch.Size([3, 5376, 256, 256])\n",
      "<class 'torch.Tensor'>\n"
     ]
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "device = 'cpu'\n",
    "n_test = 3\n",
    "image_size = 256\n",
    "class_num = len(classes)\n",
    "generator_path = './checkpoint/550000.pt'\n",
    "latent_dim = 512\n",
    "truncation = 0.7\n",
    "\n",
    "\n",
    "from model.stylegan_model import Generator\n",
    "generator = Generator(image_size, latent_dim, 8)\n",
    "generator_ckpt = torch.load(generator_path, map_location='cpu')\n",
    "generator.load_state_dict(generator_ckpt[\"g_ema\"], strict=False)\n",
    "generator.eval().to(device)\n",
    "net.eval().to(device)\n",
    "\n",
    "with torch.no_grad():\n",
    "    trunc_mean = generator.mean_latent(4096).detach().clone()\n",
    "    latent = generator.get_latent(torch.randn(n_test, latent_dim, device=device))\n",
    "    imgs_gen, features = generator([latent],\n",
    "                                   truncation=truncation,\n",
    "                                   truncation_latent=trunc_mean.to(device),\n",
    "                                   input_is_latent=True,\n",
    "                                   randomize_noise=True)\n",
    "\n",
    "    torch.cuda.empty_cache()\n",
    "    out = net(features)\n",
    "    predictions = out.data.max(1)[1].cpu().numpy()\n",
    "    \n",
    "    save_img_and_label(imgs_gen, predictions)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "d88d1844a0fbbb13bf1f929fffff709e062e8c90e068b01d73bef27e583553a9"
  },
  "kernelspec": {
   "display_name": "Python 3.9.6 64-bit ('eff': conda)",
   "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}