notebooks/scannet_explorer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Scannet vizualisation"
   ]
  },
  {
   "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",
    "import glob\n",
    "from matplotlib.colors import ListedColormap\n",
    "from omegaconf import OmegaConf\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",
    "from torch_points3d.datasets.segmentation.scannet import ScannetDataset, Scannet\n",
    "from torch_points3d.datasets.segmentation import IGNORE_LABEL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Scannet dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DIR = os.path.dirname(os.getcwd())\n",
    "dataset_options = OmegaConf.load(os.path.join(DIR,'conf/data/segmentation/scannet.yaml'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_options.data.dataroot = os.path.join(DIR,\"data\")\n",
    "dataset = ScannetDataset(dataset_options.data)\n",
    "print(dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualise the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def buil_cmap():\n",
    "    mapping = np.linspace(IGNORE_LABEL, len(Scannet.VALID_CLASS_IDS)+1, 256)\n",
    "    newcolors = np.zeros((256, 3))\n",
    "    for i, raw_label in enumerate(Scannet.VALID_CLASS_IDS):\n",
    "        newcolors[mapping >= i-0.5] = np.asarray(Scannet.SCANNET_COLOR_MAP[raw_label]) / 255.\n",
    "    return ListedColormap(newcolors)\n",
    "cmap = buil_cmap()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_random_data(event):\n",
    "    i = np.random.randint(0, len(dataset.train_dataset))\n",
    "    sample = dataset.train_dataset[i]\n",
    "    xyz = sample.pos.numpy()\n",
    "    pl = pv.Plotter(notebook=True)\n",
    "    point_cloud = pv.PolyData(xyz)\n",
    "    point_cloud['label'] = sample.y.numpy()\n",
    "    pl.add_points(point_cloud,cmap=cmap, clim=[-1, len(Scannet.VALID_CLASS_IDS)+1]) \n",
    "    \n",
    "    pan.object = pl.ren_win\n",
    "    return point_cloud"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pl = pv.Plotter(notebook=True)\n",
    "pan = pn.panel(pl.ren_win, sizing_mode='scale_both', aspect_ratio=1,orientation_widget=True,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "button = pn.widgets.Button(name='Load new model', button_type='primary')\n",
    "button.on_click(load_random_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dashboard = pn.Row(\n",
    "    pn.Column('## Visualiser for forward runs',button),\n",
    "    pan\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dashboard"
   ]
  },
  {
   "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.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}