lidar-centerpoint-tvm.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bceb08ca-ca5d-43ef-9247-4c8754eb3099",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Jupyter environment detected. Enabling Open3D WebVisualizer.\n",
      "[Open3D INFO] WebRTC GUI backend enabled.\n",
      "[Open3D INFO] WebRTCWindowSystem: HTTP handshake server disabled.\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import open3d as o3d\n",
    "import copy\n",
    "\n",
    "from rosbags.rosbag2 import Reader\n",
    "from rosbags.serde import deserialize_cdr\n",
    "from rosbags.typesys import get_types_from_idl, register_types\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "\n",
    "import tvm\n",
    "from tvm.contrib import graph_executor\n",
    "from tvm.contrib import graph_runtime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ed6a3eff-4cee-46ee-b89c-df8755992784",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/tohmae/.pyenv/versions/miniconda3-4.7.12/envs/py310tvm/lib/python3.10/site-packages/tvm/:/opt/ros/humble/lib\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "sys.path.append('/opt/ros/humble/local/lib/python3.10/dist-packages')\n",
    "sys.path.append('/opt/ros/humble/lib/python3.10/site-packages')\n",
    "\n",
    "#print(sys.path)\n",
    "print(os.environ['LD_LIBRARY_PATH'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d3610133-389a-49b1-8ca0-e51d74929229",
   "metadata": {},
   "outputs": [],
   "source": [
    "pcd_filename = \"pointcloud/test_pointcloud1.pcd\"\n",
    "timelag = 0.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6ad281c0-e587-4e06-b9aa-0cb74c59379d",
   "metadata": {},
   "outputs": [],
   "source": [
    "pc = o3d.io.read_point_cloud(pcd_filename)\n",
    "pc_msg = np.asarray(pc.points)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f5bc3d1c-61b0-41f0-a8c7-1b79264e18dd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(221733, 3)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pc_msg.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d9a63ca9-2209-4610-bf4f-580fd6e2b04c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Config:\n",
    "    point_cloud_range = [-89.6, -89.6, -3.0, 89.6, 89.6, 5.0]\n",
    "    voxel_size = [0.32, 0.32, 8.0]\n",
    "    point_feature_size = 4\n",
    "    point_dim_size = 3\n",
    "    max_voxel_size = 40000\n",
    "    max_point_in_voxel_size = 32\n",
    "    score_threshold = 0.35\n",
    "    circle_nms_dist_threshold = 1.5\n",
    "    yaw_norm_threshold = 0.0\n",
    "    downsample_factor_ = 1\n",
    "    circle_nms_dist_threshold = 1.5\n",
    "    class_names = [\"CAR\", \"PEDESTRIAN\", \"BICYCLE\"]\n",
    "    rename_car_to_truck_and_bus = True\n",
    "    has_twist = False\n",
    "    \n",
    "    \n",
    "    range_min_x_ = point_cloud_range[0]\n",
    "    range_min_y_ = point_cloud_range[1]\n",
    "    range_min_z_ = point_cloud_range[2]\n",
    "    range_max_x_ = point_cloud_range[3]\n",
    "    range_max_y_ = point_cloud_range[4]\n",
    "    range_max_z_ = point_cloud_range[5]\n",
    "\n",
    "    voxel_size_x_ = voxel_size[0]\n",
    "    voxel_size_y_ = voxel_size[1]\n",
    "    voxel_size_z_ = voxel_size[2]\n",
    "\n",
    "    grid_size_x_ = int((range_max_x_ - range_min_x_) / voxel_size_x_)\n",
    "    grid_size_y_ = int((range_max_y_ - range_min_y_) / voxel_size_y_)\n",
    "    grid_size_z_ = int((range_max_z_ - range_min_z_) / voxel_size_z_)\n",
    "    \n",
    "    offset_x_ = range_min_x_ + voxel_size_x_ / 2\n",
    "    offset_y_ = range_min_y_ + voxel_size_y_ / 2\n",
    "    offset_z_ = range_min_z_ + voxel_size_z_ / 2\n",
    "    \n",
    "    down_grid_size_x_ = int(grid_size_x_ / downsample_factor_)\n",
    "    down_grid_size_y_ = int(grid_size_y_ / downsample_factor_)\n",
    "\n",
    "    range_ = np.asarray((range_min_x_, range_min_y_, range_min_z_, range_max_x_, range_max_y_, range_max_z_))\n",
    "    grid_size_ = np.asarray((grid_size_x_, grid_size_y_, grid_size_z_))\n",
    "    recip_voxel_size_ = np.asarray((1.0 / voxel_size_x_, 1.0 / voxel_size_y_, 1.0 / voxel_size_z_))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "fc69695b-5a2b-4d00-b556-7eb91f5c8640",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = Config()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "903adb41-f7dd-4a2c-8d20-6d2e8946d874",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(560, -89.6)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "config.grid_size_x_, config.range_min_x_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "28c1bc9f-0ad8-4a41-b6aa-296cf1a04f0d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([-89.6, -89.6,  -3. ,  89.6,  89.6,   5. ]),\n",
       " array([560, 560,   1]),\n",
       " array([3.125, 3.125, 0.125]),\n",
       " 560,\n",
       " 560)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "config.range_, config.grid_size_, config.recip_voxel_size_, config.down_grid_size_x_, config.down_grid_size_y_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "67950ef9-1a62-43e8-a89c-8337f382ec3b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.32, 0.32, 8.0, -89.44, -89.44, 1.0)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "config.voxel_size_x_, config.voxel_size_y_, config.voxel_size_z_, config.offset_x_, config.offset_y_, config.offset_z_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "1dfd67f4-92a9-4c89-bf07-54b28ee12517",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pointsToVoxel(pc_msg):\n",
    "    grid_size = config.grid_size_z_ * config.grid_size_y_ * config.grid_size_x_\n",
    "    coord_to_voxel_idx = np.repeat(-1, grid_size)\n",
    "\n",
    "    df_pc = pd.DataFrame(pc_msg,columns=['x','y','z'])\n",
    "    df_pc['timelag'] = timelag\n",
    "    df_pc['coord_x'] = ((df_pc['x'] - config.range_[0]) * config.recip_voxel_size_[0]).astype('int32')\n",
    "    df_pc['coord_y'] = ((df_pc['y'] - config.range_[1]) * config.recip_voxel_size_[1]).astype('int32')\n",
    "    df_pc['coord_z'] = ((df_pc['z'] - config.range_[2]) * config.recip_voxel_size_[2]).astype('int32')\n",
    "    \n",
    "    df_pc = df_pc[(df_pc['coord_x']>=0) & (df_pc['coord_x'] < config.grid_size_[0])]\n",
    "    df_pc = df_pc[(df_pc['coord_y']>=0) & (df_pc['coord_y'] < config.grid_size_[1])]\n",
    "    df_pc = df_pc[(df_pc['coord_z']>=0) & (df_pc['coord_z'] < config.grid_size_[2])]\n",
    "    \n",
    "    df_pc['coord_idx'] = df_pc['coord_z'] * config.grid_size_y_ * config.grid_size_x_ +df_pc['coord_y'] * config.grid_size_x_ + df_pc['coord_x']\n",
    "    \n",
    "    np.random.seed(seed=32)\n",
    "    df_pc['random'] = np.random.rand(len(df_pc))\n",
    "    df_pc['rank'] = df_pc.groupby(['coord_idx'])['random'].rank(method='first')\n",
    "    df_pc = df_pc[df_pc['rank']<=32]\n",
    "    \n",
    "    df_num_points_per_voxel = df_pc['coord_idx'].value_counts().reset_index()\n",
    "    df_num_points_per_voxel.columns = ['coord_idx', 'point_cnt']\n",
    "    df_num_points_per_voxel['voxel_idx'] = df_num_points_per_voxel.index\n",
    "    df_num_points_per_voxel = df_num_points_per_voxel[['voxel_idx','coord_idx','point_cnt']]\n",
    "    voxel_cnt = len(df_num_points_per_voxel)\n",
    "    \n",
    "    df_pc = pd.merge(df_pc, df_num_points_per_voxel, how='inner', on='coord_idx')\n",
    "    df_coordinate_tmp = df_pc[~df_pc.duplicated(subset='voxel_idx')].sort_values('voxel_idx')[['voxel_idx', 'coord_z','coord_y', 'coord_x']].reset_index(drop=True)\n",
    "\n",
    "    df_coordinate = pd.DataFrame(np.arange(0, config.max_voxel_size),columns=['voxel_idx'])\n",
    "    df_coordinate = pd.merge(df_coordinate, df_coordinate_tmp, how='left', on='voxel_idx')\n",
    "    df_coordinate = df_coordinate.fillna(-1)\n",
    "    df_coordinate = df_coordinate.astype('int32')\n",
    "    \n",
    "    df_voxel = pd.DataFrame(np.repeat(np.arange(0,config.max_voxel_size), config.max_point_in_voxel_size),columns=['voxel_idx'])\n",
    "    df_voxel['rank'] = np.tile(np.arange(1,config.max_point_in_voxel_size+1), config.max_voxel_size)\n",
    "    df_voxel = pd.merge(df_voxel, df_pc, how='left', on=['voxel_idx', 'rank'])\n",
    "    nancols = ['x','y','z','timelag']\n",
    "    for nancol in nancols:\n",
    "        df_voxel[nancol] = df_voxel[nancol].fillna(0)\n",
    "    df_voxel = df_voxel[['voxel_idx','rank','x','y','z','timelag']]\n",
    "    \n",
    "    return df_pc, df_voxel, df_coordinate, df_num_points_per_voxel, voxel_cnt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e407f247-918f-42a7-9841-edc560f56418",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_pc, df_voxel, df_coordinate, df_num_points_per_voxel,voxel_cnt = pointsToVoxel(pc_msg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "f12d73e3-73d9-45e3-8d38-d1819a701d1c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(157025, 1280000, 40000, 31227, 31227)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(df_pc), len(df_voxel), len(df_coordinate), len(df_num_points_per_voxel), voxel_cnt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "15741138-790e-416a-ae79-0b8b020eed86",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(   voxel_idx  coord_z  coord_y  coord_x\n",
       " 0          0        0      293      282\n",
       " 1          1        0      260      277\n",
       " 2          2        0      324      432\n",
       " 3          3        0      278      354\n",
       " 4          4        0      242      288,\n",
       "        voxel_idx  coord_z  coord_y  coord_x\n",
       " 39995      39995       -1       -1       -1\n",
       " 39996      39996       -1       -1       -1\n",
       " 39997      39997       -1       -1       -1\n",
       " 39998      39998       -1       -1       -1\n",
       " 39999      39999       -1       -1       -1)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_coordinate.head(), df_coordinate.tail()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "2e794697-92bf-4fdb-8e5d-714b86ea33e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generatefeature(df_pc):\n",
    "    df_pc_mean = df_pc[['x','y','z', 'voxel_idx']].groupby('voxel_idx').mean().reset_index()\n",
    "    df_pc_mean.columns = ['voxel_idx','x_mean', 'y_mean', 'z_mean']\n",
    "    \n",
    "    df_pc['x_offset'] = df_pc['coord_x'] * config.voxel_size_x_ + config.offset_x_\n",
    "    df_pc['y_offset'] = df_pc['coord_y'] * config.voxel_size_y_ + config.offset_y_\n",
    "    \n",
    "    df_feature = pd.merge(df_pc, df_pc_mean, on='voxel_idx')\n",
    "    \n",
    "    df_feature['feature_0'] = df_feature['x']\n",
    "    df_feature['feature_1'] = df_feature['y']\n",
    "    df_feature['feature_2'] = df_feature['z']\n",
    "    df_feature['feature_3'] = df_feature['timelag']\n",
    "\n",
    "    df_feature['feature_4'] = df_feature['x'] - df_feature['x_mean']\n",
    "    df_feature['feature_5'] = df_feature['y'] - df_feature['y_mean']\n",
    "    df_feature['feature_6'] = df_feature['z'] - df_feature['z_mean']\n",
    "\n",
    "    df_feature['feature_7'] = df_feature['x'] - df_feature['x_offset']\n",
    "    df_feature['feature_8'] = df_feature['y'] - df_feature['y_offset']\n",
    "    \n",
    "    df_feature_new = pd.DataFrame(np.repeat(np.arange(0,config.max_voxel_size), config.max_point_in_voxel_size),columns=['voxel_idx'])\n",
    "    df_feature_new['rank'] = np.tile(np.arange(1,config.max_point_in_voxel_size+1), config.max_voxel_size)\n",
    "    df_feature_new = pd.merge(df_feature_new, df_feature, how='left', on=['voxel_idx', 'rank'])\n",
    "\n",
    "    feature_cols = ['feature_{}'.format(i) for i in range(9)]\n",
    "    for nancol in feature_cols:\n",
    "        df_feature_new[nancol] = df_feature_new[nancol].fillna(0)\n",
    "    features = df_feature_new[['feature_{}'.format(i) for i in range(9)]].values.reshape(40000,32,9).astype('float32')\n",
    "    return df_feature, features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "37198e34-506b-4d28-bc04-5abd6ce8dd5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_feature, features = generatefeature(df_pc)\n",
    "coordinates = df_coordinate[['coord_z','coord_y','coord_x']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "bbf4ff5c-e114-4703-85ae-c44dd41df97a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(157025, (40000, 32, 9), (40000, 3))"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(df_feature), features.shape, coordinates.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "a1c79490-0702-4b0d-86fa-6c82cc1fc08b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>x</th>\n",
       "      <th>y</th>\n",
       "      <th>z</th>\n",
       "      <th>timelag</th>\n",
       "      <th>coord_x</th>\n",
       "      <th>coord_y</th>\n",
       "      <th>coord_z</th>\n",
       "      <th>coord_idx</th>\n",
       "      <th>random</th>\n",
       "      <th>rank</th>\n",
       "      <th>...</th>\n",
       "      <th>z_mean</th>\n",
       "      <th>feature_0</th>\n",
       "      <th>feature_1</th>\n",
       "      <th>feature_2</th>\n",
       "      <th>feature_3</th>\n",
       "      <th>feature_4</th>\n",
       "      <th>feature_5</th>\n",
       "      <th>feature_6</th>\n",
       "      <th>feature_7</th>\n",
       "      <th>feature_8</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0.833137</td>\n",
       "      <td>4.236041</td>\n",
       "      <td>-0.029378</td>\n",
       "      <td>0.0</td>\n",
       "      <td>282</td>\n",
       "      <td>293</td>\n",
       "      <td>0</td>\n",
       "      <td>164362</td>\n",
       "      <td>0.858889</td>\n",
       "      <td>28.0</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.020831</td>\n",
       "      <td>0.833137</td>\n",
       "      <td>4.236041</td>\n",
       "      <td>-0.029378</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.010767</td>\n",
       "      <td>-0.041291</td>\n",
       "      <td>-0.008547</td>\n",
       "      <td>0.033137</td>\n",
       "      <td>-0.083959</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0.846620</td>\n",
       "      <td>4.243673</td>\n",
       "      <td>-0.033356</td>\n",
       "      <td>0.0</td>\n",
       "      <td>282</td>\n",
       "      <td>293</td>\n",
       "      <td>0</td>\n",
       "      <td>164362</td>\n",
       "      <td>0.372711</td>\n",
       "      <td>11.0</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.020831</td>\n",
       "      <td>0.846620</td>\n",
       "      <td>4.243673</td>\n",
       "      <td>-0.033356</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.024250</td>\n",
       "      <td>-0.033658</td>\n",
       "      <td>-0.012525</td>\n",
       "      <td>0.046620</td>\n",
       "      <td>-0.076327</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0.860237</td>\n",
       "      <td>4.249470</td>\n",
       "      <td>-0.036439</td>\n",
       "      <td>0.0</td>\n",
       "      <td>282</td>\n",
       "      <td>293</td>\n",
       "      <td>0</td>\n",
       "      <td>164362</td>\n",
       "      <td>0.555129</td>\n",
       "      <td>15.0</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.020831</td>\n",
       "      <td>0.860237</td>\n",
       "      <td>4.249470</td>\n",
       "      <td>-0.036439</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.037866</td>\n",
       "      <td>-0.027862</td>\n",
       "      <td>-0.015608</td>\n",
       "      <td>0.060237</td>\n",
       "      <td>-0.070530</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0.887949</td>\n",
       "      <td>4.232125</td>\n",
       "      <td>-0.028625</td>\n",
       "      <td>0.0</td>\n",
       "      <td>282</td>\n",
       "      <td>293</td>\n",
       "      <td>0</td>\n",
       "      <td>164362</td>\n",
       "      <td>0.736670</td>\n",
       "      <td>22.0</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.020831</td>\n",
       "      <td>0.887949</td>\n",
       "      <td>4.232125</td>\n",
       "      <td>-0.028625</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.065578</td>\n",
       "      <td>-0.045207</td>\n",
       "      <td>-0.007794</td>\n",
       "      <td>0.087949</td>\n",
       "      <td>-0.087875</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0.901640</td>\n",
       "      <td>4.236015</td>\n",
       "      <td>-0.030749</td>\n",
       "      <td>0.0</td>\n",
       "      <td>282</td>\n",
       "      <td>293</td>\n",
       "      <td>0</td>\n",
       "      <td>164362</td>\n",
       "      <td>0.816205</td>\n",
       "      <td>25.0</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.020831</td>\n",
       "      <td>0.901640</td>\n",
       "      <td>4.236015</td>\n",
       "      <td>-0.030749</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.079269</td>\n",
       "      <td>-0.041316</td>\n",
       "      <td>-0.009918</td>\n",
       "      <td>0.101640</td>\n",
       "      <td>-0.083985</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 26 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "          x         y         z  timelag  coord_x  coord_y  coord_z  \\\n",
       "0  0.833137  4.236041 -0.029378      0.0      282      293        0   \n",
       "1  0.846620  4.243673 -0.033356      0.0      282      293        0   \n",
       "2  0.860237  4.249470 -0.036439      0.0      282      293        0   \n",
       "3  0.887949  4.232125 -0.028625      0.0      282      293        0   \n",
       "4  0.901640  4.236015 -0.030749      0.0      282      293        0   \n",
       "\n",
       "   coord_idx    random  rank  ...    z_mean  feature_0  feature_1  feature_2  \\\n",
       "0     164362  0.858889  28.0  ... -0.020831   0.833137   4.236041  -0.029378   \n",
       "1     164362  0.372711  11.0  ... -0.020831   0.846620   4.243673  -0.033356   \n",
       "2     164362  0.555129  15.0  ... -0.020831   0.860237   4.249470  -0.036439   \n",
       "3     164362  0.736670  22.0  ... -0.020831   0.887949   4.232125  -0.028625   \n",
       "4     164362  0.816205  25.0  ... -0.020831   0.901640   4.236015  -0.030749   \n",
       "\n",
       "   feature_3  feature_4  feature_5  feature_6  feature_7  feature_8  \n",
       "0        0.0   0.010767  -0.041291  -0.008547   0.033137  -0.083959  \n",
       "1        0.0   0.024250  -0.033658  -0.012525   0.046620  -0.076327  \n",
       "2        0.0   0.037866  -0.027862  -0.015608   0.060237  -0.070530  \n",
       "3        0.0   0.065578  -0.045207  -0.007794   0.087949  -0.087875  \n",
       "4        0.0   0.079269  -0.041316  -0.009918   0.101640  -0.083985  \n",
       "\n",
       "[5 rows x 26 columns]"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_feature.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ff22c7e-7b1d-45be-bd94-c71da06ec10f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "8e67fb83-25e4-4297-8154-521099fb2250",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generatemodule(model_dir):\n",
    "    loaded_lib = tvm.runtime.load_module(model_dir + '/deploy_lib.so')\n",
    "    print(\"deploy loaded\")\n",
    "    loaded_json = open(model_dir + '/deploy_graph.json').read()\n",
    "    print(\"json loaded\")\n",
    "    loaded_params = bytearray(open(model_dir + '/deploy_param.params', \"rb\").read())\n",
    "    print(\"params loaded\")\n",
    "    \n",
    "    module = graph_runtime.create(loaded_json, loaded_lib, tvm.cpu(0))\n",
    "    module.load_params(loaded_params)\n",
    "    \n",
    "    return module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "d419d409-b674-41bd-a61b-6dd5cb120c37",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_en_dir = \"models/centerpoint_encoder\"\n",
    "model_bk_dir = \"models/centerpoint_backbone\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "487ef4ba-944e-4eb5-b705-fd11c288e5fc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "deploy loaded\n",
      "json loaded\n",
      "params loaded\n",
      "deploy loaded\n",
      "json loaded\n",
      "params loaded\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tohmae/.pyenv/versions/miniconda3-4.7.12/envs/py310tvm/lib/python3.10/site-packages/tvm/contrib/graph_runtime.py:25: UserWarning: This function has been moved to tvm.contrib.graph_executor and will be removed in the next TVM release\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "module_en = generatemodule(model_en_dir)\n",
    "module_bk = generatemodule(model_bk_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "9963892a-f2e9-46bd-90b5-8b150a6e7fb9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(({\"input_features\": [40000, 32, 9]}, {\"input_features\": \"float32\"}),\n",
       " ({\"spatial_features\": [1, 32, 560, 560]}, {\"spatial_features\": \"float32\"}))"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "module_en.get_input_info(), module_bk.get_input_info()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "63c00806-8321-49c8-bf8e-94fb27a6875c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generatefunction(model_dir):\n",
    "    lib = tvm.runtime.load_module(model_dir + '/preprocess.so')\n",
    "    print(\"deploy loaded\")\n",
    "    fun = lib[\"scatter\"]\n",
    "    return fun"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "251eb8be-37ad-4a5b-b370-0d882e4b0617",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "deploy loaded\n"
     ]
    }
   ],
   "source": [
    "scatter = generatefunction(model_bk_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "5e169692-9081-4c91-b6e7-8c8bbfca0b73",
   "metadata": {},
   "outputs": [],
   "source": [
    "def encoder(features):\n",
    "    # VE_IE\n",
    "    module_en.set_input('input_features', tvm.nd.array(features))\n",
    "    module_en.run()\n",
    "    output_en = module_en.get_output(0)\n",
    "    \n",
    "    return output_en\n",
    "\n",
    "\n",
    "def scatter_ie(input_0, coordinate):\n",
    "    output_sc_np = np.array(np.random.random(size=(1,32,560,560)),dtype=np.float32)\n",
    "    \n",
    "    input_1: tvm.runtime.NDArray = tvm.nd.array(coordinates)\n",
    "    output_sc: tvm.runtime.NDArray = tvm.nd.array(output_sc_np)\n",
    "    scatter(input_0, input_1, output_sc)\n",
    "    \n",
    "    return output_sc\n",
    "\n",
    "def backbone(input_bk):\n",
    "    module_bk.set_input('spatial_features', input_bk)\n",
    "    module_bk.run()\n",
    "    \n",
    "    output_cols = ['heatmap','offset','z','dim','rot','vel']\n",
    "    outputs = {}\n",
    "    for i, output_col in enumerate(output_cols):\n",
    "        outputs[output_col] = module_bk.get_output(i).asnumpy()[0]\n",
    "        \n",
    "    return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "d27669e1-f5e7-45cd-841d-0b0d86e45097",
   "metadata": {},
   "outputs": [],
   "source": [
    "def TSP_pipeline(features, coordinates):\n",
    "    output_en = encoder(features)\n",
    "    output_sc = scatter_ie(output_en, coordinates)\n",
    "    output = backbone(output_sc)\n",
    "    \n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "1863e82a-405c-43dc-a749-23fec7cd7e80",
   "metadata": {},
   "outputs": [],
   "source": [
    "outputs = TSP_pipeline(features, coordinates)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "e880c53b-35b7-48f8-af99-b5dc487d43b9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "heatmap (3, 560, 560)\n",
      "offset (2, 560, 560)\n",
      "z (1, 560, 560)\n",
      "dim (3, 560, 560)\n",
      "rot (2, 560, 560)\n",
      "vel (2, 560, 560)\n"
     ]
    }
   ],
   "source": [
    "for k,v in outputs.items():\n",
    "    print(k, v.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "87822fd3-1f08-4432-8933-fab576c66e2d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.7363652"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.max(outputs['heatmap'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82d7464a-a501-4b71-9a13-9429d144581a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "87d88f0d-e24d-4a25-8385-719de4d03605",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sigmoid(x):\n",
    "    return 1 / (1 + np.exp(-x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "79301460-8d49-4fdc-8f87-e1b2e2cc1583",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dist2dPow(df1, df2):\n",
    "    return np.power(df1.x-df2.x, 2) + np.power(df1.y-df2.y, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "36579b2e-bf0d-4c60-be70-f504311fa923",
   "metadata": {},
   "outputs": [],
   "source": [
    "def circleNMS(df_det_boxes3d_nonms):\n",
    "    num_boxes3d = len(df_det_boxes3d_nonms)\n",
    "    dist2d_pow_thresh = np.power(config.circle_nms_dist_threshold, 2)\n",
    "    suppress = [False] * num_boxes3d\n",
    "    keep_mask = [False] * num_boxes3d\n",
    "    num_to_keep = 0\n",
    "    \n",
    "    for i in range(num_boxes3d):\n",
    "        if suppress[i]:\n",
    "            keep_mask[i] = False\n",
    "        else:\n",
    "            keep_mask[i] = True\n",
    "            num_to_keep += 1\n",
    "            for j in range(i+1, num_boxes3d):\n",
    "                if suppress[j]:\n",
    "                    continue\n",
    "                if dist2dPow(df_det_boxes3d_nonms.loc[i], df_det_boxes3d_nonms.loc[j]) < dist2d_pow_thresh:\n",
    "                    suppress[j] = True\n",
    "    return num_to_keep, keep_mask "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "faad1ca3-c1f3-4c7b-82e1-9fc78bb1403e",
   "metadata": {},
   "outputs": [],
   "source": [
    "down_grid_size = config.down_grid_size_y_ * config.down_grid_size_x_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "d15ab3e5-4b37-48cd-bb05-ef241019beb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generateBoxes3D(outputs):\n",
    "    max_score = np.max(sigmoid(outputs['heatmap']), axis=0)\n",
    "    label = np.argmax(sigmoid(outputs['heatmap']), axis=0)\n",
    "    \n",
    "    offset_x = outputs['offset'][0,:,:]\n",
    "    offset_y = outputs['offset'][1,:,:]\n",
    "    y_i = np.repeat(np.arange(0, config.grid_size_x_), config.grid_size_y_).reshape((config.grid_size_x_, config.grid_size_x_))\n",
    "    x_i = y_i.T\n",
    "\n",
    "    x = config.voxel_size_x_ * (offset_x + x_i) + config.range_min_x_\n",
    "    y = config.voxel_size_y_ * (offset_y + y_i) + config.range_min_y_\n",
    "    \n",
    "    z = outputs['z'][0]\n",
    "    \n",
    "    w = outputs['dim'][0]\n",
    "    l = outputs['dim'][1]\n",
    "    h = outputs['dim'][2]\n",
    "    yaw_sin = outputs['rot'][0]\n",
    "    yaw_cos = outputs['rot'][1]\n",
    "    yaw_norm = np.sqrt(yaw_sin * yaw_sin + yaw_cos * yaw_cos)\n",
    "    vel_x = outputs['vel'][0]\n",
    "    vel_y = outputs['vel'][1]\n",
    "    \n",
    "    score = np.where(yaw_norm >= config.yaw_norm_threshold, max_score, 0.0)\n",
    "    length = np.exp(l)\n",
    "    width = np.exp(w)\n",
    "    height = np.exp(h)\n",
    "    yaw = np.arctan2(yaw_sin, yaw_cos)\n",
    "    \n",
    "    boxes3d = np.stack([label, score, x, y, z, length, width, height, yaw, vel_x, vel_y])\n",
    "\n",
    "    num_det_boxes3d = np.sum(max_score > config.score_threshold)\n",
    "    \n",
    "    box3d_columns = ['label','score','x','y','z','length','width','height','yaw','vel_x','vel_y']\n",
    "    index = np.where(boxes3d[1,:,:] > config.score_threshold)\n",
    "    df_det_boxes3d_nonms = pd.DataFrame(columns=box3d_columns)\n",
    "    for i, colnm in enumerate(box3d_columns):\n",
    "        df_det_boxes3d_nonms[colnm] = boxes3d[i,:,:][index]\n",
    "    \n",
    "    df_det_boxes3d_nonms = df_det_boxes3d_nonms.sort_values('score', ascending=False).reset_index(drop=True)\n",
    "    \n",
    "    num_to_keep, keep_mask = circleNMS(df_det_boxes3d_nonms)\n",
    "    \n",
    "    box3d_columns = ['label','score','x','y','z','length','width','height','yaw','vel_x','vel_y']\n",
    "    df_det_boxes3d = pd.DataFrame(columns=box3d_columns)\n",
    "    for idx, mask in enumerate(keep_mask):\n",
    "        if mask:\n",
    "#        print(idx,df_det_boxes3d_nonms.loc[idx])\n",
    "#        df_det_boxes3d.append(df_det_boxes3d_nonms.loc[idx],ignore_index=True)\n",
    "            df_det_boxes3d = pd.concat([df_det_boxes3d, pd.DataFrame(df_det_boxes3d_nonms.loc[idx]).T])\n",
    "#            print(idx)\n",
    "    \n",
    "    return df_det_boxes3d_nonms, df_det_boxes3d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "d68ff007-aeaa-4b53-9929-377d63006dfd",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_det_boxes3d_nonms, df_det_boxes3d = generateBoxes3D(outputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "d295ae6e-2e0b-47aa-bb51-0efba31446f9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>label</th>\n",
       "      <th>score</th>\n",
       "      <th>x</th>\n",
       "      <th>y</th>\n",
       "      <th>z</th>\n",
       "      <th>length</th>\n",
       "      <th>width</th>\n",
       "      <th>height</th>\n",
       "      <th>yaw</th>\n",
       "      <th>vel_x</th>\n",
       "      <th>vel_y</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0.0</td>\n",
       "      <td>0.850225</td>\n",
       "      <td>-2.508435</td>\n",
       "      <td>-3.290982</td>\n",
       "      <td>0.866475</td>\n",
       "      <td>5.068622</td>\n",
       "      <td>2.205514</td>\n",
       "      <td>1.528983</td>\n",
       "      <td>-1.284788</td>\n",
       "      <td>6.534736e+00</td>\n",
       "      <td>-1.932273e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0.0</td>\n",
       "      <td>0.842624</td>\n",
       "      <td>0.085640</td>\n",
       "      <td>-7.173362</td>\n",
       "      <td>0.920280</td>\n",
       "      <td>4.899607</td>\n",
       "      <td>2.116013</td>\n",
       "      <td>1.546768</td>\n",
       "      <td>-1.290130</td>\n",
       "      <td>1.568295e+00</td>\n",
       "      <td>-3.619276e-01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0.0</td>\n",
       "      <td>0.787792</td>\n",
       "      <td>47.795307</td>\n",
       "      <td>-25.431138</td>\n",
       "      <td>-0.946414</td>\n",
       "      <td>4.440732</td>\n",
       "      <td>2.033822</td>\n",
       "      <td>1.610885</td>\n",
       "      <td>1.864955</td>\n",
       "      <td>-3.251190e-01</td>\n",
       "      <td>1.885750e-01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>0.0</td>\n",
       "      <td>0.709717</td>\n",
       "      <td>34.679834</td>\n",
       "      <td>10.039294</td>\n",
       "      <td>0.010503</td>\n",
       "      <td>4.601593</td>\n",
       "      <td>2.009392</td>\n",
       "      <td>1.510614</td>\n",
       "      <td>0.367650</td>\n",
       "      <td>-2.001595e-09</td>\n",
       "      <td>-2.956825e-09</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>0.0</td>\n",
       "      <td>0.616896</td>\n",
       "      <td>61.304490</td>\n",
       "      <td>-29.826901</td>\n",
       "      <td>-1.233422</td>\n",
       "      <td>5.035407</td>\n",
       "      <td>2.289372</td>\n",
       "      <td>2.057414</td>\n",
       "      <td>-1.237693</td>\n",
       "      <td>1.508202e+01</td>\n",
       "      <td>-5.134261e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>1.0</td>\n",
       "      <td>0.543303</td>\n",
       "      <td>-15.579638</td>\n",
       "      <td>8.215860</td>\n",
       "      <td>1.709832</td>\n",
       "      <td>0.922758</td>\n",
       "      <td>0.836147</td>\n",
       "      <td>1.772237</td>\n",
       "      <td>-1.484769</td>\n",
       "      <td>4.384600e-01</td>\n",
       "      <td>-4.435377e-02</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>1.0</td>\n",
       "      <td>0.519902</td>\n",
       "      <td>27.119826</td>\n",
       "      <td>-45.006624</td>\n",
       "      <td>-0.043317</td>\n",
       "      <td>0.951638</td>\n",
       "      <td>0.843262</td>\n",
       "      <td>1.787209</td>\n",
       "      <td>-2.725189</td>\n",
       "      <td>2.469870e-01</td>\n",
       "      <td>5.204241e-01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>1.0</td>\n",
       "      <td>0.497980</td>\n",
       "      <td>49.559895</td>\n",
       "      <td>-4.101688</td>\n",
       "      <td>-0.912111</td>\n",
       "      <td>0.951298</td>\n",
       "      <td>0.850127</td>\n",
       "      <td>1.792570</td>\n",
       "      <td>1.971386</td>\n",
       "      <td>-8.062065e-01</td>\n",
       "      <td>5.138541e-01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>22</th>\n",
       "      <td>0.0</td>\n",
       "      <td>0.467097</td>\n",
       "      <td>-3.036579</td>\n",
       "      <td>-81.130304</td>\n",
       "      <td>2.087189</td>\n",
       "      <td>14.347609</td>\n",
       "      <td>3.302393</td>\n",
       "      <td>4.182689</td>\n",
       "      <td>0.254189</td>\n",
       "      <td>-3.153192e-02</td>\n",
       "      <td>-1.565528e-01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>26</th>\n",
       "      <td>0.0</td>\n",
       "      <td>0.436712</td>\n",
       "      <td>17.422580</td>\n",
       "      <td>29.909503</td>\n",
       "      <td>1.123179</td>\n",
       "      <td>8.649951</td>\n",
       "      <td>2.939437</td>\n",
       "      <td>3.580776</td>\n",
       "      <td>-2.826342</td>\n",
       "      <td>-2.001595e-09</td>\n",
       "      <td>-2.956825e-09</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>27</th>\n",
       "      <td>0.0</td>\n",
       "      <td>0.433634</td>\n",
       "      <td>-2.106905</td>\n",
       "      <td>-79.525296</td>\n",
       "      <td>2.310956</td>\n",
       "      <td>16.507183</td>\n",
       "      <td>3.532621</td>\n",
       "      <td>4.516965</td>\n",
       "      <td>-2.695662</td>\n",
       "      <td>3.390232e+00</td>\n",
       "      <td>7.315186e+00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>30</th>\n",
       "      <td>1.0</td>\n",
       "      <td>0.425671</td>\n",
       "      <td>8.757515</td>\n",
       "      <td>-37.324641</td>\n",
       "      <td>0.688944</td>\n",
       "      <td>0.874540</td>\n",
       "      <td>0.814139</td>\n",
       "      <td>1.513363</td>\n",
       "      <td>0.843315</td>\n",
       "      <td>-1.625371e-01</td>\n",
       "      <td>-3.192537e-02</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    label     score          x          y         z     length     width  \\\n",
       "0     0.0  0.850225  -2.508435  -3.290982  0.866475   5.068622  2.205514   \n",
       "2     0.0  0.842624   0.085640  -7.173362  0.920280   4.899607  2.116013   \n",
       "4     0.0  0.787792  47.795307 -25.431138 -0.946414   4.440732  2.033822   \n",
       "6     0.0  0.709717  34.679834  10.039294  0.010503   4.601593  2.009392   \n",
       "10    0.0  0.616896  61.304490 -29.826901 -1.233422   5.035407  2.289372   \n",
       "14    1.0  0.543303 -15.579638   8.215860  1.709832   0.922758  0.836147   \n",
       "16    1.0  0.519902  27.119826 -45.006624 -0.043317   0.951638  0.843262   \n",
       "18    1.0  0.497980  49.559895  -4.101688 -0.912111   0.951298  0.850127   \n",
       "22    0.0  0.467097  -3.036579 -81.130304  2.087189  14.347609  3.302393   \n",
       "26    0.0  0.436712  17.422580  29.909503  1.123179   8.649951  2.939437   \n",
       "27    0.0  0.433634  -2.106905 -79.525296  2.310956  16.507183  3.532621   \n",
       "30    1.0  0.425671   8.757515 -37.324641  0.688944   0.874540  0.814139   \n",
       "\n",
       "      height       yaw         vel_x         vel_y  \n",
       "0   1.528983 -1.284788  6.534736e+00 -1.932273e+00  \n",
       "2   1.546768 -1.290130  1.568295e+00 -3.619276e-01  \n",
       "4   1.610885  1.864955 -3.251190e-01  1.885750e-01  \n",
       "6   1.510614  0.367650 -2.001595e-09 -2.956825e-09  \n",
       "10  2.057414 -1.237693  1.508202e+01 -5.134261e+00  \n",
       "14  1.772237 -1.484769  4.384600e-01 -4.435377e-02  \n",
       "16  1.787209 -2.725189  2.469870e-01  5.204241e-01  \n",
       "18  1.792570  1.971386 -8.062065e-01  5.138541e-01  \n",
       "22  4.182689  0.254189 -3.153192e-02 -1.565528e-01  \n",
       "26  3.580776 -2.826342 -2.001595e-09 -2.956825e-09  \n",
       "27  4.516965 -2.695662  3.390232e+00  7.315186e+00  \n",
       "30  1.513363  0.843315 -1.625371e-01 -3.192537e-02  "
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_det_boxes3d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "6f1bfdee-326f-4eab-99cb-e2a2143996bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "add_types = {}\n",
    "\n",
    "for pathstr in [\n",
    "    'autoware_auto_msgs/autoware_auto_perception_msgs/msg/DetectedObjects.idl',\n",
    "    'autoware_auto_msgs/autoware_auto_perception_msgs/msg/DetectedObject.idl',\n",
    "    'autoware_auto_msgs/autoware_auto_perception_msgs/msg/ObjectClassification.idl',\n",
    "    'autoware_auto_msgs/autoware_auto_perception_msgs/msg/DetectedObjectKinematics.idl',\n",
    "    'autoware_auto_msgs/autoware_auto_perception_msgs/msg/Shape.idl',\n",
    "]:\n",
    "    msgpath = Path(pathstr)\n",
    "    msgdef = msgpath.read_text(encoding='utf-8')\n",
    "    add_types.update(get_types_from_idl(msgdef))\n",
    "\n",
    "register_types(add_types)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "862bf4ac-6039-4c82-a3d2-083a17832668",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rosbags.typesys.types import autoware_auto_perception_msgs__msg__ObjectClassification as ObjectClassification\n",
    "from rosbags.typesys.types import autoware_auto_perception_msgs__msg__DetectedObject as DetectedObject\n",
    "from rosbags.typesys.types import autoware_auto_perception_msgs__msg__DetectedObjectKinematics as DetectedObjectKinematics\n",
    "from rosbags.typesys.types import autoware_auto_perception_msgs__msg__Shape as Shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "f2852111-ce24-4b1d-b71b-2b64396426a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rosbags.typesys.types import geometry_msgs__msg__Point as Point\n",
    "from rosbags.typesys.types import geometry_msgs__msg__Quaternion as Quaternion\n",
    "from rosbags.typesys.types import geometry_msgs__msg__Pose as Pose\n",
    "from rosbags.typesys.types import geometry_msgs__msg__Vector3 as Vector3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e738d534-6d2f-41ca-93d6-0824f6f7231a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "c7e3f857-fbe1-49e6-84d2-a91f9228f243",
   "metadata": {},
   "outputs": [],
   "source": [
    "def getSemanticType(class_name):\n",
    "    if class_name == 'CAR':\n",
    "        return ObjectClassification.CAR\n",
    "    elif class_name == 'PEDESTRIAN':\n",
    "        return ObjectClassification.PEDESTRIAN\n",
    "    elif class_name == 'BUS':\n",
    "        return ObjectClassification.BUS\n",
    "    elif class_name == 'BICYCLE':\n",
    "        return ObjectClassification.BICYCLE\n",
    "    else:\n",
    "        return ObjectClassification.UNKNOWN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "a00dd675-8e18-4846-b028-668356573f4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def createPoint(x, y, z):\n",
    "    return Point(x,y,z)\n",
    "\n",
    "def createTranslation(x,y,z):\n",
    "    return Vector3(x,y,z)\n",
    "\n",
    "def createQuaternionFromYaw(yaw):\n",
    "    x = 0\n",
    "    y = 0\n",
    "    z = np.sin(yaw / 2)\n",
    "    w = np.cos(yaw / 2)\n",
    "    return Quaternion(x,y,z,w)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "fa0bb91d-e1b5-4308-916e-0a5886029c55",
   "metadata": {},
   "outputs": [],
   "source": [
    "def isCarLikeVehicleLabel(label):\n",
    "    return (label == ObjectClassification.CAR) | (label == ObjectClassification.TRUCK) | (label == ObjectClassification.BUS) | (label == ObjectClassification.TRAILER)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "fb2cd6d5-8758-4153-9e56-701fa05fa484",
   "metadata": {},
   "outputs": [],
   "source": [
    "def box3DToDetectedObject(box3d):\n",
    "    # def ObjectClassification\n",
    "    label = getSemanticType(config.class_names[int(box3d.label)])\n",
    "    probability = 1.0\n",
    "    \n",
    "    l = box3d.length\n",
    "    w = box3d.width\n",
    "    if label == ObjectClassification.CAR and config.rename_car_to_truck_and_bus:\n",
    "        if (w * l > 2.2 * 5.5) and (w * l <= 2.5 * 7.9):\n",
    "            label = ObjectClassification.TRUCK\n",
    "            print(\"TRUCK\")\n",
    "        elif w * l > 2.5 * 7.9:\n",
    "            label = ObjectClassification.BUS\n",
    "            print(\"BUS\")\n",
    "            \n",
    "    classification = ObjectClassification(label=label,probability=probability)\n",
    "    \n",
    "    \n",
    "    # def kinematics\n",
    "    if isCarLikeVehicleLabel(classification.label):\n",
    "        orientation_availability = DetectedObjectKinematics.SIGN_UNKNOWN\n",
    "    else:\n",
    "        orientation_availability = None\n",
    "        print(\"Not Car\")\n",
    "\n",
    "    pi = 3.14159265358979323846\n",
    "    yaw = -box3d.yaw - pi / 2\n",
    "    position = createPoint(box3d.x, box3d.y, box3d.z)\n",
    "    orientation = createQuaternionFromYaw(yaw)\n",
    "    pose = Pose(position=position, orientation=orientation)\n",
    "    kinematics = DetectedObjectKinematics(pose_with_covariance=pose, has_position_covariance=None, orientation_availability=orientation_availability, twist_with_covariance=None, has_twist=None, has_twist_covariance=None)\n",
    "    \n",
    "    # def shape\n",
    "    shape_type = Shape.BOUNDING_BOX\n",
    "    shape_dimensions = createTranslation(box3d.length, box3d.width, box3d.height)\n",
    "    shape = Shape(type=shape_type, dimensions = shape_dimensions,footprint=None)\n",
    "    \n",
    "    # def DetectedObject\n",
    "    obj = DetectedObject(existence_probability= box3d.score, classification=classification, kinematics=kinematics, shape=shape)\n",
    "    \n",
    "    return obj"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "687ecb0c-5b31-41e8-9faa-c96cc8f1e415",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "2\n",
      "4\n",
      "6\n",
      "10\n",
      "14\n",
      "Not Car\n",
      "16\n",
      "Not Car\n",
      "18\n",
      "Not Car\n",
      "22\n",
      "BUS\n",
      "26\n",
      "BUS\n",
      "27\n",
      "BUS\n",
      "30\n",
      "Not Car\n"
     ]
    }
   ],
   "source": [
    "objs = []\n",
    "for index, box3d in df_det_boxes3d.iterrows():\n",
    "    print(index)\n",
    "    obj = box3DToDetectedObject(box3d)\n",
    "    objs.append(obj)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "7f633563-6df1-41cb-b2e1-61f9588f0b48",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "geometry_msgs__msg__Quaternion(x=0, y=0, z=-0.14251736293737718, w=0.9897923020822984, __msgtype__='geometry_msgs/msg/Quaternion')\n",
      "[[ 0.9593776   0.28212518  0.        ]\n",
      " [-0.28212518  0.9593776  -0.        ]\n",
      " [-0.          0.          1.        ]]\n",
      "[[ 0.9593776   0.28212518  0.         -2.50843517]\n",
      " [-0.28212518  0.9593776  -0.         -3.29098198]\n",
      " [-0.          0.          1.          0.8664754 ]\n",
      " [ 0.          0.          0.          1.        ]]\n"
     ]
    }
   ],
   "source": [
    "position = objs[0].kinematics.pose_with_covariance.position\n",
    "orientation = objs[0].kinematics.pose_with_covariance.orientation\n",
    "print(orientation)\n",
    "pose_affine = np.eye(4)\n",
    "R = o3d.geometry.get_rotation_matrix_from_quaternion([orientation.w, orientation.x, orientation.y, orientation.z])\n",
    "print(R)\n",
    "pose_affine[:3,:3] = R\n",
    "pose_affine[0,3] = position.x\n",
    "pose_affine[1,3] = position.y\n",
    "pose_affine[2,3] = position.z\n",
    "print(pose_affine)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "7c60b3da-3042-41a5-a695-d73b1a5787d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def getVertices(shape, pose_affine):\n",
    "    pcd = o3d.geometry.PointCloud()\n",
    "    pcd.points.append([-shape.dimensions.x / 2.0, -shape.dimensions.y / 2.0,  shape.dimensions.z / 2.0])\n",
    "    pcd.points.append([-shape.dimensions.x / 2.0,  shape.dimensions.y / 2.0,  shape.dimensions.z / 2.0])\n",
    "    pcd.points.append([-shape.dimensions.x / 2.0,  shape.dimensions.y / 2.0, -shape.dimensions.z / 2.0])\n",
    "    pcd.points.append([ shape.dimensions.x / 2.0,  shape.dimensions.y / 2.0,  shape.dimensions.z / 2.0])\n",
    "    pcd.points.append([ shape.dimensions.x / 2.0,  shape.dimensions.y / 2.0, -shape.dimensions.z / 2.0])\n",
    "    pcd.points.append([ shape.dimensions.x / 2.0, -shape.dimensions.y / 2.0,  shape.dimensions.z / 2.0])\n",
    "    pcd.points.append([ shape.dimensions.x / 2.0, -shape.dimensions.y / 2.0, -shape.dimensions.z / 2.0])\n",
    "    pcd.points.append([-shape.dimensions.x / 2.0, -shape.dimensions.y / 2.0, -shape.dimensions.z / 2.0])\n",
    "    pcd_new = copy.deepcopy(pcd).transform(pose_affine)\n",
    "    return np.asarray(pcd_new.points)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "a885fc5c-0279-4220-9df4-f1f33b075c02",
   "metadata": {},
   "outputs": [],
   "source": [
    "def streamFace(v1, v2, v3):\n",
    "    return \"3 {} {} {}\".format(v1,v2,v3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "b41838a4-54af-44fd-b345-38b93ec39a54",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dumpDetectionsAsMech(objs, output_path):\n",
    "    index = 0\n",
    "    num_detections = len(objs)\n",
    "    vertices_stream = []\n",
    "    faces_stream = []\n",
    "    with open(output_path, 'w') as f:\n",
    "        print(\"ply\", file=f)\n",
    "        print(\"format ascii 1.0\", file=f)\n",
    "        print(\"comment created by lidar_centerpoint\", file=f)\n",
    "        print(\"element vertex {}\".format(8*num_detections), file=f)\n",
    "        print(\"property float x\", file=f)\n",
    "        print(\"property float y\", file=f)\n",
    "        print(\"property float z\", file=f)\n",
    "        print(\"element face {}\".format(12*num_detections), file=f)\n",
    "        print(\"property list uchar uint vertex_indices\", file=f)\n",
    "        print(\"end_header\", file=f)\n",
    "        \n",
    "        for obj in objs:\n",
    "            position = obj.kinematics.pose_with_covariance.position\n",
    "            orientation = obj.kinematics.pose_with_covariance.orientation\n",
    "            shape = obj.shape\n",
    "            \n",
    "            pose_affine = np.eye(4)\n",
    "            R = o3d.geometry.get_rotation_matrix_from_quaternion([orientation.w, orientation.x, orientation.y, orientation.z])\n",
    "\n",
    "            pose_affine[:3,:3] = R\n",
    "            pose_affine[0,3] = position.x\n",
    "            pose_affine[1,3] = position.y\n",
    "            pose_affine[2,3] = position.z\n",
    "            \n",
    "            vertices = getVertices(shape, pose_affine)\n",
    "            \n",
    "            for vertex in vertices:\n",
    "                vertices_stream.append(\"{} {} {}\".format(vertex[0], vertex[1], vertex[2]))\n",
    "            \n",
    "            \n",
    "            faces_stream.append(streamFace(index + 1, index + 3, index + 4))\n",
    "            faces_stream.append(streamFace(index + 3, index + 5, index + 6))\n",
    "            faces_stream.append(streamFace(index + 0, index + 7, index + 5))\n",
    "            faces_stream.append(streamFace(index + 7, index + 2, index + 4))\n",
    "            faces_stream.append(streamFace(index + 5, index + 3, index + 1))\n",
    "            faces_stream.append(streamFace(index + 7, index + 0, index + 2))\n",
    "            faces_stream.append(streamFace(index + 2, index + 1, index + 4))\n",
    "            faces_stream.append(streamFace(index + 4, index + 3, index + 6))\n",
    "            faces_stream.append(streamFace(index + 5, index + 7, index + 6))\n",
    "            faces_stream.append(streamFace(index + 6, index + 7, index + 4))\n",
    "            faces_stream.append(streamFace(index + 0, index + 5, index + 1))\n",
    "            index += 8\n",
    "        print('\\n'.join(vertices_stream),file=f)\n",
    "        print('\\n'.join(faces_stream),file=f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "a4ba5b9e-2a20-4e8d-926e-773b324b7b13",
   "metadata": {},
   "outputs": [],
   "source": [
    "dumpDetectionsAsMech(objs,'aaa.ply')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "de771dc4-38f5-4315-9628-bae618709ce0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;33m[Open3D WARNING] Read PLY failed: unable to read file: aaa.ply\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] GLFW Error: X11: The DISPLAY environment variable is missing\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] Failed to initialize GLFW\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] [DrawGeometries] Failed creating OpenGL window.\u001b[0;m\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "RPly: Unexpected end of file\n",
      "RPly: Error reading 'vertex_indices' of 'face' number 132\n"
     ]
    }
   ],
   "source": [
    "mesh = o3d.io.read_triangle_mesh('aaa.ply')\n",
    "mesh_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0, origin=[0, 0, 0])\n",
    "detection_lines = o3d.geometry.LineSet.create_from_triangle_mesh(mesh)\n",
    "detection_lines.paint_uniform_color([1.0, 0.0, 1.0])\n",
    "\n",
    "o3d.visualization.draw_geometries([mesh_frame, pc, detection_lines])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "7f140cd9-5c75-4804-9ae7-a9af7cc3e0dc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TriangleMesh with 96 points and 132 triangles."
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33a315bb-6b13-4a33-8f31-624c6411c3fd",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f5d0d19-0da4-4f85-aaec-fa8fd9ab88ca",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e50fa2da-bc5d-4815-abc7-e7eb7f59dab2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0354659-882b-49b0-8dd4-db92f23c5fcc",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "6dad1107-a677-4382-a4ec-99408ace3bde",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'hoge' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[52], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mhoge\u001b[49m()\n",
      "\u001b[0;31mNameError\u001b[0m: name 'hoge' is not defined"
     ]
    }
   ],
   "source": [
    "hoge()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cda52f88-21e7-4863-af0a-cda1a2277a26",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}