infer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CoMoGan\n",
    "Notebook to test our model after training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets as widgets\n",
    "import pytorch_lightning as pl\n",
    "import pathlib\n",
    "import torch\n",
    "import yaml\n",
    "import os\n",
    "\n",
    "from math import pi\n",
    "from PIL import Image\n",
    "from munch import Munch\n",
    "from threading import Timer\n",
    "from IPython.display import clear_output\n",
    "from torchvision.transforms import ToPILImage\n",
    "\n",
    "from data import create_dataset\n",
    "from torchvision.transforms import ToTensor\n",
    "from data.base_dataset import get_transform\n",
    "from networks import find_model_using_name, create_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the model with a checkpoint\n",
    "Choose the directory that contains the checkpoint that you want."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load names of directories inside /logs\n",
    "p = pathlib.Path('./logs')\n",
    "list_run_id = [str(x).split('/')[1] for x in p.iterdir() if x.is_dir()]\n",
    "\n",
    "w_run = widgets.Dropdown(options=list_run_id,\n",
    "                         description='Select RUN_ID',\n",
    "                         disabled=False,\n",
    "                         style=dict(description_width='initial'))\n",
    "\n",
    "w_check = None\n",
    "root_dir = None\n",
    "\n",
    "def on_value_change_check(change):\n",
    "    global w_check, w_run, root_dir\n",
    "    \n",
    "    clear_output(wait=True)\n",
    "    \n",
    "    root_dir = os.path.join('logs', w_run.value, 'tensorboard', 'default', 'version_0')\n",
    "    p = pathlib.Path(root_dir + '/checkpoints')\n",
    "    # Load a list of checkpoints, use the last one by default\n",
    "    list_checkpoint = [str(x).split('checkpoints/')[1] for x in p.iterdir() if 'iter' in str(x)]\n",
    "    list_checkpoint.sort(reverse = True, \n",
    "                         key=lambda x: int(x.split('_')[1].split('.pth')[0]))\n",
    "    \n",
    "    w_check = widgets.Dropdown(options=list_checkpoint,\n",
    "                               description='Select checkpoint',\n",
    "                               disabled=False,\n",
    "                               style=dict(description_width='initial'))\n",
    "    display(widgets.HBox([w_run, w_check]))\n",
    "\n",
    "    \n",
    "on_value_change_check({'new':w_run.value})\n",
    "w_run.observe(on_value_change_check, names='value')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RUN_ID = w_run.value\n",
    "CHECKPOINT = w_check.value\n",
    "\n",
    "# Load parameters\n",
    "with open(os.path.join(root_dir, 'hparams.yaml')) as cfg_file:\n",
    "    opt = Munch(yaml.safe_load(cfg_file))\n",
    "\n",
    "opt.no_flip = True\n",
    "# Load parameters to the model, load the checkpoint\n",
    "model = create_model(opt)\n",
    "model = model.load_from_checkpoint((os.path.join(root_dir, 'checkpoints/', CHECKPOINT)))\n",
    "# Transfer the model to the GPU\n",
    "model.to('cuda');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the validation dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Switch train_only settings into validation\n",
    "opt.dataroot = opt.dataroot.replace('train/', 'val/sunny/Day/')\n",
    "# Load paths of all files contain in /Day\n",
    "p = pathlib.Path(opt.dataroot)\n",
    "dataset_paths = [str(x).replace(opt.dataroot,'') for x in p.iterdir()]\n",
    "dataset_paths.sort()\n",
    "\n",
    "sequence_name = {}\n",
    "# Make a dict with each sequence name as a key and\n",
    "# a list of paths to the images of the sequence as a value\n",
    "for file in dataset_paths:\n",
    "    # Keep only the sequence part contained in the name of the image\n",
    "    strip_name = file.split('_')[0]\n",
    "    \n",
    "    if strip_name not in sequence_name:\n",
    "        sequence_name[strip_name] = [file]\n",
    "    else:\n",
    "        sequence_name[strip_name].append(file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Select a sequence, an image and the moment of the day\n",
    "Select the sequence on which you want to work before choosing which image should be used.<br>\n",
    "Select the moment of the day, by choosing the angle of the sun $\\phi$ between [0,2$\\pi$],\n",
    "which maps to a sun elevation ∈ [+30◦,−40◦]\n",
    "<ul>\n",
    "<li>0 means day</li>\n",
    "<li>$\\pi$/2 means dusk</li>\n",
    "<li>$\\pi$ means night</li>\n",
    "<li>$\\pi$ + $\\pi$/2 means dawn</li>\n",
    "<li>2$\\pi$ means day (again)</li>\n",
    "</ul>"
   ]
  },
  {

    
          
            
          
    
    
  
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def drop_list():\n",
    "    # Select the sequence on which you want to make your test\n",
    "    return widgets.Dropdown(options=sequence_name.keys(),\n",
    "                         description='Select sequence',\n",
    "                         disabled=False,\n",
    "                         style=dict(description_width='initial'))\n",
    "def slider_img():\n",
    "    # Select an image from the sequence\n",
    "    return widgets.IntSlider(min=0, max=len(sequence_name[w_seq.value]) - 1,\n",
    "                           description='Select image')\n",
    "def slider_time():\n",
    "    # Select time\n",
    "    return widgets.FloatSlider(value=0, min=0, max=pi*2, step=0.01,\n",
    "                             description='Select time',\n",
    "                             readout_format='.2f')\n",
    "\n",
    "def debounce(wait):\n",
    "    # Decorator that will debounce the call to a function\n",
    "    def decorator(fn):\n",
    "        timer = None\n",
    "        def debounced(*args, **kwargs):\n",
    "            nonlocal timer\n",
    "            def call_it():\n",
    "                fn(*args, **kwargs)\n",
    "            if timer is not None:\n",
    "                timer.cancel()\n",
    "            timer = Timer(wait, call_it)\n",
    "            timer.start()\n",
    "        return debounced\n",
    "    return decorator\n",
    "    \n",
    "def inference(seq, index_img, phi, output = True):\n",
    "    global sequence_name, w_img_time, w_seq, opt, out\n",
    "    # Load the image\n",
    "    A_path = os.path.join(opt.dataroot, sequence_name[seq.value][index_img])\n",
    "    A_img = Image.open(A_path).convert('RGB')\n",
    "    # Apply image transformation\n",
    "    A = get_transform(opt, convert=False)(A_img)\n",
    "    # Normalization between -1 and 1\n",
    "    img_real = (((ToTensor()(A)) * 2) - 1).unsqueeze(0)\n",
    "    # Forward our image into the model with the specified ɸ\n",
    "    img_fake = model.forward(img_real.cuda(0), phi.cuda(0)) \n",
    "    # Encapsulate the initial image beside our result\n",
    "    new_im = Image.new('RGB', (A_img.width * 2, A_img.height))\n",
    "    new_im.paste(A_img, (0, 0))\n",
    "    new_im.paste(ToPILImage()((img_fake[0].cpu() + 1) / 2), (A_img.width, 0))\n",
    "    # Clear the output and display the widgets and the images\n",
    "    if output:\n",
    "        out.clear_output(wait=True)\n",
    "        with out:\n",
    "            # Resize the output\n",
    "            O_img = new_im.resize((new_im.width // 2, new_im.height // 2))\n",
    "            display(w_img_time, O_img)\n",
    "        display(out)\n",
    "    \n",
    "    return new_im\n",
    "\n",
    "@debounce(0.2)\n",
    "def on_value_change_img(change):\n",
    "    global w_seq, w_time\n",
    "    inference(w_seq, change['new'], torch.tensor(w_time.value))\n",
    "    \n",
    "@debounce(0.2)\n",
    "def on_value_change_time(change):\n",
    "    global w_seq, w_img\n",
    "    inference(w_seq, w_img.value, torch.tensor(change['new']))\n",
    "    \n",
    "def on_value_change_seq(change):\n",
    "    global w_seq, w_img, w_time\n",
    "    w_img = slider_img()\n",
    "    w_time = slider_time()\n",
    "    inference(w_seq, w_img.value, torch.tensor(w_time.value))\n",
    "    \n",
    "w_seq = drop_list()\n",
    "w_img = slider_img()\n",
    "w_time = slider_time()\n",
    "w_img_time = widgets.VBox([w_seq, widgets.HBox([w_img, w_time])])\n",
    "# Set the size of the output cell\n",
    "out = widgets.Output(layout=widgets.Layout(width='auto', height='240px'))\n",
    "\n",
    "inference(w_seq, w_img.value, torch.tensor(w_time.value))\n",
    "w_img.observe(on_value_change_img, names='value')\n",
    "w_time.observe(on_value_change_time, names='value')\n",
    "w_seq.observe(on_value_change_seq, names='value')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sequence to video\n",
    "The code below translates a sequence of images into a video.<br>\n",
    "By default, the 'Select time' slider is on 'Variable phi' ($\\phi$), in this case the time parameter will progressively increase from 0 to 2$\\pi$.<br>\n",
    "2$\\pi$ is reached at the end of the video. If you move the slider, you can select a fixed $\\phi$.<br>\n",
    "Require to install two more packages <code>pip install imageio imageio-ffmpeg<code>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Video\n",
    "import numpy as np\n",
    "import imageio\n",
    "\n",
    "w_button = widgets.Button(description='Start', tooltip='Start the inference of a sequence',\n",
    "                          icon='play')\n",
    "\n",
    "phi_opt = ['Variable phi'] + [str(round(f, 2)) for f in np.arange(0, pi*2, 0.01)]\n",
    "w_vid_time = widgets.SelectionSlider(options=phi_opt, value=phi_opt[0], description='Select time : ')\n",
    "\n",
    "w_vid_seq = drop_list()\n",
    "\n",
    "w_button_seq = widgets.VBox([widgets.HBox([w_vid_seq, w_vid_time]), w_button])\n",
    "display(w_button_seq)\n",
    "\n",
    "def get_video(bt):\n",
    "    global sequence_name, w_vid_seq, w_vid_time, w_button_seq\n",
    "    \n",
    "    clear_output(wait=True)\n",
    "    display(w_button_seq)\n",
    "    seq_size = len(sequence_name[w_vid_seq.value])\n",
    "    # Display progress bar\n",
    "    w_prog = widgets.IntProgress(value=0, min=0, max=seq_size, description='Loading:')\n",
    "    display(w_prog)\n",
    "    # Create a videos directory to save our video\n",
    "    save_name = str(pathlib.Path('.').absolute()) + '/videos/'\n",
    "    os.makedirs(save_name, exist_ok=True)\n",
    "    # If variable phi\n",
    "    if w_vid_time.value == 'Variable phi':\n",
    "        # Write our video in the project folder\n",
    "        save_name += 'comogan_{}_phi_{}.mp4'.format(w_vid_seq.value.replace('segment-', 'seg_'),\n",
    "                                                   'variable')\n",
    "    else:\n",
    "        save_name += 'comogan_{}_phi_{}.mp4'.format(w_vid_seq.value.replace('segment-', 'seg_'),\n",
    "                                                   w_vid_time.value.replace('.', '_'))\n",
    "    writer = imageio.get_writer(save_name, fps=10)\n",
    "    # Loop on the images contained in the sequence\n",
    "    for i in range(seq_size):\n",
    "        if w_vid_time.value == 'Variable phi':\n",
    "            # Inference of the i image in sequence_name[key]\n",
    "            phi_var = 2*pi/seq_size * i\n",
    "            my_img = inference(w_vid_seq, i, torch.tensor(phi_var), False)\n",
    "        else:\n",
    "            my_img = inference(w_vid_seq, i, torch.tensor(float(w_vid_time.value)), False)\n",
    "        writer.append_data(np.array(my_img))\n",
    "        # Progress bar\n",
    "        w_prog.value += 1\n",
    "    \n",
    "    writer.close()\n",
    "    display(Video(save_name, embed=True))\n",
    "\n",
    "w_button.on_click(get_video)"
   ]
  }
 ],
 "metadata": {
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}