notebooks/224-3D-segmentation-point-clouds/224-3D-segmentation-point-clouds.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part Segmentation of 3D Point Clouds with OpenVINO™  \n",
    "\n",
    "This notebook demonstrates how to process [point cloud](https://en.wikipedia.org/wiki/Point_cloud) data and run 3D Part Segmentation with OpenVINO. We use the [PointNet](https://arxiv.org/abs/1612.00593) pre-trained model to detect each part of a chair and return its category.\n",
    "## PointNet\n",
    "PointNet was proposed by Charles Ruizhongtai Qi, a researcher at Stanford University in 2016: arXiv:1612.00593 <[PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation](https://arxiv.org/abs/1612.00593)>. The motivation behind the research is to classify and segment 3D representations of images. They use a data structure called point cloud, which is a set of points that represents a 3D shape or object. PointNet provides a unified architecture for applications ranging from object classification, part segmentation, to scene semantic parsing. It is highly efficient and effective, showing strong performance on par or even better than state of the art. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from openvino.runtime import Core\n",
    "\n",
    "sys.path.append(\"../utils\")\n",
    "from notebook_utils import download_file"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare the Model\n",
    "Download the pre-trained PointNet ONNX model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the data and model directories, model source URL and model filename\n",
    "MODEL_DIR = \"model\"\n",
    "os.makedirs(MODEL_DIR, exist_ok=True)\n",
    "download_file(\"https://storage.googleapis.com/ailia-models/pointnet_pytorch/chair_100.onnx\", directory=Path(MODEL_DIR), show_progress=False)\n",
    "onnx_model_path = Path(MODEL_DIR) / \"chair_100.onnx\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Convert the ONNX model to OpenVINO IR. An OpenVINO IR (Intermediate Representation) model consists of an `.xml` file, containing information about network topology, and a `.bin` file, containing the weights and biases binary data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ir_model_xml = onnx_model_path.with_suffix(\".xml\")\n",
    "ir_model_bin = onnx_model_path.with_suffix(\".bin\")\n",
    "\n",
    "if not ir_model_xml.exists():\n",
    "    !mo --input_model $onnx_model_path --output_dir $MODEL_DIR --data_type FP16"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Processing Module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_data(point_file):\n",
    "    \"\"\"\n",
    "    Load the point cloud data and convert it to ndarray\n",
    "\n",
    "    Parameters:\n",
    "        point_file: string, path of .pts data\n",
    "    \"\"\"\n",
    "\n",
    "    point_set = np.loadtxt(point_file).astype(np.float32)\n",
    "\n",
    "    # normailization\n",
    "    point_set = point_set - np.expand_dims(np.mean(point_set, axis=0), 0)  # center\n",
    "    dist = np.max(np.sqrt(np.sum(point_set ** 2, axis=1)), 0)\n",
    "    point_set = point_set / dist  # scale\n",
    "\n",
    "    return point_set\n",
    "\n",
    "\n",
    "def visualize(point_set):\n",
    "    \"\"\"\n",
    "    Create a 3D view for data visualization\n",
    "\n",
    "    Parameters:\n",
    "        point_set: ndarray, the coordinate data in X Y Z format\n",
    "    \"\"\"\n",
    "\n",
    "    fig = plt.figure(dpi=192, figsize=(4, 4))\n",
    "    ax = fig.add_subplot(111, projection='3d')\n",
    "    X = point_set[:, 0]\n",
    "    Y = point_set[:, 2]\n",
    "    Z = point_set[:, 1]\n",
    "\n",
    "    # Scale the view of each axis to adapt to the coordinate data distribution\n",
    "    max_range = np.array([X.max() - X.min(), Y.max() - Y.min(), Z.max() - Z.min()]).max() * 0.5\n",
    "    mid_x = (X.max() + X.min()) * 0.5\n",
    "    mid_y = (Y.max() + Y.min()) * 0.5\n",
    "    mid_z = (Z.max() + Z.min()) * 0.5\n",
    "    ax.set_xlim(mid_x - max_range, mid_x + max_range)\n",
    "    ax.set_ylim(mid_y - max_range, mid_y + max_range)\n",
    "    ax.set_zlim(mid_z - max_range, mid_z + max_range)\n",
    "    \n",
    "    plt.tick_params(labelsize=5)\n",
    "    ax.set_xlabel('X', fontsize=10)\n",
    "    ax.set_ylabel('Y', fontsize=10)\n",
    "    ax.set_zlabel('Z', fontsize=10)\n",
    "\n",
    "    return ax"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize the original 3D data\n",
    "The point cloud data can be downloaded from [ShapeNet](https://shapenet.org/), a large-scale dataset of 3D shapes. Here we select the 3D data of a chair for example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "point_data = \"data/chair.pts\"\n",
    "points = load_data(point_data)\n",
    "X = points[:, 0]\n",
    "Y = points[:, 2]\n",
    "Z = points[:, 1]\n",
    "ax = visualize(points)\n",
    "ax.scatter3D(X, Y, Z, s=5, cmap=\"jet\", marker=\"o\", label='chair')\n",
    "ax.set_title('3D Visualization')\n",
    "plt.legend(loc='upper right', fontsize=8)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run inference\n",
    "Run inference and visualize the results of 3D segmentation. \n",
    "- The input data is a point cloud with `1 batch size`，`3 axis value` (x, y, z) and `arbitrary number of points` (dynamic shape).\n",
    "- The output data is a mask with `1 batch size` and `4 classifcation confidence` for each input point."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parts of a chair\n",
    "classes = ['back', 'seat', 'leg', 'arm']\n",
    "\n",
    "# Preprocess the input data\n",
    "point = points.transpose(1, 0)\n",
    "point = np.expand_dims(point, axis=0)\n",
    "\n",
    "# Read model\n",
    "ie = Core()\n",
    "model = ie.read_model(model=ir_model_xml)\n",
    "\n",
    "print(f\"input shape: {model.input().partial_shape}\")\n",
    "print(f\"output shape: {model.output(0).partial_shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inference\n",
    "compiled_model = ie.compile_model(model=model, device_name=\"CPU\")\n",
    "output_layer = compiled_model.output(0)\n",
    "result = compiled_model([point])[output_layer]\n",
    "\n",
    "# Find the label map for all points of chair with highest confidence\n",
    "pred = np.argmax(result[0], axis=1)\n",
    "ax = visualize(point)\n",
    "for i, name in enumerate([0, 1, 2, 3]):\n",
    "    XCur = []\n",
    "    YCur = []\n",
    "    ZCur = []\n",
    "    for j, nameCur in enumerate(pred):\n",
    "        if name == nameCur:\n",
    "            XCur.append(X[j])\n",
    "            YCur.append(Y[j])\n",
    "            ZCur.append(Z[j])\n",
    "    XCur = np.array(XCur)\n",
    "    YCur = np.array(YCur)\n",
    "    ZCur = np.array(ZCur)\n",
    "\n",
    "    # add current point of the part\n",
    "    ax.scatter(XCur, YCur, ZCur, s=5, cmap=\"jet\", marker=\"o\", label=classes[i])\n",
    "\n",
    "ax.set_title('3D Segmentation Visualization')\n",
    "plt.legend(loc='upper right', fontsize=8)\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.9 64-bit",
   "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.9"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "08af04953a73b86b66cc089a637d3d397b0b73ad05ea59846f770cc21ccdacba"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}