open3d_pointnet_inference.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "from random import randrange\n",
    "from IPython.display import clear_output\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.parallel\n",
    "import torch.optim as optim\n",
    "import torch.utils.data\n",
    "import torchvision.datasets as dset\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision.utils as vutils\n",
    "from torch.autograd import Variable\n",
    "from datasets import PartDataset\n",
    "from pointnet import PointNetCls\n",
    "import torch.nn.functional as F\n",
    "import matplotlib.pyplot as plt\n",
    "import open3d as o3\n",
    "import download\n",
    "from open3d import JVisualizer\n",
    "if torch.cuda.is_available():\n",
    "    import torch.backends.cudnn as cudnn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# General parameters\n",
    "NUM_POINTS = 10000\n",
    "MODEL_PATH = './cls_model.pth'\n",
    "DATA_FOLDER = './shapenetcore_partanno_segmentation_benchmark_v0'\n",
    "\n",
    "# download dataset and pre-trained model\n",
    "download.download_contents()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create dataset object\n",
    "test_dataset_seg = PartDataset(\n",
    "    root=DATA_FOLDER,\n",
    "    train=False,\n",
    "    classification=False,\n",
    "    npoints=NUM_POINTS)\n",
    "\n",
    "# Problem ontology\n",
    "classes_dict = {'Airplane': 0, 'Bag': 1, 'Cap': 2, 'Car': 3, 'Chair': 4, \n",
    "                'Earphone': 5, 'Guitar': 6, 'Knife': 7, 'Lamp': 8, 'Laptop': 9,\n",
    "                'Motorbike': 10, 'Mug': 11, 'Pistol': 12, 'Rocket': 13, \n",
    "                'Skateboard': 14, 'Table': 15}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the classification network from pre-trained model\n",
    "classifier = PointNetCls(k=len(classes_dict.items()), num_points=NUM_POINTS)\n",
    "if torch.cuda.is_available():\n",
    "    classifier.cuda()\n",
    "    classifier.load_state_dict(torch.load(MODEL_PATH))\n",
    "else:\n",
    "    classifier.load_state_dict(torch.load(MODEL_PATH, map_location='cpu'))\n",
    "classifier.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simple point cloud coloring mapping\n",
    "def read_pointnet_colors(seg_labels):\n",
    "    map_label_to_rgb = {\n",
    "        1: [0, 255, 0],\n",
    "        2: [0, 0, 255],\n",
    "        3: [255, 0, 0],\n",
    "        4: [255, 0, 255],  # purple\n",
    "        5: [0, 255, 255],  # cyan\n",
    "        6: [255, 255, 0],  # yellow\n",
    "    }\n",
    "    colors = np.array([map_label_to_rgb[label] for label in seg_labels])\n",
    "    return colors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Three.js based visualizer\n",
    "visualizer = JVisualizer()\n",
    "\n",
    "# Basic inference and visualization loop\n",
    "MAX_SAMPLES = 15\n",
    "for samples in range(MAX_SAMPLES):\n",
    "    random_index = randrange(len(test_dataset_seg))\n",
    "    print('[Sample {} / {}]'.format(random_index, len(test_dataset_seg)))\n",
    "    \n",
    "    # clean visualization\n",
    "    visualizer.clear()\n",
    "    clear_output()\n",
    "\n",
    "    # get next sample\n",
    "    point_set, seg = test_dataset_seg.__getitem__(random_index)\n",
    "\n",
    "    # create cloud for visualization\n",
    "    cloud = o3.PointCloud()\n",
    "    cloud.points = o3.Vector3dVector(point_set)\n",
    "    cloud.colors = o3.Vector3dVector(read_pointnet_colors(seg.numpy()))\n",
    "\n",
    "    # perform inference in GPU\n",
    "    points = Variable(point_set.unsqueeze(0))\n",
    "    points = points.transpose(2, 1)\n",
    "    if torch.cuda.is_available():\n",
    "        points = points.cuda()\n",
    "    pred_logsoft, _ = classifier(points)\n",
    "\n",
    "    # move data back to cpu for visualization\n",
    "    pred_logsoft_cpu = pred_logsoft.data.cpu().numpy().squeeze()\n",
    "    pred_soft_cpu = np.exp(pred_logsoft_cpu)\n",
    "    pred_class = np.argmax(pred_soft_cpu)\n",
    "\n",
    "    # let's visualize the input sample\n",
    "    visualizer.add_geometry(cloud)\n",
    "    visualizer.show()\n",
    "    \n",
    "    # Visualize probabilities\n",
    "    plt.xticks(list(classes_dict.values()), list(classes_dict.keys()),rotation=90)\n",
    "    plt.xlabel('Classes')\n",
    "    plt.ylabel('Probabilities')\n",
    "    plt.plot(list(classes_dict.values()), pred_soft_cpu)\n",
    "    plt.show()\n",
    "\n",
    "    input('Your object is a [{}] with probability {:0.3}. Press enter to continue!'\n",
    "          .format(list(classes_dict.keys())[pred_class], pred_soft_cpu[pred_class]))\n",
    "    "
   ]
  }
 ],
 "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}