notebooks/votenet_debug.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import os\n",
    "import sys\n",
    "import panel as pn\n",
    "import numpy as np\n",
    "import pyvista as pv\n",
    "pv.set_plot_theme(\"document\")\n",
    "import glob\n",
    "from matplotlib.colors import ListedColormap\n",
    "from omegaconf import OmegaConf\n",
    "from torch_geometric.data import Data\n",
    "import random\n",
    "import torch\n",
    "\n",
    "pn.extension('vtk')\n",
    "os.system('/usr/bin/Xvfb :99 -screen 0 1024x768x24 &')\n",
    "os.environ['DISPLAY'] = ':99'\n",
    "os.environ['PYVISTA_OFF_SCREEN'] = 'True'\n",
    "os.environ['PYVISTA_USE_PANEL'] = 'True'\n",
    "\n",
    "DIR = os.path.dirname(os.getcwd())\n",
    "sys.path.append(DIR)\n",
    "from torch_points3d.datasets.segmentation.scannet import SCANNET_COLOR_MAP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "VIZ_REL_DIR = \"outputs/2020-07-24/10-55-05/eval/2020-07-28_14-59-26/viz\"\n",
    "VIZ_DIR = os.path.join(DIR, VIZ_REL_DIR)\n",
    "all_viz_data = glob.glob(os.path.join(VIZ_DIR,'*.pt'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_idx = 2\n",
    "sample_idx = 0\n",
    "data = torch.load(all_viz_data[file_idx])\n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data.pred_boxes[0][0].classname"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = torch.load(all_viz_data[file_idx])\n",
    "confidence_threshold = 0.7\n",
    "def load_random_sample(event):\n",
    "    i = np.random.randint(0, data.pos.shape[0])\n",
    "     \n",
    "    pl1 = pv.Plotter(notebook=True)\n",
    "    pl2 = pv.Plotter(notebook=True)\n",
    "    \n",
    "    # Color by points with a label\n",
    "    pl1.add_points(data.pos[i].numpy(), color=\"gray\", opacity=0.4) \n",
    "    pl2.add_points(data.pos[i].numpy(), color=\"gray\", opacity=0.4) \n",
    "    \n",
    "    gt_boxes = data.gt_boxes[i]\n",
    "    for k in range(len(gt_boxes)):\n",
    "        xmin, ymin, zmin = gt_boxes[k].corners3d.min(0)\n",
    "        xmax, ymax, zmax = gt_boxes[k].corners3d.max(0)\n",
    "        box = pv.Box([xmin,xmax,ymin,ymax,zmin,zmax])\n",
    "        color = np.asarray(SCANNET_COLOR_MAP[gt_boxes[k].classname]) / 255.\n",
    "        pl1.add_mesh(box, color=color, show_edges=True, opacity=0.5)\n",
    "        \n",
    "    pred_boxes = data.pred_boxes[i]\n",
    "    for k in range(len(pred_boxes)):\n",
    "        if pred_boxes[k].score < confidence_threshold: continue\n",
    "        xmin, ymin, zmin = pred_boxes[k].corners3d.min(0)\n",
    "        xmax, ymax, zmax = pred_boxes[k].corners3d.max(0)\n",
    "        box = pv.Box([xmin,xmax,ymin,ymax,zmin,zmax])\n",
    "        color = np.asarray(SCANNET_COLOR_MAP[pred_boxes[k].classname]) / 255.\n",
    "        pl2.add_mesh(box, color=color, show_edges=True, opacity=0.5)\n",
    "\n",
    "    pan1.object = pl1.ren_win\n",
    "    pan2.object = pl2.ren_win"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pl1 = pv.Plotter(notebook=True)\n",
    "pl2 = pv.Plotter(notebook=True)\n",
    "pan1 = pn.panel(pl1.ren_win, sizing_mode='scale_both', aspect_ratio=1,orientation_widget=False,)\n",
    "pan2 = pn.panel(pl2.ren_win, sizing_mode='scale_both', aspect_ratio=1,orientation_widget=False,)\n",
    "button = pn.widgets.Button(name='Load new model', button_type='primary')\n",
    "button.on_click(load_random_sample)\n",
    "pn.Row(\n",
    "    pn.Column('## Votenet visualizer',button,'Threshold = 0.7', width=200),\n",
    "    pn.Column(pan1,'Ground truth'),pn.Column(pan2, 'Box Predictions'),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}