Population Painter.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import imageio\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "from PIL import Image, ImageEnhance\n",
    "%matplotlib widget\n",
    "import matplotlib.pylab as plt\n",
    "\n",
    "from checkpointer import Checkpointer\n",
    "from defaults import get_cfg_defaults\n",
    "from model import Model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Loading config...')\n",
    "config_file='configs/popgan.yaml'\n",
    "cfg = get_cfg_defaults()\n",
    "cfg.merge_from_file(config_file)\n",
    "cfg.freeze()\n",
    "\n",
    "print('Initializing model...')\n",
    "torch.cuda.set_device(0)\n",
    "layer_count = cfg.MODEL.LAYER_COUNT\n",
    "model = Model(\n",
    "    startf=cfg.MODEL.START_CHANNEL_COUNT,\n",
    "    layer_count=cfg.MODEL.LAYER_COUNT,\n",
    "    maxf=cfg.MODEL.MAX_CHANNEL_COUNT,\n",
    "    latent_size=cfg.MODEL.LATENT_SPACE_SIZE,\n",
    "    truncation_psi=cfg.MODEL.TRUNCATIOM_PSI,\n",
    "    truncation_cutoff=cfg.MODEL.TRUNCATIOM_CUTOFF,\n",
    "    mapping_layers=cfg.MODEL.MAPPING_LAYERS,\n",
    "    channels=cfg.MODEL.CHANNELS,\n",
    "    generator=cfg.MODEL.GENERATOR,\n",
    "    encoder=cfg.MODEL.ENCODER)\n",
    "model.cuda(0)\n",
    "model.eval()\n",
    "model.requires_grad_(False)\n",
    "\n",
    "print('Loading checkpoint...')\n",
    "model_dict = {\n",
    "    'discriminator_s': model.encoder,\n",
    "    'generator_s': model.decoder,\n",
    "    'mapping_tl_s': model.mapping_tl,\n",
    "    'mapping_fl_s': model.mapping_fl,\n",
    "    'dlatent_avg': model.dlatent_avg\n",
    "}\n",
    "checkpointer = Checkpointer(cfg, model_dict)\n",
    "extra_checkpoint_data = checkpointer.load(file_name='training_artifacts/popgan/model_tmp_lod8_e200.pth')\n",
    "\n",
    "print('READY!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def encode(x, pop):\n",
    "    x_pop = torch.cat((x, pop), 1)\n",
    "    zlist = []\n",
    "    for i in range(x_pop.shape[0]):\n",
    "        Z, _ = model.encode(x_pop[i][None, ...], layer_count - 1, 1)\n",
    "        zlist.append(Z)\n",
    "    Z = torch.cat(zlist)\n",
    "    Z = Z.repeat(1, model.mapping_fl.num_layers, 1)\n",
    "    return Z\n",
    "\n",
    "noises = [[\n",
    "        0,\n",
    "        torch.randn([1, 1, 2, 2]),\n",
    "        torch.randn([1, 1, 4, 4]),\n",
    "        torch.randn([1, 1, 8, 8]),\n",
    "        torch.randn([1, 1, 16, 16]),\n",
    "        torch.randn([1, 1, 32, 32]),\n",
    "        torch.randn([1, 1, 64, 64]),\n",
    "        torch.randn([1, 1, 128, 128]),\n",
    "        torch.randn([1, 1, 256, 256]),\n",
    "        torch.randn([1, 1, 512, 512]),\n",
    "        torch.randn([1, 1, 1024, 1024])\n",
    "    ] for i in range(10)]\n",
    "\n",
    "def sample_z(seed):\n",
    "    with torch.no_grad():\n",
    "        rng = np.random.RandomState(seed)\n",
    "        z = rng.standard_normal(cfg.MODEL.LATENT_SPACE_SIZE)\n",
    "        z = torch.from_numpy(z).float().cuda().unsqueeze(0)\n",
    "    return z\n",
    "\n",
    "def z2w(z):\n",
    "    with torch.no_grad():\n",
    "        w = model.mapping_fl(z)\n",
    "    return w\n",
    "\n",
    "def decode(w, pop, noise_i=0):\n",
    "    return torch.clamp(model.decoder(w, pop, layer_count - 1, 1, noise=noises[noise_i]), -1., 1.)\n",
    "\n",
    "def mix_styles(style_shallow, style_deep, layer_range, strength):\n",
    "    style = style_shallow.clone()\n",
    "    style[layer_range, :] = (style_deep[layer_range, :] * strength) + (style_shallow[layer_range, :] * (1-strength))\n",
    "    return style\n",
    "\n",
    "def tensor_to_numpy(img):\n",
    "    img = img * 0.5 + 0.5\n",
    "    img = img.cpu().squeeze().numpy()\n",
    "    img = np.clip(img, 0, 1)\n",
    "    if len(img.shape) == 3:\n",
    "        img = np.moveaxis(img, 0, 2)\n",
    "    img = (img *255).astype(np.uint8)\n",
    "    return img\n",
    "\n",
    "def tensor_to_PIL(img):\n",
    "    img = Image.fromarray(tensor_to_numpy(img))\n",
    "    img = ImageEnhance.Brightness(img).enhance(1.4)\n",
    "    img = ImageEnhance.Contrast(img).enhance(1.4)\n",
    "    return img\n",
    "\n",
    "def load_image(path):\n",
    "    im_raw = Image.open(path)\n",
    "    im = im_raw.resize((1024, 1024))\n",
    "    im = np.asarray(im)\n",
    "    x = numpy_to_tensor(im)\n",
    "    return im_raw, x\n",
    "\n",
    "def numpy_to_tensor(im):\n",
    "    if len(im.shape) < 3:\n",
    "        im = np.expand_dims(im, 2)\n",
    "    im = im.transpose((2, 0, 1))\n",
    "    x = torch.tensor(np.asarray(im, dtype=np.float32), requires_grad=False).cuda() / 127.5 - 1.\n",
    "    return x\n",
    "\n",
    "def numpy_to_PIL(img):\n",
    "    img = np.clip(img, 0, 1)\n",
    "    img = (img *255).astype(np.uint8)\n",
    "    img = Image.fromarray(img)\n",
    "    img = img.resize((1024, 1024))\n",
    "    return img\n",
    "\n",
    "def save_mp4_from_images(dst, images):\n",
    "    writer = imageio.get_writer(dst, fps=10)\n",
    "    for img in images:\n",
    "        writer.append_data(img)\n",
    "    writer.close()\n",
    "    \n",
    "def db_to_tensor(im):\n",
    "    x = torch.tensor(np.asarray(im, dtype=np.float32), requires_grad=False).cuda(gpu) / 127.5 - 1.\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z_seed = 7\n",
    "noise_seed = 4\n",
    "brush_size = 2\n",
    "brush_strength = 0.5\n",
    "img_size = 1024\n",
    "pop_size = 64\n",
    "disp_size = 512\n",
    "\n",
    "def generate_pop(w, np_pop, noise_seed):\n",
    "    pop = numpy_to_tensor(np_pop*255)\n",
    "    x_rec = decode(w, pop.unsqueeze(0), noise_seed)\n",
    "    img = tensor_to_PIL(x_rec)\n",
    "    img = np.array(img)\n",
    "    return img\n",
    "\n",
    "fig, axes = plt.subplots(1,2, figsize=(10, 5))\n",
    "for a in axes:\n",
    "    a.get_xaxis().set_visible(False)\n",
    "    a.get_yaxis().set_visible(False)\n",
    "\n",
    "z = sample_z(z_seed)\n",
    "w = z2w(z)\n",
    "np_pop = np.zeros((pop_size, pop_size))\n",
    "resized_pop = np.array(Image.fromarray(np_pop).resize((img_size, img_size)))\n",
    "img = generate_pop(w, resized_pop, noise_seed)\n",
    "disp_img = np.array(Image.fromarray(img).resize((disp_size, disp_size)))\n",
    "axes[0].imshow(np_pop)\n",
    "axes[1].imshow(disp_img)\n",
    "\n",
    "debug = []\n",
    "events = []\n",
    "def onclick(event):\n",
    "    try:\n",
    "        if event.button.name in ['LEFT', 'RIGHT']:\n",
    "            direction = 1 if event.button.name == 'LEFT' else -1\n",
    "            x_min = int(max(0, event.xdata - brush_size))\n",
    "            x_max = int(min(pop_size, event.xdata + brush_size))\n",
    "            y_min = int(max(0, event.ydata - brush_size))\n",
    "            y_max = int(min(pop_size, event.ydata + brush_size))\n",
    "            \n",
    "            debug.append({\n",
    "                'event': event.button.name,\n",
    "                'direction': direction,\n",
    "                'x_min': x_min,\n",
    "                'x_max': x_max,\n",
    "                'y_min': y_min,\n",
    "                'y_max': y_max,\n",
    "                'pop_shape': np_pop[x_min:x_max, y_min:y_max].shape,\n",
    "                'brush_effect': brush_strength * direction\n",
    "            })\n",
    "\n",
    "            np_pop[y_min:y_max, x_min:x_max] += (brush_strength * direction)\n",
    "            np_pop[:, :] = np.clip(np_pop[:, :], 0., 1.)\n",
    "            axes[0].imshow(np_pop)\n",
    "    except Exception as e:\n",
    "        debug.append(e)\n",
    "        \n",
    "def onpress(event):\n",
    "    try:\n",
    "        events.append(event)\n",
    "        if event.key == 'w':\n",
    "            resized_pop = np.array(Image.fromarray(np_pop).resize((img_size, img_size)))\n",
    "            img = generate_pop(w, resized_pop, noise_seed)\n",
    "            disp_img = np.array(Image.fromarray(img).resize((disp_size, disp_size)))\n",
    "            axes[1].imshow(disp_img)\n",
    "    except Exception as e:\n",
    "        debug.append(e)\n",
    "\n",
    "fig.canvas.mpl_connect('button_press_event', onclick)\n",
    "fig.canvas.mpl_connect('key_press_event', onpress)\n",
    "fig.show()\n",
    "fig.tight_layout()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py37_pytorch",
   "language": "python",
   "name": "py37_pytorch"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}