tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Searching Efficient 3D Architectures with Sparse Point-Voxel Convolution"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial, we will introduce how to efficiently segment LiDAR point clouds with our pre-trained SPVNAS."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first clone the codebase:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone https://github.com/mit-han-lab/spvnas.git\n",
    "\n",
    "import os\n",
    "import sys\n",
    "sys.path.append(os.path.join(os.getcwd(), 'spvnas'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's then install some libraries. This step might take around 5 minutes on Google Colab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!sudo apt-get install libsparsehash-dev 1>/dev/null\n",
    "!pip install --upgrade plotly 1>/dev/null\n",
    "!pip install --upgrade torchpack 1>/dev/null\n",
    "!pip install --upgrade git+https://github.com/mit-han-lab/torchsparse.git 1>/dev/null"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's import some libraries and define constants for visualization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "from torchsparse import SparseTensor\n",
    "from torchsparse.utils.quantize import sparse_quantize\n",
    "from torchsparse.utils.collate import sparse_collate\n",
    "\n",
    "COLOR_MAP = np.array(['#f59664', '#f5e664', '#963c1e', '#b41e50', '#ff0000',\n",
    "                      '#1e1eff', '#c828ff', '#5a1e96', '#ff00ff', '#ff96ff',\n",
    "                      '#4b004b', '#4b00af', '#00c8ff', '#3278ff', '#00af00',\n",
    "                      '#003c87', '#50f096', '#96f0ff', '#0000ff', '#ffffff'])\n",
    "\n",
    "LABEL_MAP = np.array([19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 0, 1, 19,\n",
    "                      19, 19, 2, 19, 19, 3, 19, 4, 19, 19, 19, 19, 19,\n",
    "                      19, 19, 19, 19, 5, 6, 7, 19, 19, 19, 19, 19, 19,\n",
    "                      19, 8, 19, 19, 19, 9, 19, 19, 19, 10, 11, 12, 13,\n",
    "                      19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19,\n",
    "                      19, 19, 19, 19, 19, 14, 15, 16, 19, 19, 19, 19, 19,\n",
    "                      19, 19, 17, 18, 19, 19, 19, 19, 19, 19, 19, 19, 19,\n",
    "                      19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19,\n",
    "                      19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19,\n",
    "                      19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19,\n",
    "                      19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19,\n",
    "                      19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19,\n",
    "                      19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19,\n",
    "                      19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19,\n",
    "                      19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19,\n",
    "                      19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19,\n",
    "                      19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19,\n",
    "                      19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19,\n",
    "                      19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19,\n",
    "                      19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's load some real lidar data and pre-process it for inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the LiDAR data and its label\n",
    "lidar = np.fromfile('spvnas/assets/000000.bin', dtype=np.float32)\n",
    "label = np.fromfile('spvnas/assets/000000.label', dtype=np.int32)\n",
    "lidar = lidar.reshape(-1, 4)\n",
    "label = LABEL_MAP[label & 0xFFFF]\n",
    "\n",
    "# Filter out ignored points\n",
    "lidar = lidar[label != 19]\n",
    "label = label[label != 19]\n",
    "\n",
    "# Quantize coordinates\n",
    "coords = np.round(lidar[:, :3] / 0.05)\n",
    "coords -= coords.min(0, keepdims=1)\n",
    "feats = lidar\n",
    "\n",
    "# Filter out duplicate points\n",
    "coords, indices, inverse = sparse_quantize(coords, return_index=True, return_inverse=True)\n",
    "coords = torch.tensor(coords, dtype=torch.int)\n",
    "feats = torch.tensor(feats[indices], dtype=torch.float)\n",
    "\n",
    "inputs = SparseTensor(coords=coords, feats=feats)\n",
    "inputs = sparse_collate([inputs]).cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we import the pretrained SPVNAS from our model zoo to run the inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from model_zoo import spvnas_specialized\n",
    "\n",
    "# Load the pre-trained model from model zoo\n",
    "model = spvnas_specialized('SemanticKITTI_val_SPVNAS@65GMACs').cuda()\n",
    "model.eval()\n",
    "\n",
    "# Run the inference\n",
    "outputs = model(inputs)\n",
    "outputs = outputs.argmax(1).cpu().numpy()\n",
    "\n",
    "# Map the prediction back to original points\n",
    "outputs = outputs[inverse]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we visualize the predictions from SPVNAS in an interactive window. Enjoy!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "def configure_plotly_browser_state():\n",
    "    import IPython\n",
    "    display(IPython.core.display.HTML('''\n",
    "        <script src=\"/static/components/requirejs/require.js\"></script>\n",
    "        <script>\n",
    "          requirejs.config({\n",
    "            paths: {\n",
    "              base: '/static/base',\n",
    "              plotly: 'https://cdn.plot.ly/plotly-latest.min.js?noext',\n",
    "            },\n",
    "          });\n",
    "        </script>\n",
    "        '''))\n",
    "\n",
    "import plotly\n",
    "import plotly.graph_objs as go\n",
    "\n",
    "trace = go.Scatter3d(\n",
    "    x=lidar[:, 0],\n",
    "    y=lidar[:, 1],\n",
    "    z=lidar[:, 2],\n",
    "    mode='markers',\n",
    "    marker={\n",
    "        'size': 1,\n",
    "        'opacity': 0.8,\n",
    "        'color': COLOR_MAP[outputs].tolist(),\n",
    "    }\n",
    ")\n",
    "\n",
    "configure_plotly_browser_state()\n",
    "plotly.offline.init_notebook_mode(connected=False)\n",
    "\n",
    "layout = go.Layout(\n",
    "    margin={'l': 0, 'r': 0, 'b': 0, 't': 0},\n",
    "    scene=dict(aspectmode='manual', aspectratio=dict(x=1, y=1, z=0.2))\n",
    ")\n",
    "\n",
    "plotly.offline.iplot(go.Figure(data=[trace], layout=layout))"
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}