{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "bf950162-2118-4887-84b5-a41adc3e0853",
      "metadata": {
        "id": "bf950162-2118-4887-84b5-a41adc3e0853"
      },
      "source": [
        "# Implicit shape interpolation.\n",
        "\n",
        "In this lab work we will focus on interpolating between shapes $\\mathcal{S}_0$ and $\\mathcal{S}_1$. We will make so using the Level Set Equation (see e.g. <a href=\"https://ieeexplore.ieee.org/abstract/document/938899\">link</a>).\n",
        "We now parameterize not only a single shape, but the whole path $F_\\theta(x, t)$ between two shapes, i.e. $F_{\\theta}(x, 0)$ is the SDF of the first shape $\\mathcal{S}_0$ and $F_\\theta(x, 1)$ is the SDF of the second shape $\\mathcal{S}_1$\n",
        "\n",
        "Given a temporal signed distance function $F_\\theta(x, t)$, the level set equation encodes the motion of the level set following a prescribed vector field $V$ :\n",
        "$$\n",
        "\\frac{\\partial F(x, t)}{dt} + \\langle \\nabla F, V \\rangle = 0\n",
        "$$\n",
        "with $F(x, 0)$ set to be the signed distance field of the initial shape $\\mathcal{S}_0$.\n",
        "It has been used with neural signed distance functions only in 2023 (<a href=\"https://openaccess.thecvf.com/content/ICCV2023/papers/Novello_Neural_Implicit_Surface_Evolution_ICCV_2023_paper.pdf\"> link </a>). This lab will focus on replicating the code of the paper."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "VxQGThKwzl1M",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VxQGThKwzl1M",
        "outputId": "6ef89f4f-de61-46a0-82e2-3f767a1a1ee9"
      },
      "outputs": [],
      "source": [
        "!wget https://www.lix.polytechnique.fr/~pierson/cours/TP6.zip"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "md_Bc2mv0LOG",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "md_Bc2mv0LOG",
        "outputId": "16c0f74e-2fc3-4761-cf7e-bba2f5ba3a5c"
      },
      "outputs": [],
      "source": [
        "!unzip -o TP6.zip\n",
        "!ls"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "id": "11ba5897-c43c-4487-bade-e31c794a2c03",
      "metadata": {
        "id": "11ba5897-c43c-4487-bade-e31c794a2c03"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import time\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from matplotlib import pyplot as plt\n",
        "from slices import display_multi_slices, get_sdf_image\n",
        "from tqdm.auto import tqdm\n",
        "from synthetic_data import sphere_data, square_data\n",
        "from marching_squares import get_img_marching_squares, display_multiple_shapes_sdf\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8cca9b23-4b03-4e88-ab02-7ad8c32a8276",
      "metadata": {
        "id": "8cca9b23-4b03-4e88-ab02-7ad8c32a8276"
      },
      "source": [
        "## SIREN Neural network\n",
        "\n",
        "The networks are built with <a href=\"https://arxiv.org/pdf/2006.09661\">SIREN</a>. The network activation are sinusoidal functions, and a Siren layer is a network of the form\n",
        "$$\n",
        "\\phi(x) = \\sin(\\omega_0*(W x + b))\n",
        "$$\n",
        "\n",
        "with $\\omega_0 = 30$.\n",
        "\n",
        "The network is initialized such that the first layer weight follow a $\\mathcal{U}([-1/d_{in}, 1/d_{in}])$ (bias are initialized same), with $d_{in}$ the input feature, and the other follow a law $\\mathcal{U}([-\\frac{\\sqrt{6/d_{in}}}{\\omega_0}, \\frac{\\sqrt{6/d_{in}}}{\\omega_0}])$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "id": "7b08d16e-16d8-44e2-8923-74aabee2eafc",
      "metadata": {
        "id": "7b08d16e-16d8-44e2-8923-74aabee2eafc"
      },
      "outputs": [],
      "source": [
        "class SirenLayer(nn.Module):\n",
        "    def __init__(self, in_features, out_features, bias=True,\n",
        "                 is_first=False, omega_0=30):\n",
        "        super().__init__()\n",
        "        self.omega_0 = omega_0\n",
        "        self.is_first = is_first\n",
        "\n",
        "        self.in_features = in_features\n",
        "        ## Create the layer, and initialize it. You can do it in init_weights\n",
        "        \n",
        "        self.init_weights()\n",
        "\n",
        "    def init_weights(self):\n",
        "        with torch.no_grad():\n",
        "            if self.is_first:\n",
        "                ## Initialization of first layer type\n",
        "            else:\n",
        "                ## Other initialization\n",
        "\n",
        "    def forward(self, input):\n",
        "        ## Logic\n",
        "        return \n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c51c83f7-f900-49ad-b7d1-9304f7462005",
      "metadata": {
        "id": "c51c83f7-f900-49ad-b7d1-9304f7462005"
      },
      "source": [
        "### Neural network logic\n",
        "A siren network is of the form :\n",
        "$$\n",
        "\\Phi(x) = W * (\\phi_0  \\circ \\phi_1 \\circ \\phi_2 ....)(x) + b\n",
        "$$\n",
        "\n",
        "Build the network. The SIREN layers have constant output dimension hidden_dim. Optionnally, you can have skip connections on certain layers (you can ignore it for the beginning). Note that here, $x$ is in reality $x, t$ where $x$ is the coordinates in space and $t$ the coordinates in time."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "id": "282f9d2f-267b-4c7e-b5e2-734d95642c6d",
      "metadata": {
        "id": "282f9d2f-267b-4c7e-b5e2-734d95642c6d"
      },
      "outputs": [],
      "source": [
        "class SirenNet(torch.nn.Module):\n",
        "    def __init__(self, dim_in, dim_hidden, dim_out, num_layers, skip = [], omega_0 = 30.):\n",
        "        super().__init__()\n",
        "        self.num_layers = num_layers\n",
        "        self.dim_hidden = dim_hidden\n",
        "        self.skip = [i in skip for i in range(num_layers)]\n",
        "        self.omega_0 = omega_0\n",
        "\n",
        "        ## Create layer.\n",
        "        ## Last layer is a simple linear layer. Don't forget to intialize your weights as before!\n",
        "\n",
        "    def forward(self, x):\n",
        "        ## Network logic \n",
        "        ## You can ignore skip connections at the beginning\n",
        "        return "
      ]
    },
    {
      "cell_type": "markdown",
      "id": "142f5c59-6e14-41cd-98d4-ea0d0d6baede",
      "metadata": {
        "id": "142f5c59-6e14-41cd-98d4-ea0d0d6baede"
      },
      "source": [
        "# Loss functions\n",
        "\n",
        "A few loss functions are required to make the network work properly."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d1520f91-44bc-4a73-b53d-38b581c54e16",
      "metadata": {
        "id": "d1520f91-44bc-4a73-b53d-38b581c54e16"
      },
      "source": [
        "### Shape data attachement loss\n",
        "The data attachement loss contains two terms. First the network should be $0$ on the shape:\n",
        "\n",
        "$$\n",
        "\\mathcal{L}_{data} = \\sum_{x \\in \\mathcal{S}_0} ||F_\\theta(x, 0)||^2\n",
        "$$\n",
        "\n",
        "Second, the gradient of the SDF should be equal to the normal on the surface:\n",
        "\n",
        "$$\n",
        "\\mathcal{L}_{normal} = \\sum_{(x,n) \\in \\mathcal{S}_0} ||1 - \\langle n, \\nabla_\\theta F_\\theta(x, 0)||^2\n",
        "$$\n",
        "\n",
        "Write the code for $\\mathcal{L}_{data}$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "id": "8f411800-9c1d-482e-8694-d1df3f457f0f",
      "metadata": {
        "id": "8f411800-9c1d-482e-8694-d1df3f457f0f"
      },
      "outputs": [],
      "source": [
        "def loss_data(sdf_pc):\n",
        "    return "
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d619c0ad-c1e4-41b5-b310-eca1ba3f42a8",
      "metadata": {
        "id": "d619c0ad-c1e4-41b5-b310-eca1ba3f42a8"
      },
      "source": [
        "The computation of the gradient will be given. Write the code for $\\mathcal{L}_{normal}$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "id": "75f44af9-ca80-4aa5-802d-948e0b0e1c74",
      "metadata": {
        "id": "75f44af9-ca80-4aa5-802d-948e0b0e1c74"
      },
      "outputs": [],
      "source": [
        "def gradient(y, x, grad_outputs=None):\n",
        "    if grad_outputs is None:\n",
        "        grad_outputs = torch.ones_like(y)\n",
        "    grad = torch.autograd.grad(y, [x], grad_outputs=grad_outputs, create_graph=True)[0]\n",
        "    return grad\n",
        "\n",
        "def loss_normalign(u,v):\n",
        "    ##\n",
        "    return "
      ]
    },
    {
      "cell_type": "markdown",
      "id": "bfd78b08-321f-46d8-8fab-032266a0b4ea",
      "metadata": {
        "id": "bfd78b08-321f-46d8-8fab-032266a0b4ea"
      },
      "source": [
        "Now, given pc, normals as input, and a batch_size, sample batch size points and return $100*\\mathcal{L}_{data} + \\mathcal{L}_{normal}$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "id": "0344f266-0d73-49a3-beb6-386fcbd064ca",
      "metadata": {
        "id": "0344f266-0d73-49a3-beb6-386fcbd064ca"
      },
      "outputs": [],
      "source": [
        "def loss_shape_data(net, pc, normals, batch_size=2000, dim_space=2):\n",
        "    #sample batch size points and the corresponding normals\n",
        "    sample_pc = #\n",
        "    sample_pc.requires_grad = True\n",
        "    sample_nc = # has to be sampled with same indices!!!\n",
        "\n",
        "    sdf_pc = net(sample_pc)\n",
        "    #spatial gradients\n",
        "    grad_pc = gradient(sdf_pc, sample_pc)[:,0:2]\n",
        "\n",
        "    ## compute loss\n",
        "    loss_pc = 100*loss_data(sdf_pc) + loss_normalign(grad_pc, sample_nc)\n",
        "    return loss_pc"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7fca4087-1343-47b5-aa78-596d66c31a6b",
      "metadata": {
        "id": "7fca4087-1343-47b5-aa78-596d66c31a6b"
      },
      "source": [
        "### SDF data attachement loss\n",
        "Like in deepSDF, we learn the SDF based on the ground truth SDF in ambient space (we suppose we have a way to compute it).\n",
        "\n",
        "The loss is given as\n",
        "$$\n",
        "\\mathcal{L}_{amb} =  \\mathbb{E}_{x \\sim X} [||F_\\theta(x, 0) - \\text{SDF}_{\\text{gt}}(x)||²]\n",
        "$$\n",
        "where $X$ is the ambient space.\n",
        "\n",
        "In practice, we have already sample points before as in previous lab. Code the loss function that randomly selects them and compute the loss. Return $10*\\mathcal{L}_{amb}$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "id": "45a7b170-21ba-485a-9429-9f93a2f46f43",
      "metadata": {
        "id": "45a7b170-21ba-485a-9429-9f93a2f46f43"
      },
      "outputs": [],
      "source": [
        "def loss_amb(net, pc_hint, gt_sdf, batch_size=2000):\n",
        "    ## Sample random points in pc_hint and compute the SDF loss (as in previous lab but no clamp)\n",
        "    \n",
        "    return 10*#loss"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b57abd89-43a2-452d-83d6-7d243b42fb13",
      "metadata": {
        "id": "b57abd89-43a2-452d-83d6-7d243b42fb13"
      },
      "source": [
        "### Eikonal loss\n",
        "\n",
        "The final loss is based on the fact that the SDF respects the eikonal equation (the gradient of the SDF is always one):\n",
        "$$\n",
        "||\\nabla_x \\text{SDF}_x || = 1\n",
        "$$\n",
        "The loss is given by:\n",
        "$$\n",
        "\\mathbb{E}_{x, t} [|1-F(x, t)|]\n",
        "$$\n",
        "This expectation is approximated using a Monte Carlo estimation, at each step, points and time steps are sampled in the ambient space and we evaluate the sum at these points."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "id": "db4c6f20-c96c-4c87-bfab-97cf88e8e740",
      "metadata": {
        "id": "db4c6f20-c96c-4c87-bfab-97cf88e8e740"
      },
      "outputs": [],
      "source": [
        "def loss_eikonal_pts(grad_pts):\n",
        "    ## Get the eikonal loss given spatial gradients\n",
        "    return\n",
        "\n",
        "def loss_eikonal(net, batch_size, dim_space=2):\n",
        "    ## Sample random points in space ([-1, 1]^2) and in time ([0, 1]), and compute eikonal loss\n",
        "    pts_random = #\n",
        "    pts_random.requires_grad = True\n",
        "\n",
        "    sdf_random = #net(pts_random)\n",
        "\n",
        "    grad_tot_random = gradient(sdf_random, pts_random)\n",
        "    grad_spatial = grad_tot_random[:,0:2]\n",
        "    return loss_eikonal_pts(grad_spatial)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8156f23c-52d0-4441-8dd9-15f381a2198f",
      "metadata": {
        "id": "8156f23c-52d0-4441-8dd9-15f381a2198f"
      },
      "source": [
        "### Levet set loss\n",
        "\n",
        "In the case of a constant vector field, the equation reads as:\n",
        "\n",
        "$$\n",
        "\\mathcal{L}_{lse} = \\mathbb{E}_{(x, t)} [|| \\frac{\\partial F_\\theta(x, t)}{\\partial t} + \\langle \\nabla F_\\theta(x, t), V \\rangle||^2]\n",
        "$$\n",
        "\n",
        "This expectation is approximated using a Monte Carlo estimation, at each step, points and time steps are sampled in the ambient space and we evaluate the sum at these points."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "id": "e2d5b9e5-a7ec-40ca-8557-a771898348b3",
      "metadata": {
        "id": "e2d5b9e5-a7ec-40ca-8557-a771898348b3"
      },
      "outputs": [],
      "source": [
        "def loss_lse_eq(net, vf, batch_size, dim_space=2):\n",
        "    ## Sample random points in space ([-1, 1]^2) and in time ([0, 1]), and compute lse loss\n",
        "    pts_random = #\n",
        "    pts_random.requires_grad = True\n",
        "\n",
        "    sdf_random = net(pts_random)\n",
        "\n",
        "    grad_tot_random = gradient(sdf_random, pts_random)\n",
        "    #spatal gradient\n",
        "    grad_random = grad_tot_random[:,0:2]\n",
        "\n",
        "    #temporal gradient\n",
        "    tempgrad_random = grad_tot_random[:,2]\n",
        "\n",
        "    ## compute the loss\n",
        "    loss_lse = \n",
        "    return loss_lse"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "id": "05475c33-ebd8-4d4d-8a6d-f9452475e3ba",
      "metadata": {
        "id": "05475c33-ebd8-4d4d-8a6d-f9452475e3ba"
      },
      "outputs": [],
      "source": [
        "def evaluate_loss_cst_vf(net, pc, normals, hints_pc, gtsdf, vf,\n",
        "                           lpc, leik, lh, llse,\n",
        "                           lambda_pc = 1, lambda_eik = 2, lambda_hint = 1, lambda_lse = 1, batch_size = 2000):\n",
        "\n",
        "\n",
        "\n",
        "    # compute and store standard losses\n",
        "    loss_pc = loss_shape_data(net, pc, normals, batch_size)\n",
        "\n",
        "    loss_hint = loss_amb(net, hints_pc, gtsdf, batch_size)\n",
        "\n",
        "    loss_eik = loss_eikonal(net, batch_size)\n",
        "\n",
        "    loss_lse = loss_lse_eq(net, vf, batch_size)\n",
        "\n",
        "    # append all the losses\n",
        "    lpc.append(float(loss_pc))\n",
        "    leik.append(float(loss_eik))\n",
        "    lh.append(float(loss_hint))\n",
        "    llse.append(float(loss_lse))\n",
        "\n",
        "    # sum the losses of reach of this set of points\n",
        "    loss = lambda_pc*loss_pc + lambda_eik*loss_eik + lambda_hint*loss_hint + lambda_lse*loss_lse\n",
        "\n",
        "    return loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "id": "0e0efe94-1b71-42bc-9b8b-3296e87b87a4",
      "metadata": {
        "id": "0e0efe94-1b71-42bc-9b8b-3296e87b87a4"
      },
      "outputs": [],
      "source": [
        "def optimize_nise_vf(net, pc0, nc0, hints0, gtsdf0,\n",
        "                           vf, lpc, leik, lh, llse,\n",
        "                           lambda_pc = 1, lambda_eik = 2, lambda_hint = 1, lambda_lse = 2, batch_size = 2000, nepochs=100, plot_loss = True):\n",
        "\n",
        "    optim = torch.optim.Adam(params=net.parameters(), lr=2e-5)\n",
        "\n",
        "    tinit = time.time()\n",
        "    pbar = tqdm(total=nepochs,\n",
        "                desc=\"Training\")\n",
        "    for batch in range(nepochs):\n",
        "        optim.zero_grad()\n",
        "\n",
        "        loss = evaluate_loss_cst_vf(net, pc0, nc0, hints0, gtsdf0, vf, lpc, leik, lh, llse, lambda_pc, lambda_eik, lambda_hint, lambda_lse, batch_size)\n",
        "\n",
        "        loss.backward()\n",
        "        optim.step()\n",
        "        if batch % 100 == 99 or batch == 0:\n",
        "        #     print(f\"Epoch {epoch}/{nepochs} - loss : {loss.item()}\")\n",
        "            pbar.set_postfix({\"loss\": loss.item()})\n",
        "        pbar.update(1)\n",
        "\n",
        "    tend = time.time()\n",
        "\n",
        "    print(\"Optimizing NN took\", '{:.2f}'.format(tend-tinit),\"s.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "60280ae4-4c78-4bbb-a06a-8aeddc623a46",
      "metadata": {
        "id": "60280ae4-4c78-4bbb-a06a-8aeddc623a46"
      },
      "source": [
        "### Data\n",
        "We work for the moment on simple 2D canonical shapes for example a circle and a square, but feel free to test on your own shapes. You will need points sampled on the shape with their normals $(x_i, n_i)$ and point sampled in the ambient space with their groundtruth signed distances $(\\text{hint}_j, \\text{SDF}_{\\text{GT}}(\\text{hint}_j))$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e8b91bc9-6b50-44ff-b61e-0b7d9e91346e",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 250
        },
        "id": "e8b91bc9-6b50-44ff-b61e-0b7d9e91346e",
        "outputId": "df140348-934f-4fde-9d78-dfb226f58293"
      },
      "outputs": [],
      "source": [
        "pc0,nc0,pts_hints0,gt_sdf_hints0,gt_grad_hints0, coords, sdf_coords = sphere_data(10000,100000)\n",
        "\n",
        "#constant vector field (a single vector)\n",
        "vf = torch.tensor([0.3,0.], device=device)\n",
        "\n",
        "\n",
        "fig, axs =  plt.subplots(1, 4, figsize=(8, 15))\n",
        "pc_numpy = pc0.detach().cpu().numpy()\n",
        "axs[0].scatter(pc_numpy[::100, 0], pc_numpy[::100, 1])\n",
        "axs[0].set_aspect(1.0/axs[0].get_data_ratio(), adjustable='box')\n",
        "axs[0].set_xlim([-1, 1])\n",
        "axs[0].set_ylim([-1, 1])\n",
        "axs[0].set_title(\"Sphere points\")\n",
        "sdf_numpy = sdf_coords.detach().cpu().numpy()\n",
        "axs[1].imshow(get_sdf_image(sdf_numpy))\n",
        "axs[1].contour(sdf_numpy, 10, colors='k', linewidths=.4, linestyles='solid')\n",
        "axs[1].set_xticks([])\n",
        "axs[1].set_yticks([])\n",
        "axs[1].set_title(\"Ground truth SDF\")\n",
        "rec_img = get_img_marching_squares(sdf_numpy)\n",
        "axs[2].imshow(rec_img)\n",
        "axs[2].set_xticks([])\n",
        "axs[2].set_yticks([])\n",
        "axs[2].set_title(\"Reconstructed shape\")\n",
        "\n",
        "def get_vector_field_meshgrid(X_value, Y_value):\n",
        "    # Create a grid of points\n",
        "    x = np.linspace(-1, 1, 20)\n",
        "    y = np.linspace(-1, 1, 20)\n",
        "    X, Y = np.meshgrid(x, y)\n",
        "\n",
        "    # Define the constant vector field (e.g., constant vector [1, 1] everywhere)\n",
        "    U = np.ones_like(X)*X_value/(10*max(X_value, Y_value))  # constant x-component\n",
        "    V = np.ones_like(Y)*Y_value/(10*max(X_value, Y_value))  # constant y-component\n",
        "    return X,Y,U,V\n",
        "\n",
        "X,Y,U,V = get_vector_field_meshgrid(vf[0].item(), vf[1].item())\n",
        "axs[3].quiver(X, Y, U, V, angles='xy', scale_units='xy', scale=1)\n",
        "axs[3].set_xlim((-1, 1))\n",
        "axs[3].set_ylim((-1, 1))\n",
        "axs[3].set_aspect(\"equal\")\n",
        "axs[3].set_xticks([])\n",
        "axs[3].set_yticks([])\n",
        "axs[3].set_title('Constant Vector Field')\n",
        "axs[3].grid()\n",
        "plt.tight_layout()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b5b9a7cd-7978-4917-9497-5733b8a05ccf",
      "metadata": {
        "id": "b5b9a7cd-7978-4917-9497-5733b8a05ccf"
      },
      "source": [
        "Check that the eikonal loss is correct"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "75259c4d-7871-45cb-9bb6-60f21e0dfc3f",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "75259c4d-7871-45cb-9bb6-60f21e0dfc3f",
        "outputId": "07cee8a1-35e5-44c1-ed76-079c2e48bdb2"
      },
      "outputs": [],
      "source": [
        "np.isclose(loss_eikonal_pts(gt_grad_hints0).item(), 0) ##Should be true"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d0543e26-c522-40b9-8241-27817ac5dcab",
      "metadata": {
        "id": "d0543e26-c522-40b9-8241-27817ac5dcab"
      },
      "source": [
        "Run the training and check the loss and results. We advice setting dim_hdden = 128, and a depth of 3 to start (faster). You can try improving the results later."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ba046f6d-23d7-4d10-b9a0-b3099cdd323f",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 208,
          "referenced_widgets": [
            "d4ca0db0b9d84682b24c7735edf63d73",
            "ec63e5747b5c42a8962dc1f7db379e10",
            "31992866c73e4b689a22cb3a80789f9a",
            "763ea50710c14251b69163387530aa56",
            "f166f9e6e45440218710de5a418e22d1",
            "15ad03f4bfe848e9bb51e0d4401a9b9e",
            "32cfa13bde2c4147ae307ac25609e819",
            "5158ba6dd7da4f8ba602aedf49c63bac",
            "a55e3546319a4cefbaee60142ac0d4e5",
            "da24c622a7254f9a960e71f62e48bd25",
            "a77c8c33fb2747aeb3466ddb3e05b6e2"
          ]
        },
        "id": "ba046f6d-23d7-4d10-b9a0-b3099cdd323f",
        "outputId": "27f4d744-1b00-4b32-8479-d57062543fad"
      },
      "outputs": [],
      "source": [
        "npl = 128\n",
        "depth = 3\n",
        "net = SirenNet(dim_in=3,dim_out=1,num_layers=depth,dim_hidden=npl,skip=[]).to(device)\n",
        "lpc, leik, lh, llse = [], [], [], []\n",
        "\n",
        "pc_t0 = torch.concat((pc0,torch.zeros((pc0.size(0),1), device = device)),dim=1)\n",
        "pts_hintst0 = torch.concat((pts_hints0,torch.zeros((pts_hints0.size(0),1), device = device)),dim=1)\n",
        "\n",
        "\n",
        "optimize_nise_vf(net, pc_t0, nc0, pts_hintst0, gt_sdf_hints0, vf,\n",
        "                       lpc, leik, lh, llse,nepochs=5000)\n",
        "\n",
        "torch.save(net, \"net_vf_{}_{}.net\".format(npl, depth))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f980ff55-b070-46e4-9a46-7bf3f8d204cc",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 388
        },
        "id": "f980ff55-b070-46e4-9a46-7bf3f8d204cc",
        "outputId": "e2b8f679-b069-42b2-cebf-a2770d795b0b"
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(6,4))\n",
        "plt.yscale('log')\n",
        "plt.plot(lpc, label = 'Point cloud loss ({:.2f})'.format(lpc[-1]))\n",
        "plt.plot(leik, label = 'Eikonal loss ({:.2f})'.format(leik[-1]))\n",
        "plt.plot(lh, label = 'Learning points loss ({:.2f})'.format(lh[-1]))\n",
        "plt.plot(llse, label = 'LSE loss ({:.2f})'.format(llse[-1]))\n",
        "plt.xlabel(\"Epochs\")\n",
        "plt.legend()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "503a8468-e7b9-4cbe-a6d8-499ff4d8a204",
      "metadata": {
        "id": "503a8468-e7b9-4cbe-a6d8-499ff4d8a204"
      },
      "source": [
        "## Results\n",
        "We display the SDF and the resulting reconstructed shapes. You can find a potential example in the img/ folder"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8810b38e-af78-4b1c-a0e6-19692930b561",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 422
        },
        "id": "8810b38e-af78-4b1c-a0e6-19692930b561",
        "outputId": "7168b3d4-a780-4e28-b513-79fa6febb24f"
      },
      "outputs": [],
      "source": [
        "display_multi_slices(net, resolution=200,figsize=(14, 5))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5d2e01d2-43bc-47ea-8f9c-a149579abe0b",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 422
        },
        "id": "5d2e01d2-43bc-47ea-8f9c-a149579abe0b",
        "outputId": "c9291fed-c063-48e4-b8ff-25a63b253d8e"
      },
      "outputs": [],
      "source": [
        "display_multiple_shapes_sdf(net, resolution=200, figsize=(14, 5))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3243326e-6373-49f6-bd59-0fb07d8d8954",
      "metadata": {
        "id": "3243326e-6373-49f6-bd59-0fb07d8d8954"
      },
      "source": [
        "## Morphing between shapes\n",
        "\n",
        "\n",
        "The second test case consists into interpolate between two shapes, with signed distance functions $f_O$ at time 0 and $f_1$ at time 1. The vector field is not known in that case, but Novello et al. propose a simple and elegant solution, based on the following equality:\n",
        "$$\n",
        "V(x,t) = -(f_1(x) - F_{\\theta}(x,t))\\frac{\\nabla F_{\\theta}(x,t)}{\\|\\nabla F_{\\theta}(x,t)\\|}\n",
        "$$\n",
        "with\n",
        "$$\n",
        "F_{\\theta}(x,0)=f_0(x)\n",
        "$$\n",
        "\n",
        " Then by substitution in the LSE we have:\n",
        "$$\n",
        "\\frac{\\partial F_{\\theta}(x,t)}{\\partial t} - (f_1(x) - F_{\\theta}(x,t))\\|\\nabla F_{\\theta}(x,t)\\| = 0\n",
        "$$\n",
        "and the loss:\n",
        "$$\n",
        "\\mathcal L_{LSE}(\\theta) = \\mathbb{E}_{(x,t)}\\Big[\\Big\\|\\frac{\\partial F_{\\theta}(x,t)}{\\partial t}- (f_1(x) - F_{\\theta}(x,t))\\|\\nabla F_{\\theta}(x,t)\\|\\Big\\|^2\\Big]\n",
        "$$\n",
        "Again the loss is evaluated using a Monte Carlo strategy"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "id": "0b297e85-bd61-4529-a2b5-23c20ce1f9ae",
      "metadata": {
        "id": "0b297e85-bd61-4529-a2b5-23c20ce1f9ae"
      },
      "outputs": [],
      "source": [
        "def loss_lse_morph(net, pc_hint, gt_sdf, batch_size):\n",
        "    ## Sample batch size points with position hintspc, and add random time in [0, 1]\n",
        "\n",
        "    random_hints = #\n",
        "    random_hints.requires_grad = True\n",
        "\n",
        "    ## Evaluate SDF and gradient. Spatial gradient (nabla) is accessible via [:, 0:2], temporal derivative via [:, 2]\n",
        "    sdf_random_hints = net(random_hints)\n",
        "    grad_tot_ = gradient(sdf_random_hints, random_hints)\n",
        "    spatial_grad = grad_tot_[:,0:2]\n",
        "    #temporal gradient\n",
        "    temp_deriv = grad_tot_[:,2]\n",
        "\n",
        "    #compute the loss\n",
        "    l_lse = \n",
        "    return l_lse"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "id": "a099822c-b01f-4839-8214-c234fbcf5001",
      "metadata": {
        "id": "a099822c-b01f-4839-8214-c234fbcf5001"
      },
      "outputs": [],
      "source": [
        "def evaluate_loss_morphing(net, pc0, normals0, hints_pc0, gtsdf0,\n",
        "                           pc1, normals1, hints_pc1, gtsdf1,\n",
        "                           lpc, leik, lh, llse,\n",
        "                           lambda_pc = 100, lambda_eik = 2e2, lambda_hint = 1e2, lambda_lse = 1e2, batch_size = 2000):\n",
        "\n",
        "\n",
        "    # compute and store standard losses\n",
        "    loss_pc = loss_shape_data(net, pc0, normals0, batch_size) + loss_shape_data(net, pc1, normals1, batch_size)\n",
        "    loss_hint = loss_amb(net, hints_pc0, gtsdf0, batch_size) + loss_amb(net, hints_pc1, gtsdf1, batch_size)\n",
        "\n",
        "    loss_eik = loss_eikonal(net, batch_size)\n",
        "\n",
        "    loss_lse = loss_lse_morph(net, hints_pc1, gtsdf1, batch_size)\n",
        "\n",
        "    # append all the losses\n",
        "    lpc.append(float(loss_pc))\n",
        "    leik.append(float(loss_eik))\n",
        "    lh.append(float(loss_hint))\n",
        "    llse.append(float(loss_lse))\n",
        "\n",
        "    # sum the losses of reach of this set of points\n",
        "    loss = lambda_pc*loss_pc + lambda_eik*loss_eik + lambda_hint*loss_hint + lambda_lse*loss_lse\n",
        "\n",
        "    return loss"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "51c2bbc6-6e06-4161-8f81-f7fdea1750df",
      "metadata": {
        "id": "51c2bbc6-6e06-4161-8f81-f7fdea1750df"
      },
      "source": [
        "### Interpolate between square and circle"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "id": "1d3b6fda-ceb3-45ea-b778-ada69956588e",
      "metadata": {
        "id": "1d3b6fda-ceb3-45ea-b778-ada69956588e"
      },
      "outputs": [],
      "source": [
        "#interpolate between a circle and a square\n",
        "pc0,nc0,pts_hints0,gt_sdf_hints0,gt_grad_hints0,_,sdf_coords0 = sphere_data(10000,100000)\n",
        "pc1,nc1,pts_hints1,gt_sdf_hints1,gt_grad_hints1,_,sdf_coords1 = square_data(10000,100000)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8ff70dc4-7426-415d-9c77-2a1067f962a7",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 265
        },
        "id": "8ff70dc4-7426-415d-9c77-2a1067f962a7",
        "outputId": "05c16f2c-5584-40a5-bd1b-93f07f04bd3c"
      },
      "outputs": [],
      "source": [
        "fig, axs =  plt.subplots(1, 3, figsize=(8, 10))\n",
        "pc_numpy = pc1.detach().cpu().numpy()\n",
        "axs[0].scatter(pc_numpy[::100, 0], pc_numpy[::100, 1])\n",
        "axs[0].set_aspect(1.0/axs[0].get_data_ratio(), adjustable='box')\n",
        "axs[0].set_xlim([-1, 1])\n",
        "axs[0].set_ylim([-1, 1])\n",
        "axs[0].set_title(\"Square points\")\n",
        "sdf_numpy = sdf_coords1.detach().cpu().numpy()\n",
        "axs[1].imshow(get_sdf_image(sdf_numpy))\n",
        "axs[1].contour(sdf_numpy, 10, colors='k', linewidths=.4, linestyles='solid')\n",
        "axs[1].set_xticks([])\n",
        "axs[1].set_yticks([])\n",
        "axs[1].set_title(\"Ground truth SDF\")\n",
        "rec_img = get_img_marching_squares(sdf_numpy)\n",
        "axs[2].imshow(rec_img)\n",
        "axs[2].set_xticks([])\n",
        "axs[2].set_yticks([])\n",
        "axs[2].set_title(\"Reconstructed shape\")\n",
        "a = 1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "id": "f7043a59-f5fb-4066-b37d-24443620ca8f",
      "metadata": {
        "id": "f7043a59-f5fb-4066-b37d-24443620ca8f"
      },
      "outputs": [],
      "source": [
        "def optimize_nise_morphing(net, pc0, nc0, pc1, nc1, hints0, gtsdf0,\n",
        "                           hints1, gtsdf1, lpc, leik, lh, llse,\n",
        "                           lambda_pc = 1, lambda_eik = 2, lambda_hint = 1, lambda_lse = 2, batch_size = 2000, nepochs=100, plot_loss = True):\n",
        "\n",
        "    optim = torch.optim.Adam(params=net.parameters(), lr=2e-5)\n",
        "    #optim = torch.optim.LBFGS(params=net.parameters())\n",
        "    tinit = time.time()\n",
        "    pbar = tqdm(total=nepochs,\n",
        "                desc=\"Training\")\n",
        "    for batch in range(nepochs):\n",
        "        optim.zero_grad()\n",
        "\n",
        "        loss = evaluate_loss_morphing(net, pc0, nc0, pc1, nc1, hints0, gtsdf0, hints1, gtsdf1, lpc, leik, lh, llse, lambda_pc, lambda_eik, lambda_hint, lambda_lse, batch_size)\n",
        "\n",
        "        loss.backward()\n",
        "        optim.step()\n",
        "\n",
        "        if batch % 100 == 99 or batch == 0:\n",
        "            pbar.set_postfix({\"loss\": loss.item()})\n",
        "        pbar.update(1)\n",
        "\n",
        "    tend = time.time()\n",
        "\n",
        "    print(\"Optimizing NN took\", '{:.2f}'.format(tend-tinit),\"s.\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "cc059598-e2c0-4045-b8ec-d279325d0a44",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 153,
          "referenced_widgets": [
            "061eafcbd47746deac63904c124c5882",
            "f8f23ec6ff72468596bc2a433d6807a1",
            "af6b8a57b73b4353aecf8c5d0fb74460",
            "eb44f40cd7b44c17a5a1540489cf6f22",
            "acff691c0e1048d1aae574000b15670b",
            "30f338dfedcb4ec3bd63b9d1ff6f68f7",
            "329739d78ca245aa9ced82958842d0bb",
            "7726ab68098049f78232a2df895658ea",
            "033d4c41b79a42b29ce395cf3aa2e664",
            "ee0524227a8341589fce110493410182",
            "7b76b328fe3c422e9bd69193eb122893"
          ]
        },
        "id": "cc059598-e2c0-4045-b8ec-d279325d0a44",
        "outputId": "8d863973-3d3d-4940-8cfa-07ff4bce3ef9"
      },
      "outputs": [],
      "source": [
        "npl = 128\n",
        "depth = 3\n",
        "net_morph = SirenNet(dim_in=3,dim_out=1,num_layers=depth,dim_hidden=npl,skip=[]).to(device)\n",
        "lpc, leik, lh, llse = [], [], [], []\n",
        "\n",
        "pc_t0 = torch.concat((pc0,torch.zeros((pc0.size(0),1), device = device)),dim=1)\n",
        "pc_t1 = torch.concat((pc1,torch.ones((pc1.size(0),1), device = device)),dim=1)\n",
        "pts_hintst0 = torch.concat((pts_hints0,torch.zeros((pts_hints0.size(0),1), device = device)),dim=1)\n",
        "pts_hintst1 = torch.concat((pts_hints1,torch.ones((pts_hints1.size(0),1), device = device)),dim=1)\n",
        "\n",
        "optimize_nise_morphing(net_morph, pc_t0, nc0, pts_hintst0, gt_sdf_hints0,\n",
        "                       pc_t1, nc1, pts_hintst1, gt_sdf_hints1, lpc, leik, lh, llse,nepochs=5000)\n",
        "\n",
        "torch.save(net, \"net_morph_{}_{}.net\".format(npl, depth))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "414db36f-a4cc-4d08-9682-596da420b010",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 405
        },
        "id": "414db36f-a4cc-4d08-9682-596da420b010",
        "outputId": "714e1bb2-31d0-4ae8-910e-ca26908972da"
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(6,4))\n",
        "plt.yscale('log')\n",
        "plt.plot(lpc, label = 'Point cloud loss ({:.2f})'.format(lpc[-1]))\n",
        "plt.plot(leik, label = 'Eikonal loss ({:.2f})'.format(leik[-1]))\n",
        "plt.plot(lh, label = 'Learning points loss ({:.2f})'.format(lh[-1]))\n",
        "plt.plot(llse, label = 'LSE loss ({:.2f})'.format(llse[-1]))\n",
        "plt.xlabel(\"Epochs\")\n",
        "plt.legend()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "389bcf2b-7e04-48f7-90ed-05227f1f0520",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 422
        },
        "id": "389bcf2b-7e04-48f7-90ed-05227f1f0520",
        "outputId": "53860a2d-9fef-4ec1-df74-609686757e52"
      },
      "outputs": [],
      "source": [
        "display_multi_slices(net_morph, resolution=200, figsize=(14, 5))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3112976b-684d-4c79-8155-7f9f3d227f36",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 422
        },
        "id": "3112976b-684d-4c79-8155-7f9f3d227f36",
        "outputId": "d293b397-cda5-4b03-dcaa-e9f85a35a57a"
      },
      "outputs": [],
      "source": [
        "display_multiple_shapes_sdf(net_morph, resolution=200, figsize=(14, 5))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e4795c08-4fd3-48be-a02d-bd775ef6339f",
      "metadata": {
        "id": "e4795c08-4fd3-48be-a02d-bd775ef6339f"
      },
      "source": [
        "# Bonus: Interpolation of real shapes\n",
        "\n",
        "Once the method is working (at least partially), try using real shapes. We suggest you to download an image from the following <a href=\" https://2dshapesstructure.github.io/\">website</a>.\n",
        "Then you can extract the SDF by first extracting distances <a href=\"https://docs.scipy.org/doc/scipy/reference/generated/scipy.ndimage.distance_transform_edt.html\">scipy.ndimage.distance_transform_edt</a>, and then find how to compute the <b>signed</b> distance field.\n",
        "\n",
        "You can also extract the contour using scikit-image <a href=\"https://scikit-image.org/docs/stable/api/skimage.measure.html#skimage.measure.find_contours\">measure.find_contours</a> and compute the SDF as in previous TD. Good luck!"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "cdc85100",
      "metadata": {},
      "source": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.8.18"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "033d4c41b79a42b29ce395cf3aa2e664": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "061eafcbd47746deac63904c124c5882": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_f8f23ec6ff72468596bc2a433d6807a1",
              "IPY_MODEL_af6b8a57b73b4353aecf8c5d0fb74460",
              "IPY_MODEL_eb44f40cd7b44c17a5a1540489cf6f22"
            ],
            "layout": "IPY_MODEL_acff691c0e1048d1aae574000b15670b"
          }
        },
        "15ad03f4bfe848e9bb51e0d4401a9b9e": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "30f338dfedcb4ec3bd63b9d1ff6f68f7": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "31992866c73e4b689a22cb3a80789f9a": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_5158ba6dd7da4f8ba602aedf49c63bac",
            "max": 5000,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_a55e3546319a4cefbaee60142ac0d4e5",
            "value": 5000
          }
        },
        "329739d78ca245aa9ced82958842d0bb": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "32cfa13bde2c4147ae307ac25609e819": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "5158ba6dd7da4f8ba602aedf49c63bac": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "763ea50710c14251b69163387530aa56": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_da24c622a7254f9a960e71f62e48bd25",
            "placeholder": "​",
            "style": "IPY_MODEL_a77c8c33fb2747aeb3466ddb3e05b6e2",
            "value": " 5000/5000 [01:04&lt;00:00, 98.19it/s, loss=0.00516]"
          }
        },
        "7726ab68098049f78232a2df895658ea": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "7b76b328fe3c422e9bd69193eb122893": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "a55e3546319a4cefbaee60142ac0d4e5": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "a77c8c33fb2747aeb3466ddb3e05b6e2": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "acff691c0e1048d1aae574000b15670b": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "af6b8a57b73b4353aecf8c5d0fb74460": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_7726ab68098049f78232a2df895658ea",
            "max": 5000,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_033d4c41b79a42b29ce395cf3aa2e664",
            "value": 5000
          }
        },
        "d4ca0db0b9d84682b24c7735edf63d73": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_ec63e5747b5c42a8962dc1f7db379e10",
              "IPY_MODEL_31992866c73e4b689a22cb3a80789f9a",
              "IPY_MODEL_763ea50710c14251b69163387530aa56"
            ],
            "layout": "IPY_MODEL_f166f9e6e45440218710de5a418e22d1"
          }
        },
        "da24c622a7254f9a960e71f62e48bd25": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "eb44f40cd7b44c17a5a1540489cf6f22": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_ee0524227a8341589fce110493410182",
            "placeholder": "​",
            "style": "IPY_MODEL_7b76b328fe3c422e9bd69193eb122893",
            "value": " 5000/5000 [01:17&lt;00:00, 72.64it/s, loss=0.0864]"
          }
        },
        "ec63e5747b5c42a8962dc1f7db379e10": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_15ad03f4bfe848e9bb51e0d4401a9b9e",
            "placeholder": "​",
            "style": "IPY_MODEL_32cfa13bde2c4147ae307ac25609e819",
            "value": "Training: 100%"
          }
        },
        "ee0524227a8341589fce110493410182": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "f166f9e6e45440218710de5a418e22d1": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "f8f23ec6ff72468596bc2a433d6807a1": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_30f338dfedcb4ec3bd63b9d1ff6f68f7",
            "placeholder": "​",
            "style": "IPY_MODEL_329739d78ca245aa9ced82958842d0bb",
            "value": "Training: 100%"
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}