Population Painter Colab.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SCALAE\n",
    "\n",
    "Interactive implementation of https://github.com/LendelTheGreat/SCALAE\n",
    "\n",
    "Run all and scroll to the bottom!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Code and model downloading\n",
    "!pip install yacs\n",
    "!git clone https://github.com/LendelTheGreat/SCALAE.git\n",
    "!gdown --id 1pqjjx8zRSPsTzPXVQiFf3PgmGnJpLSQn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Model setup\n",
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "from IPython.display import HTML, Image, clear_output\n",
    "from google.colab.output import eval_js\n",
    "from base64 import b64decode, b64encode\n",
    "import PIL\n",
    "import PIL.ImageOps\n",
    "import io\n",
    "\n",
    "SCALAE_PATH = '/content/SCALAE'\n",
    "if SCALAE_PATH not in sys.path:\n",
    "    sys.path.append(SCALAE_PATH)\n",
    "from checkpointer import Checkpointer\n",
    "from defaults import get_cfg_defaults\n",
    "from model import Model\n",
    "\n",
    "\n",
    "print('Loading config...')\n",
    "config_file=f'{SCALAE_PATH}/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()\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='/content/scalae_pop2sat_1024_e200.pth')\n",
    "\n",
    "print('READY!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Helper functions\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 sample_noise(seed):\n",
    "  rng = np.random.RandomState(seed)\n",
    "  noise = [0] + [torch.from_numpy(rng.standard_normal((1, 1, 2 ** i, 2 ** i)).astype(np.float32)) for i in range(1, 11)]\n",
    "  return noise\n",
    "\n",
    "def decode(w, pop, noise):\n",
    "  return torch.clamp(model.decoder(w, pop, layer_count - 1, 1, noise=noise), -1., 1.)\n",
    "\n",
    "def numpy_to_tensor(img):\n",
    "  if len(img.shape) < 3:\n",
    "    img = np.expand_dims(img, 2)\n",
    "  img = img.transpose((2, 0, 1))\n",
    "  x = torch.tensor(np.asarray(img, dtype=np.float32), requires_grad=False).cuda() / 0.5 - 1.\n",
    "  return x\n",
    "\n",
    "def tensor_to_numpy(img):\n",
    "  img = img * 0.5 + 0.5\n",
    "  img = img.cpu().squeeze().numpy()\n",
    "  if len(img.shape) == 3:\n",
    "    img = np.moveaxis(img, 0, 2)\n",
    "  return img\n",
    "\n",
    "def generate(w, np_pop, noise):\n",
    "  pop = numpy_to_tensor(np_pop).unsqueeze(0)\n",
    "  img = torch.clamp(model.decoder(w, pop, layer_count - 1, 1, noise=noise), -1., 1.)\n",
    "  img = tensor_to_numpy(img)\n",
    "  return img\n",
    "\n",
    "def numpy_to_PIL(img, img_size=256):\n",
    "  img = np.clip(img, 0, 1)\n",
    "  img = (img *255).astype(np.uint8)\n",
    "  img = PIL.Image.fromarray(img)\n",
    "  img = img.resize((img_size, img_size))\n",
    "  return img\n",
    "\n",
    "canvas_html = \"\"\"\n",
    "<div class=\"slidecontainer\">\n",
    "  <label>Brush strengt (corresponds to population density)</label>\n",
    "  <input style=\"width:%dpx\" type=\"range\" min=\"0\" max=\"10\" value=\"10\" class=\"slider\" id=\"pop_slider\">\n",
    "</div>\n",
    "<div class=\"slidecontainer\">\n",
    "  <label>Brush width</label>\n",
    "  <input style=\"width:%dpx\" type=\"range\" min=\"5\" max=\"100\" value=\"20\" class=\"slider\" id=\"line_slider\">\n",
    "</div>\n",
    "<div>\n",
    "  <canvas width=%d height=%d></canvas>\n",
    "  <img src='data:image/jpeg;base64,%s'/>\n",
    "</div>\n",
    "<button style=\"height:20px;width:%dpx\" id=\"1\">Generate Fake Satellite Image</button>\n",
    "<label></label>\n",
    "<script>\n",
    "var label =  document.querySelector('label')\n",
    "var canvas = document.querySelector('canvas')\n",
    "var ctx = canvas.getContext('2d')\n",
    "var line_slider = document.getElementById(\"line_slider\");\n",
    "ctx.lineWidth = line_slider.value\n",
    "line_slider.oninput = function() {\n",
    "  ctx.lineWidth = this.value;\n",
    "}\n",
    "var pop_slider = document.getElementById(\"pop_slider\");\n",
    "ctx.strokeStyle = \"#\" + (pop_slider.value * 25).toString(16) + (pop_slider.value * 25).toString(16) + (pop_slider.value * 25).toString(16);\n",
    "pop_slider.oninput = function() {\n",
    "  ctx.strokeStyle = \"#\" + (this.value * 25).toString(16) + (this.value * 25).toString(16) + (this.value * 25).toString(16);\n",
    "}\n",
    "pop_image = new Image();\n",
    "pop_image.src = 'data:image/png;base64,%s';\n",
    "pop_image.onload = function(){\n",
    "  ctx.drawImage(pop_image, 0, 0);\n",
    "}\n",
    "var button_generate = document.getElementById('1')\n",
    "var mouse = {x: 0, y: 0}\n",
    "canvas.addEventListener('mousemove', function(e) {\n",
    "  mouse.x = e.pageX - this.offsetLeft\n",
    "  mouse.y = e.pageY - this.offsetTop\n",
    "})\n",
    "canvas.onmousedown = ()=>{\n",
    "  ctx.beginPath()\n",
    "  ctx.moveTo(mouse.x, mouse.y)\n",
    "  canvas.addEventListener('mousemove', onPaint)\n",
    "}\n",
    "canvas.onmouseup = ()=>{\n",
    "  canvas.removeEventListener('mousemove', onPaint)\n",
    "}\n",
    "var onPaint = ()=>{\n",
    "  ctx.lineTo(mouse.x, mouse.y)\n",
    "  ctx.stroke()\n",
    "}\n",
    "var data = new Promise(resolve=>{\n",
    "  button_generate.onclick = ()=>{\n",
    "    resolve(canvas.toDataURL('image/png'))\n",
    "  }\n",
    "})\n",
    "</script>\n",
    "\"\"\"\n",
    "\n",
    "def setup_html(img_size, pop_file, gen_file):\n",
    "  pop_str = b64encode(imgfile_to_string(pop_file)).decode(\"utf-8\")\n",
    "  gen_str = b64encode(imgfile_to_string(gen_file)).decode(\"utf-8\")\n",
    "  return HTML(canvas_html % (img_size, img_size, img_size, img_size, gen_str, img_size, pop_str))\n",
    "\n",
    "def imgfile_to_string(file):\n",
    "  an_image = PIL.Image.open(file)\n",
    "  output = io.BytesIO()\n",
    "  an_image.save(output, format=\"png\")\n",
    "  image_as_string = output.getvalue()\n",
    "  return image_as_string\n",
    "\n",
    "def run(img_size=256, seed=0):\n",
    "  pop = np.zeros((1024, 1024))\n",
    "  numpy_to_PIL(pop, img_size).save(\"pop.png\")\n",
    "\n",
    "  z = sample_z(seed)\n",
    "  w = z2w(z)\n",
    "  noise = sample_noise(seed)\n",
    "  \n",
    "  img = generate(w, pop, noise)\n",
    "  img = numpy_to_PIL(img, img_size)\n",
    "  img.save(\"img.png\")\n",
    "\n",
    "  while True:\n",
    "    clear_output()\n",
    "    display(setup_html(img_size, 'pop.png', \"img.png\"))\n",
    "\n",
    "    data = eval_js(\"data\").split(',')[1]\n",
    "\n",
    "    binary = b64decode(data)\n",
    "    with open(\"pop.png\", 'wb') as f:\n",
    "      f.write(binary)\n",
    "\n",
    "    pop = PIL.Image.open(\"pop.png\")\n",
    "    pop = pop.resize((1024, 1024))\n",
    "    pop = np.array(pop) / 255\n",
    "    pop = pop[:, :, 0]\n",
    "\n",
    "    img = generate(w, pop, noise)\n",
    "    img = numpy_to_PIL(img, img_size)\n",
    "    img.save(\"img.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Instructions\n",
    "\n",
    "On the left you can paint in the black canvas. This is the population density input to the model. On the right you see the generated model output.\n",
    "\n",
    "Click the `Generate Fake Satellite Image` button at the bottom to generate an image.\n",
    "\n",
    "You can adjust the painting brush with the sliders on top."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run(img_size=512, seed=42)"
   ]
  }
 ],
 "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
}