interpolation_visualization.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Decoder(\n",
       "  (lin0): Linear(in_features=19, out_features=512, bias=True)\n",
       "  (lin1): Linear(in_features=512, out_features=512, bias=True)\n",
       "  (lin2): Linear(in_features=512, out_features=512, bias=True)\n",
       "  (lin3): Linear(in_features=512, out_features=493, bias=True)\n",
       "  (lin4): Linear(in_features=512, out_features=512, bias=True)\n",
       "  (lin5): Linear(in_features=512, out_features=512, bias=True)\n",
       "  (lin6): Linear(in_features=512, out_features=512, bias=True)\n",
       "  (lin7): Linear(in_features=512, out_features=512, bias=True)\n",
       "  (lin8): Linear(in_features=512, out_features=1, bias=True)\n",
       "  (relu): ReLU()\n",
       "  (th): Tanh()\n",
       ")"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from ipywidgets import FloatSlider\n",
    "from IPython.display import display\n",
    "import torch\n",
    "import numpy as np\n",
    "from deep_sdf import data, utils, mesh\n",
    "import deep_sdf.workspace as ws\n",
    "import json \n",
    "import os\n",
    "import meshplot as mp\n",
    "import trimesh\n",
    "import sys\n",
    "\n",
    "#experiment_directory = \"examples/torus/Run_w_con_loss_one_variable_usual\"  # Path to the saved decoder model\n",
    "experiment_directory = \"examples/torus\"  # Path to the saved decoder model\n",
    "# Load the decoder\n",
    "specs_filename = os.path.join(experiment_directory, \"specs.json\")\n",
    "\n",
    "if not os.path.isfile(specs_filename):\n",
    "    raise Exception(\n",
    "        'The experiment directory does not include specifications file \"specs.json\"'\n",
    "    )\n",
    "\n",
    "specs = json.load(open(specs_filename))\n",
    "\n",
    "arch = __import__(\"networks.\" + specs[\"NetworkArch\"], fromlist=[\"Decoder\"])\n",
    "\n",
    "latent_size = specs[\"CodeLength\"]\n",
    "\n",
    "decoder = arch.Decoder(latent_size, **specs[\"NetworkSpecs\"])\n",
    "\n",
    "decoder = torch.nn.DataParallel(decoder)\n",
    "\n",
    "saved_model_state = torch.load(\n",
    "    os.path.join(\n",
    "        experiment_directory, ws.model_params_subdir, \"latest\" + \".pth\"\n",
    "    )\n",
    ")\n",
    "saved_model_epoch = saved_model_state[\"epoch\"]\n",
    "\n",
    "decoder.load_state_dict(saved_model_state[\"model_state_dict\"])\n",
    "\n",
    "decoder = decoder.module.cuda()\n",
    "decoder.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Meshplot left an annoying print statement in their code. Using this context manager to supress it...\n",
    "class HiddenPrints:\n",
    "    def __enter__(self):\n",
    "        self._original_stdout = sys.stdout\n",
    "        sys.stdout = open(os.devnull, 'w')\n",
    "\n",
    "    def __exit__(self, exc_type, exc_val, exc_tb):\n",
    "        sys.stdout.close()\n",
    "        sys.stdout = self._original_stdout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c9a0fa59cce747f3b7f26201fa255a9e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(FloatSlider(value=0.0, description='z[0]', max=1.0, min=-1.0, step=0.2), FloatSlider(val…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mThe Kernel crashed while executing code in the current cell or a previous cell. \n",
      "\u001b[1;31mPlease review the code in the cell(s) to identify a possible cause of the failure. \n",
      "\u001b[1;31mClick <a href='https://aka.ms/vscodeJupyterKernelCrash'>here</a> for more info. \n",
      "\u001b[1;31mView Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details."
     ]
    }
   ],
   "source": [
    "# Initial latent vectors\n",
    "z = torch.zeros(16)\n",
    "plot=None\n",
    "@mp.interact(**{f'z[{i}]': FloatSlider(min=-1, max=1, step=0.2, value=0) for i in range(16)})\n",
    "def show(**kwargs):\n",
    "    global plot\n",
    "    global z\n",
    "    z = torch.tensor([kwargs[f'z[{i}]'] for i in range(16)])\n",
    "    z = z.to('cuda')\n",
    "\n",
    "    # Generate the 3D shape using the decoder\n",
    "    mesh_filename = \"sdf_utils/interpolation\"\n",
    "    with torch.no_grad():\n",
    "        mesh.create_mesh(decoder, z, mesh_filename, N=256, max_batch=int(2 ** 18))\n",
    "    \n",
    "    # Load the mesh using trimesh\n",
    "    loaded_mesh = trimesh.load(\"sdf_utils/interpolation.ply\")\n",
    "\n",
    "    # Extract vertices and faces\n",
    "    verts = np.array(loaded_mesh.vertices)\n",
    "    faces = np.array(loaded_mesh.faces)\n",
    "\n",
    "    # Set the color to white for all vertices\n",
    "    white_color = np.array([1.0, 1.0, 1.0])  # RGB values for white \n",
    "\n",
    "    if plot is None:\n",
    "        plot = mp.plot(verts, faces, c = white_color, return_plot=True)\n",
    "    else:\n",
    "        with HiddenPrints():\n",
    "            plot.update_object(vertices=verts, faces=faces)\n",
    "        display(plot._renderer)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "inr_sdf",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}