data_augmentation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# KITTI Dataset augmentation visualization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset structure\n",
    "\n",
    "```\n",
    "📦gta\n",
    "┗ 📂training\n",
    "  ┣ 📂calib\n",
    "  ┃ ┣ 📜<file_id>.txt\n",
    "  ┃ ┗ 📜 ...\n",
    "  ┣ 📂image_2\n",
    "  ┃ ┣ 📜<file_id>.png\n",
    "  ┃ ┗ 📜 ...\n",
    "  ┣ 📂label_2\n",
    "  ┃ ┣ 📜<file_id>.txt\n",
    "  ┃ ┗ 📜 ...\n",
    "  ┗ 📂velodyne\n",
    "    ┣ 📜<file_id>.bin\n",
    "    ┗ 📜 ...\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Notebook initialized with ipy backend.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import mayavi.mlab as mlab\n",
    "import numpy as np\n",
    "from copy import deepcopy\n",
    "import pickle\n",
    "import logging\n",
    "import math\n",
    "\n",
    "import pcdet.datasets.augmentor\n",
    "\n",
    "from pcdet.config import cfg, cfg_from_yaml_file\n",
    "from pcdet.datasets import build_dataloader\n",
    "from pcdet.utils import common_utils\n",
    "from pcdet.models import build_network\n",
    "from tools.visual_utils import visualize_utils as V\n",
    "\n",
    "from PIL import Image, ImageDraw\n",
    "\n",
    "mlab.init_notebook()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate pickles from the kitti data python -m pcdet.datasets.kitti.kitti_dataset create_kitti_infos tools/cfgs/dataset_configs/kitti_dataset.yaml\n",
    "# os.system(\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing the LiDAR point cloud with labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = common_utils.create_logger(log_level=logging.INFO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-12-03 12:56:10,137   INFO  Database filter by min points Car: 302 => 253\n",
      "2021-12-03 12:56:10,138   INFO  Database filter by difficulty Car: 253 => 214\n",
      "2021-12-03 12:56:10,139   INFO  Database filter by difficulty Pedestrian: 46 => 45\n",
      "2021-12-03 12:56:10,139   INFO  Database filter by difficulty Cyclist: 17 => 12\n",
      "2021-12-03 12:56:10,194   INFO  Loading KITTI dataset\n",
      "2021-12-03 12:56:10,202   INFO  Total samples for KITTI dataset: 100\n",
      "2021-12-03 12:56:10,203   INFO  Total number of samples: \t100\n",
      "2021-12-03 12:56:10,204   INFO  Loading samples\n",
      "2021-12-03 12:56:10,291   INFO  Loaded sample index: \t1\n",
      "2021-12-03 12:56:10,346   INFO  Loaded sample index: \t2\n",
      "2021-12-03 12:56:10,393   INFO  Loaded sample index: \t3\n",
      "2021-12-03 12:56:10,431   INFO  Loaded sample index: \t4\n",
      "2021-12-03 12:56:10,482   INFO  Loaded sample index: \t5\n",
      "2021-12-03 12:56:10,534   INFO  Loaded sample index: \t6\n",
      "2021-12-03 12:56:10,597   INFO  Loaded sample index: \t7\n",
      "2021-12-03 12:56:10,655   INFO  Loaded sample index: \t8\n",
      "2021-12-03 12:56:10,720   INFO  Loaded sample index: \t9\n",
      "2021-12-03 12:56:10,769   INFO  Loaded sample index: \t10\n",
      "2021-12-03 12:56:10,803   INFO  Loaded sample index: \t11\n",
      "2021-12-03 12:56:10,834   INFO  Loaded sample index: \t12\n",
      "2021-12-03 12:56:10,881   INFO  Loaded sample index: \t13\n",
      "2021-12-03 12:56:10,922   INFO  Loaded sample index: \t14\n",
      "2021-12-03 12:56:10,947   INFO  Loaded sample index: \t15\n",
      "2021-12-03 12:56:10,993   INFO  Loaded sample index: \t16\n",
      "2021-12-03 12:56:11,036   INFO  Loaded sample index: \t17\n",
      "2021-12-03 12:56:11,067   INFO  Loaded sample index: \t18\n",
      "2021-12-03 12:56:11,118   INFO  Loaded sample index: \t19\n",
      "2021-12-03 12:56:11,162   INFO  Loaded sample index: \t20\n",
      "2021-12-03 12:56:11,214   INFO  Loaded sample index: \t21\n",
      "2021-12-03 12:56:11,259   INFO  Loaded sample index: \t22\n",
      "2021-12-03 12:56:11,308   INFO  Loaded sample index: \t23\n",
      "2021-12-03 12:56:11,373   INFO  Loaded sample index: \t24\n",
      "2021-12-03 12:56:11,425   INFO  Loaded sample index: \t25\n",
      "2021-12-03 12:56:11,477   INFO  Loaded sample index: \t26\n",
      "2021-12-03 12:56:11,505   INFO  Loaded sample index: \t27\n",
      "2021-12-03 12:56:11,561   INFO  Loaded sample index: \t28\n",
      "2021-12-03 12:56:11,603   INFO  Loaded sample index: \t29\n",
      "2021-12-03 12:56:11,633   INFO  Loaded sample index: \t30\n",
      "2021-12-03 12:56:11,663   INFO  Loaded sample index: \t31\n",
      "2021-12-03 12:56:11,746   INFO  Loaded sample index: \t32\n",
      "2021-12-03 12:56:11,789   INFO  Loaded sample index: \t33\n",
      "2021-12-03 12:56:11,831   INFO  Loaded sample index: \t34\n",
      "2021-12-03 12:56:11,887   INFO  Loaded sample index: \t35\n",
      "2021-12-03 12:56:11,949   INFO  Loaded sample index: \t36\n",
      "2021-12-03 12:56:11,981   INFO  Loaded sample index: \t37\n",
      "2021-12-03 12:56:12,012   INFO  Loaded sample index: \t38\n",
      "2021-12-03 12:56:12,044   INFO  Loaded sample index: \t39\n",
      "2021-12-03 12:56:12,089   INFO  Loaded sample index: \t40\n",
      "2021-12-03 12:56:12,137   INFO  Loaded sample index: \t41\n",
      "2021-12-03 12:56:12,175   INFO  Loaded sample index: \t42\n",
      "2021-12-03 12:56:12,198   INFO  Loaded sample index: \t43\n",
      "2021-12-03 12:56:12,234   INFO  Loaded sample index: \t44\n",
      "2021-12-03 12:56:12,294   INFO  Loaded sample index: \t45\n",
      "2021-12-03 12:56:12,349   INFO  Loaded sample index: \t46\n",
      "2021-12-03 12:56:12,421   INFO  Loaded sample index: \t47\n",
      "2021-12-03 12:56:12,470   INFO  Loaded sample index: \t48\n",
      "2021-12-03 12:56:12,512   INFO  Loaded sample index: \t49\n",
      "2021-12-03 12:56:12,562   INFO  Loaded sample index: \t50\n",
      "2021-12-03 12:56:12,598   INFO  Loaded sample index: \t51\n",
      "2021-12-03 12:56:12,642   INFO  Loaded sample index: \t52\n",
      "2021-12-03 12:56:12,690   INFO  Loaded sample index: \t53\n",
      "2021-12-03 12:56:12,758   INFO  Loaded sample index: \t54\n",
      "2021-12-03 12:56:12,797   INFO  Loaded sample index: \t55\n",
      "2021-12-03 12:56:12,858   INFO  Loaded sample index: \t56\n",
      "2021-12-03 12:56:12,903   INFO  Loaded sample index: \t57\n",
      "2021-12-03 12:56:12,942   INFO  Loaded sample index: \t58\n",
      "2021-12-03 12:56:12,975   INFO  Loaded sample index: \t59\n",
      "2021-12-03 12:56:13,022   INFO  Loaded sample index: \t60\n",
      "2021-12-03 12:56:13,086   INFO  Loaded sample index: \t61\n",
      "2021-12-03 12:56:13,156   INFO  Loaded sample index: \t62\n",
      "2021-12-03 12:56:13,210   INFO  Loaded sample index: \t63\n",
      "2021-12-03 12:56:13,276   INFO  Loaded sample index: \t64\n",
      "2021-12-03 12:56:13,334   INFO  Loaded sample index: \t65\n",
      "2021-12-03 12:56:13,370   INFO  Loaded sample index: \t66\n",
      "2021-12-03 12:56:13,406   INFO  Loaded sample index: \t67\n",
      "2021-12-03 12:56:13,431   INFO  Loaded sample index: \t68\n",
      "2021-12-03 12:56:13,454   INFO  Loaded sample index: \t69\n",
      "2021-12-03 12:56:13,477   INFO  Loaded sample index: \t70\n",
      "2021-12-03 12:56:13,496   INFO  Loaded sample index: \t71\n",
      "2021-12-03 12:56:13,520   INFO  Loaded sample index: \t72\n",
      "2021-12-03 12:56:13,566   INFO  Loaded sample index: \t73\n",
      "2021-12-03 12:56:13,613   INFO  Loaded sample index: \t74\n",
      "2021-12-03 12:56:13,645   INFO  Loaded sample index: \t75\n",
      "2021-12-03 12:56:13,726   INFO  Loaded sample index: \t76\n",
      "2021-12-03 12:56:13,785   INFO  Loaded sample index: \t77\n",
      "2021-12-03 12:56:13,817   INFO  Loaded sample index: \t78\n",
      "2021-12-03 12:56:13,854   INFO  Loaded sample index: \t79\n",
      "2021-12-03 12:56:13,892   INFO  Loaded sample index: \t80\n",
      "2021-12-03 12:56:13,919   INFO  Loaded sample index: \t81\n",
      "2021-12-03 12:56:13,987   INFO  Loaded sample index: \t82\n",
      "2021-12-03 12:56:14,027   INFO  Loaded sample index: \t83\n",
      "2021-12-03 12:56:14,072   INFO  Loaded sample index: \t84\n",
      "2021-12-03 12:56:14,138   INFO  Loaded sample index: \t85\n",
      "2021-12-03 12:56:14,211   INFO  Loaded sample index: \t86\n",
      "2021-12-03 12:56:14,277   INFO  Loaded sample index: \t87\n",
      "2021-12-03 12:56:14,323   INFO  Loaded sample index: \t88\n",
      "2021-12-03 12:56:14,366   INFO  Loaded sample index: \t89\n",
      "2021-12-03 12:56:14,414   INFO  Loaded sample index: \t90\n",
      "2021-12-03 12:56:14,452   INFO  Loaded sample index: \t91\n",
      "2021-12-03 12:56:14,486   INFO  Loaded sample index: \t92\n",
      "2021-12-03 12:56:14,521   INFO  Loaded sample index: \t93\n",
      "2021-12-03 12:56:14,557   INFO  Loaded sample index: \t94\n",
      "2021-12-03 12:56:14,582   INFO  Loaded sample index: \t95\n",
      "2021-12-03 12:56:14,608   INFO  Loaded sample index: \t96\n",
      "2021-12-03 12:56:14,637   INFO  Loaded sample index: \t97\n",
      "2021-12-03 12:56:14,676   INFO  Loaded sample index: \t98\n",
      "2021-12-03 12:56:14,713   INFO  Loaded sample index: \t99\n",
      "2021-12-03 12:56:14,745   INFO  Loaded sample index: \t100\n"
     ]
    }
   ],
   "source": [
    "cfg_from_yaml_file('tools/cfgs/dataset_configs/kitti_dataset.yaml', cfg)\n",
    "# cfg_from_yaml_file('tools/cfgs/kitti_models/pointpillar_augs.yaml', cfg)\n",
    "\n",
    "cfg.DATA_PATH = 'data/kitti'\n",
    "\n",
    "train_set, train_loader, train_sampler = build_dataloader(\n",
    "    dataset_cfg=cfg,\n",
    "    class_names=['Car', 'Pedestrian', 'Cyclist'],\n",
    "    batch_size=1,\n",
    "    dist=False,\n",
    "    workers=4,\n",
    "    logger=logger,\n",
    "    training=True,\n",
    "    merge_all_iters_to_one_epoch=False,\n",
    "    total_epochs=0\n",
    ")\n",
    "\n",
    "logger.info(f'Total number of samples: \\t{len(train_set)}')\n",
    "\n",
    "data_dict_list = []\n",
    "logger.info('Loading samples')\n",
    "for idx, data_dict in enumerate(train_set):\n",
    "    logger.info(f'Loaded sample index: \\t{idx + 1}')\n",
    "    data_dict = train_set.collate_batch([data_dict])\n",
    "    data_dict_list.append(data_dict)\n",
    "    #if idx >= 35: break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {},
   "outputs": [],
   "source": [
    "scene = data_dict_list[8]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_pc(data_dict):\n",
    "    currfig = mlab.figure(size=(1280,720))\n",
    "\n",
    "    V.draw_scenes(\n",
    "        points=data_dict['points'][:, 1:], gt_boxes=data_dict['gt_boxes'][0]\n",
    "    )\n",
    "\n",
    "    cam = currfig.scene.camera\n",
    "    cam.zoom(2)\n",
    "    return currfig\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_bev(data_dict):\n",
    "    currfig = mlab.figure(size=(1280,1280))\n",
    "\n",
    "    V.draw_scenes(\n",
    "        points=data_dict['points'][:, 1:], gt_boxes=data_dict['gt_boxes'][0]\n",
    "    )\n",
    "\n",
    "    mlab.view(azimuth=-179, elevation=0, distance=120.0, roll=90.0)\n",
    "    \n",
    "    cam = currfig.scene.camera\n",
    "    cam.zoom(2)\n",
    "\n",
    "    return currfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-11-23 16:35:33,941   INFO  Saved mosaic_0.png\n",
      "2021-11-23 16:35:33,941   INFO  Saved mosaic_0.png\n",
      "2021-11-23 16:35:36,253   INFO  Saved mosaic_1.png\n",
      "2021-11-23 16:35:36,253   INFO  Saved mosaic_1.png\n",
      "2021-11-23 16:35:38,655   INFO  Saved mosaic_2.png\n",
      "2021-11-23 16:35:38,655   INFO  Saved mosaic_2.png\n",
      "2021-11-23 16:35:41,100   INFO  Saved mosaic_3.png\n",
      "2021-11-23 16:35:41,100   INFO  Saved mosaic_3.png\n",
      "2021-11-23 16:35:43,731   INFO  Saved mosaic_4.png\n",
      "2021-11-23 16:35:43,731   INFO  Saved mosaic_4.png\n",
      "2021-11-23 16:35:46,132   INFO  Saved mosaic_5.png\n",
      "2021-11-23 16:35:46,132   INFO  Saved mosaic_5.png\n",
      "2021-11-23 16:35:49,183   INFO  Saved mosaic_6.png\n",
      "2021-11-23 16:35:49,183   INFO  Saved mosaic_6.png\n",
      "2021-11-23 16:35:52,240   INFO  Saved mosaic_7.png\n",
      "2021-11-23 16:35:52,240   INFO  Saved mosaic_7.png\n",
      "2021-11-23 16:35:55,626   INFO  Saved mosaic_8.png\n",
      "2021-11-23 16:35:55,626   INFO  Saved mosaic_8.png\n",
      "2021-11-23 16:35:58,527   INFO  Saved mosaic_9.png\n",
      "2021-11-23 16:35:58,527   INFO  Saved mosaic_9.png\n",
      "2021-11-23 16:36:03,284   INFO  Saved mosaic_10.png\n",
      "2021-11-23 16:36:03,284   INFO  Saved mosaic_10.png\n",
      "2021-11-23 16:36:06,681   INFO  Saved mosaic_11.png\n",
      "2021-11-23 16:36:06,681   INFO  Saved mosaic_11.png\n",
      "2021-11-23 16:36:08,282   INFO  Saved mosaic_12.png\n",
      "2021-11-23 16:36:08,282   INFO  Saved mosaic_12.png\n",
      "2021-11-23 16:36:10,051   INFO  Saved mosaic_13.png\n",
      "2021-11-23 16:36:10,051   INFO  Saved mosaic_13.png\n",
      "2021-11-23 16:36:12,601   INFO  Saved mosaic_14.png\n",
      "2021-11-23 16:36:12,601   INFO  Saved mosaic_14.png\n",
      "2021-11-23 16:36:15,602   INFO  Saved mosaic_15.png\n",
      "2021-11-23 16:36:15,602   INFO  Saved mosaic_15.png\n",
      "2021-11-23 16:36:18,623   INFO  Saved mosaic_16.png\n",
      "2021-11-23 16:36:18,623   INFO  Saved mosaic_16.png\n",
      "2021-11-23 16:36:20,683   INFO  Saved mosaic_17.png\n",
      "2021-11-23 16:36:20,683   INFO  Saved mosaic_17.png\n",
      "2021-11-23 16:36:22,704   INFO  Saved mosaic_18.png\n",
      "2021-11-23 16:36:22,704   INFO  Saved mosaic_18.png\n",
      "2021-11-23 16:36:26,049   INFO  Saved mosaic_19.png\n",
      "2021-11-23 16:36:26,049   INFO  Saved mosaic_19.png\n",
      "2021-11-23 16:36:28,521   INFO  Saved mosaic_20.png\n",
      "2021-11-23 16:36:28,521   INFO  Saved mosaic_20.png\n",
      "2021-11-23 16:36:32,431   INFO  Saved mosaic_21.png\n",
      "2021-11-23 16:36:32,431   INFO  Saved mosaic_21.png\n",
      "2021-11-23 16:36:34,195   INFO  Saved mosaic_22.png\n",
      "2021-11-23 16:36:34,195   INFO  Saved mosaic_22.png\n",
      "2021-11-23 16:36:35,949   INFO  Saved mosaic_23.png\n",
      "2021-11-23 16:36:35,949   INFO  Saved mosaic_23.png\n",
      "2021-11-23 16:36:39,003   INFO  Saved mosaic_24.png\n",
      "2021-11-23 16:36:39,003   INFO  Saved mosaic_24.png\n",
      "2021-11-23 16:36:41,771   INFO  Saved mosaic_25.png\n",
      "2021-11-23 16:36:41,771   INFO  Saved mosaic_25.png\n",
      "2021-11-23 16:36:43,321   INFO  Saved mosaic_26.png\n",
      "2021-11-23 16:36:43,321   INFO  Saved mosaic_26.png\n",
      "2021-11-23 16:36:45,895   INFO  Saved mosaic_27.png\n",
      "2021-11-23 16:36:45,895   INFO  Saved mosaic_27.png\n",
      "2021-11-23 16:36:47,428   INFO  Saved mosaic_28.png\n",
      "2021-11-23 16:36:47,428   INFO  Saved mosaic_28.png\n",
      "2021-11-23 16:36:48,925   INFO  Saved mosaic_29.png\n",
      "2021-11-23 16:36:48,925   INFO  Saved mosaic_29.png\n",
      "2021-11-23 16:36:50,457   INFO  Saved mosaic_30.png\n",
      "2021-11-23 16:36:50,457   INFO  Saved mosaic_30.png\n",
      "2021-11-23 16:36:54,083   INFO  Saved mosaic_31.png\n",
      "2021-11-23 16:36:54,083   INFO  Saved mosaic_31.png\n",
      "2021-11-23 16:36:56,960   INFO  Saved mosaic_32.png\n",
      "2021-11-23 16:36:56,960   INFO  Saved mosaic_32.png\n",
      "2021-11-23 16:36:58,710   INFO  Saved mosaic_33.png\n",
      "2021-11-23 16:36:58,710   INFO  Saved mosaic_33.png\n",
      "2021-11-23 16:37:00,472   INFO  Saved mosaic_34.png\n",
      "2021-11-23 16:37:00,472   INFO  Saved mosaic_34.png\n",
      "2021-11-23 16:37:03,086   INFO  Saved mosaic_35.png\n",
      "2021-11-23 16:37:03,086   INFO  Saved mosaic_35.png\n",
      "2021-11-23 16:37:06,228   INFO  Saved mosaic_36.png\n",
      "2021-11-23 16:37:06,228   INFO  Saved mosaic_36.png\n",
      "2021-11-23 16:37:08,004   INFO  Saved mosaic_37.png\n",
      "2021-11-23 16:37:08,004   INFO  Saved mosaic_37.png\n",
      "2021-11-23 16:37:14,123   INFO  Saved mosaic_38.png\n",
      "2021-11-23 16:37:14,123   INFO  Saved mosaic_38.png\n",
      "2021-11-23 16:37:17,000   INFO  Saved mosaic_39.png\n",
      "2021-11-23 16:37:17,000   INFO  Saved mosaic_39.png\n",
      "2021-11-23 16:37:18,805   INFO  Saved mosaic_40.png\n",
      "2021-11-23 16:37:18,805   INFO  Saved mosaic_40.png\n",
      "2021-11-23 16:37:20,338   INFO  Saved mosaic_41.png\n",
      "2021-11-23 16:37:20,338   INFO  Saved mosaic_41.png\n",
      "2021-11-23 16:37:21,964   INFO  Saved mosaic_42.png\n",
      "2021-11-23 16:37:21,964   INFO  Saved mosaic_42.png\n",
      "2021-11-23 16:37:25,011   INFO  Saved mosaic_43.png\n",
      "2021-11-23 16:37:25,011   INFO  Saved mosaic_43.png\n",
      "2021-11-23 16:37:33,979   INFO  Saved mosaic_44.png\n",
      "2021-11-23 16:37:33,979   INFO  Saved mosaic_44.png\n",
      "2021-11-23 16:37:38,541   INFO  Saved mosaic_45.png\n",
      "2021-11-23 16:37:38,541   INFO  Saved mosaic_45.png\n",
      "2021-11-23 16:37:41,816   INFO  Saved mosaic_46.png\n",
      "2021-11-23 16:37:41,816   INFO  Saved mosaic_46.png\n",
      "2021-11-23 16:37:47,751   INFO  Saved mosaic_47.png\n",
      "2021-11-23 16:37:47,751   INFO  Saved mosaic_47.png\n",
      "2021-11-23 16:37:50,441   INFO  Saved mosaic_48.png\n",
      "2021-11-23 16:37:50,441   INFO  Saved mosaic_48.png\n",
      "2021-11-23 16:37:57,352   INFO  Saved mosaic_49.png\n",
      "2021-11-23 16:37:57,352   INFO  Saved mosaic_49.png\n"
     ]
    }
   ],
   "source": [
    "for idx, scene in enumerate(data_dict_list):\n",
    "    currfig = show_bev(scene)\n",
    "    mlab.savefig('figs/mosaic/mosaic_%d.png' % idx, size=(1280,1280))\n",
    "    logger.info('Saved mosaic_%d.png' % idx)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "collage = Image.new(\"RGBA\", (3000,3000), color=(255,255,255,255))\n",
    "\n",
    "c=0\n",
    "for i in range(0,3000,500):\n",
    "    for j in range(0,3000,500):\n",
    "        file = 'figs/mosaic/mosaic_%d.png' % c\n",
    "        photo = Image.open(file).convert(\"RGBA\")\n",
    "        photo = photo.resize((500,500))        \n",
    "        \n",
    "        collage.paste(photo, (i,j))\n",
    "        c+=1\n",
    "collage.save('figs/mosaic/collage.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'show_bev' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-7-ca53c535e26d>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mcurrfig\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mshow_bev\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata_dict_list\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      2\u001b[0m \u001b[0mmlab\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msavefig\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'3d_renders/original.obj'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0mmlab\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msavefig\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'figs/original.png'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msize\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1280\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m720\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0mcurrfig\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'show_bev' is not defined"
     ]
    }
   ],
   "source": [
    "currfig = show_bev(data_dict_list[2])\n",
    "mlab.savefig('3d_renders/original.obj')\n",
    "mlab.savefig('figs/original.png', size=(1280,720))\n",
    "currfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing the original scene\n",
    "![Original Camera Image](data/kitti/training/image_2/000007.png)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "18ee4df8a37f46e1ae109af72dcec7b6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x02X\\x00\\x00\\x02X\\x08\\x02\\x00\\x00\\x001\\x04\\x0f\\x8b\\x…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "currfig = show_pc(data_dict_list[7])\n",
    "mlab.savefig('3d_renders/aug_pipeline_output.obj')\n",
    "mlab.savefig('figs/aug_pipeline_output.png', size=(1280,720))\n",
    "currfig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Augmentation methods"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Random flip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def random_flip_along_x(gt_boxes, points):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        gt_boxes: (N, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]\n",
    "        points: (M, 3 + C)\n",
    "    Returns:\n",
    "    \"\"\"\n",
    "    enable = True\n",
    "    if enable:\n",
    "        gt_boxes[:, 1] = -gt_boxes[:, 1]\n",
    "        gt_boxes[:, 6] = -gt_boxes[:, 6]\n",
    "        points[:, 1] = -points[:, 1]\n",
    "\n",
    "        if gt_boxes.shape[1] > 8:\n",
    "            gt_boxes[:, 8] = -gt_boxes[:, 8]\n",
    "\n",
    "    return gt_boxes, points\n",
    "\n",
    "def random_flip_along_y(gt_boxes, points):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        gt_boxes: (N, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]\n",
    "        points: (M, 3 + C)\n",
    "    Returns:\n",
    "    \"\"\"\n",
    "    enable = True\n",
    "    if enable:\n",
    "        gt_boxes[:, 0] = -gt_boxes[:, 0]\n",
    "        gt_boxes[:, 6] = -(gt_boxes[:, 6] + np.pi)\n",
    "        points[:, 0] = -points[:, 0]\n",
    "\n",
    "        if gt_boxes.shape[1] > 7:\n",
    "            gt_boxes[:, 7] = -gt_boxes[:, 7]\n",
    "\n",
    "    return gt_boxes, points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 692 µs, sys: 906 µs, total: 1.6 ms\n",
      "Wall time: 914 µs\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "flipped_y = deepcopy(scene)\n",
    "flipped_y['gt_boxes'][0], flipped_y['points'][:,1:] = random_flip_along_y(flipped_y['gt_boxes'][0], flipped_y['points'][:,1:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e702212c71af4e30be6127853bf3f124",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x02X\\x00\\x00\\x02X\\x08\\x02\\x00\\x00\\x001\\x04\\x0f\\x8b\\x…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "currfig = show_pc(flipped_y)\n",
    "mlab.savefig('figs/flipped_y.png', size=(1280,720))\n",
    "mlab.savefig('3d_renders/flipped_y.obj')\n",
    "currfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.83 ms, sys: 469 µs, total: 2.3 ms\n",
      "Wall time: 2.77 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "flipped_x = deepcopy(scene)\n",
    "flipped_x['gt_boxes'][0], flipped_x['points'][:,1:] = random_flip_along_x(flipped_x['gt_boxes'][0], flipped_x['points'][:,1:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1e5093de84884574af404d38c9e5aa4f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x02X\\x00\\x00\\x02X\\x08\\x02\\x00\\x00\\x001\\x04\\x0f\\x8b\\x…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "currfig = show_pc(flipped_x)\n",
    "mlab.savefig('figs/flipped_x.png', size=(1280,720))\n",
    "mlab.savefig('3d_renders/flipped_x.obj')\n",
    "currfig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Global rotation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def global_rotation(gt_boxes, points, rot_range):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        gt_boxes: (N, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]\n",
    "        points: (M, 3 + C),\n",
    "        rot_range: [min, max]\n",
    "    Returns:\n",
    "    \"\"\"\n",
    "    noise_rotation = np.random.uniform(rot_range[0], rot_range[1])\n",
    "    points = common_utils.rotate_points_along_z(points[np.newaxis, :, :], np.array([noise_rotation]))[0]\n",
    "    gt_boxes[:, 0:3] = common_utils.rotate_points_along_z(gt_boxes[np.newaxis, :, 0:3], np.array([noise_rotation]))[0]\n",
    "    gt_boxes[:, 6] += noise_rotation\n",
    "    if gt_boxes.shape[1] > 8:\n",
    "        gt_boxes[:, 7:9] = common_utils.rotate_points_along_z(\n",
    "            np.hstack((gt_boxes[:, 7:9], np.zeros((gt_boxes.shape[0], 1))))[np.newaxis, :, :],\n",
    "            np.array([noise_rotation])\n",
    "        )[0][:, 0:2]\n",
    "\n",
    "    return gt_boxes, points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 25.1 ms, sys: 1.34 ms, total: 26.4 ms\n",
      "Wall time: 14.6 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "rotated = deepcopy(scene)\n",
    "rotated['gt_boxes'][0], rotated['points'][:,1:] = global_rotation(rotated['gt_boxes'][0], rotated['points'][:,1:], [-0.78539816, 0.78539816])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f92ed76a7ea74ff29eefbb724bfa2092",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x02X\\x00\\x00\\x02X\\x08\\x02\\x00\\x00\\x001\\x04\\x0f\\x8b\\x…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "currfig = show_pc(rotated)\n",
    "mlab.savefig('figs/rotated_global.png', size=(1280,720))\n",
    "mlab.savefig('3d_renders/rotated_global.obj')\n",
    "currfig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Global scaling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def global_scaling(gt_boxes, points, scale_range):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading]\n",
    "        points: (M, 3 + C),\n",
    "        scale_range: [min, max]\n",
    "    Returns:\n",
    "    \"\"\"\n",
    "    if scale_range[1] - scale_range[0] < 1e-3:\n",
    "        return gt_boxes, points\n",
    "    noise_scale = np.random.uniform(scale_range[0], scale_range[1])\n",
    "    points[:, :3] *= noise_scale\n",
    "    gt_boxes[:, :6] *= noise_scale\n",
    "    return gt_boxes, points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.09 ms, sys: 1.53 ms, total: 2.62 ms\n",
      "Wall time: 1.57 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "scaled = deepcopy(scene)\n",
    "scaled['gt_boxes'][0], scaled['points'][:,1:] = global_scaling(scaled['gt_boxes'][0], scaled['points'][:, 1:], [0.4, 0.5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "87b7b22a903141358e639ff35d5e2409",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x02X\\x00\\x00\\x02X\\x08\\x02\\x00\\x00\\x001\\x04\\x0f\\x8b\\x…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "currfig = show_pc(scaled)\n",
    "mlab.savefig('figs/scaled_global.png', size=(1280,720))\n",
    "mlab.savefig('3d_renders/scaled_global.obj')\n",
    "currfig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Global translation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def random_translation_along_x(gt_boxes, points, offset_range):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]\n",
    "        points: (M, 3 + C),\n",
    "        offset_range: [min max]]\n",
    "    Returns:\n",
    "    \"\"\"\n",
    "    offset = np.random.uniform(offset_range[0], offset_range[1])\n",
    "\n",
    "    points[:, 0] += offset\n",
    "    gt_boxes[:, 0] += offset\n",
    "\n",
    "    if gt_boxes.shape[1] > 7:\n",
    "        gt_boxes[:, 7] += offset\n",
    "\n",
    "    return gt_boxes, points\n",
    "\n",
    "def random_translation_along_y(gt_boxes, points, offset_range):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]\n",
    "        points: (M, 3 + C),\n",
    "        offset_range: [min max]]\n",
    "    Returns:\n",
    "    \"\"\"\n",
    "    offset = np.random.uniform(offset_range[0], offset_range[1])\n",
    "\n",
    "    points[:, 1] += offset\n",
    "    gt_boxes[:, 1] += offset\n",
    "\n",
    "    if gt_boxes.shape[1] > 8:\n",
    "        gt_boxes[:, 8] += offset\n",
    "\n",
    "    return gt_boxes, points\n",
    "\n",
    "def random_translation_along_z(gt_boxes, points, offset_range):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]\n",
    "        points: (M, 3 + C),\n",
    "        offset_range: [min max]]\n",
    "    Returns:\n",
    "    \"\"\"\n",
    "    offset = np.random.uniform(offset_range[0], offset_range[1])\n",
    "\n",
    "    points[:, 2] += offset\n",
    "    gt_boxes[:, 2] += offset\n",
    "\n",
    "    return gt_boxes, points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.73 ms, sys: 0 ns, total: 1.73 ms\n",
      "Wall time: 1.02 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "translation_x = deepcopy(scene)\n",
    "translation_x['gt_boxes'][0], translation_x['points'][:,1:] = random_translation_along_x(translation_x['gt_boxes'][0], translation_x['points'][:, 1:], [15, 20])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cb0e0009db9c494189bdfa97b1d5e0ce",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x02X\\x00\\x00\\x02X\\x08\\x02\\x00\\x00\\x001\\x04\\x0f\\x8b\\x…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "currfig = show_pc(translation_x)\n",
    "mlab.savefig('figs/translation_global_x.png', size=(1280,720))\n",
    "mlab.savefig('3d_renders/translation_global_x.obj')\n",
    "currfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.53 ms, sys: 1.04 ms, total: 2.57 ms\n",
      "Wall time: 1.49 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "translation_y = deepcopy(scene)\n",
    "translation_y['gt_boxes'][0], translation_y['points'][:,1:] = random_translation_along_y(translation_y['gt_boxes'][0], translation_y['points'][:, 1:], [15, 20])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aa3941093ea44463a23b906c58459038",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x02X\\x00\\x00\\x02X\\x08\\x02\\x00\\x00\\x001\\x04\\x0f\\x8b\\x…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "currfig = show_pc(translation_y)\n",
    "mlab.savefig('figs/translation_global_y.png', size=(1280,720))\n",
    "mlab.savefig('3d_renders/translation_global_y.obj')\n",
    "currfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.38 ms, sys: 858 µs, total: 2.24 ms\n",
      "Wall time: 1.31 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "translation_z = deepcopy(scene)\n",
    "translation_z['gt_boxes'][0], translation_z['points'][:,1:] = random_translation_along_z(translation_z['gt_boxes'][0], translation_z['points'][:, 1:], [15, 20])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c6e0e3006cbc4350a8f760e10ef14151",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x02X\\x00\\x00\\x02X\\x08\\x02\\x00\\x00\\x001\\x04\\x0f\\x8b\\x…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "currfig = show_pc(translation_z)\n",
    "mlab.savefig('figs/translation_global_z.png', size=(1280,720))\n",
    "mlab.savefig('3d_renders/translation_global_z.obj')\n",
    "currfig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Global frustum dropout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def global_frustum_dropout_top(gt_boxes, points, intensity_range):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]],\n",
    "        points: (M, 3 + C),\n",
    "        intensity: [min, max]\n",
    "    Returns:\n",
    "    \"\"\"\n",
    "    intensity = np.random.uniform(intensity_range[0], intensity_range[1])\n",
    "\n",
    "    threshold = np.max(points[:, 2]) - intensity * (np.max(points[:, 2]) - np.min(points[:, 2]))\n",
    "    points = points[points[:,2] < threshold]\n",
    "    gt_boxes = gt_boxes[gt_boxes[:,2] < threshold]\n",
    "\n",
    "    return gt_boxes, points\n",
    "\n",
    "def global_frustum_dropout_bottom(gt_boxes, points, intensity_range):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]],\n",
    "        points: (M, 3 + C),\n",
    "        intensity: [min, max]\n",
    "    Returns:\n",
    "    \"\"\"\n",
    "    intensity = np.random.uniform(intensity_range[0], intensity_range[1])\n",
    "\n",
    "    threshold = np.min(points[:, 2]) + intensity * (np.max(points[:, 2]) - np.min(points[:, 2]))\n",
    "    points = points[points[:,2] > threshold]\n",
    "    gt_boxes = gt_boxes[gt_boxes[:,2] > threshold]\n",
    "\n",
    "    return gt_boxes, points\n",
    "\n",
    "def global_frustum_dropout_left(gt_boxes, points, intensity_range):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]],\n",
    "        points: (M, 3 + C),\n",
    "        intensity: [min, max]\n",
    "    Returns:\n",
    "    \"\"\"\n",
    "    intensity = np.random.uniform(intensity_range[0], intensity_range[1])\n",
    "\n",
    "    threshold = np.max(points[:, 1]) - intensity * (np.max(points[:, 1]) - np.min(points[:, 1]))\n",
    "    points = points[points[:,1] < threshold]\n",
    "    gt_boxes = gt_boxes[gt_boxes[:,1] < threshold]\n",
    "\n",
    "    return gt_boxes, points\n",
    "\n",
    "def global_frustum_dropout_right(gt_boxes, points, intensity_range):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]],\n",
    "        points: (M, 3 + C),\n",
    "        intensity: [min, max]\n",
    "    Returns:\n",
    "    \"\"\"\n",
    "    intensity = np.random.uniform(intensity_range[0], intensity_range[1])\n",
    "\n",
    "    threshold = np.min(points[:, 1]) + intensity * (np.max(points[:, 1]) - np.min(points[:, 1]))\n",
    "    points = points[points[:,1] > threshold]\n",
    "    gt_boxes = gt_boxes[gt_boxes[:,1] > threshold]\n",
    "\n",
    "    return gt_boxes, points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 0 ns, sys: 2.68 ms, total: 2.68 ms\n",
      "Wall time: 1.9 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "frustum_dropout = deepcopy(scene)\n",
    "frustum_dropout['gt_boxes'], frustum_dropout['points'] = global_frustum_dropout_top(frustum_dropout['gt_boxes'][0], frustum_dropout['points'][:, 1:], [0.4, 0.5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "720cc95fbf3341888654bd02c1b9c9d9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x02X\\x00\\x00\\x02X\\x08\\x02\\x00\\x00\\x001\\x04\\x0f\\x8b\\x…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "V.draw_scenes(\n",
    "        points=frustum_dropout['points'], gt_boxes=frustum_dropout['gt_boxes']\n",
    "    )\n",
    "\n",
    "mlab.savefig('figs/frustum_dropout_top.png')\n",
    "mlab.savefig('3d_renders/frustum_dropout_top.obj')\n",
    "currfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.5 ms, sys: 652 µs, total: 2.15 ms\n",
      "Wall time: 1.33 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "frustum_dropout = deepcopy(scene)\n",
    "frustum_dropout['gt_boxes'], frustum_dropout['points'] = global_frustum_dropout_bottom(frustum_dropout['gt_boxes'][0], frustum_dropout['points'][:, 1:], [0.4, 0.5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "739052468ea24857b8ad1de504fc3da8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x02X\\x00\\x00\\x02X\\x08\\x02\\x00\\x00\\x001\\x04\\x0f\\x8b\\x…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "V.draw_scenes(\n",
    "        points=frustum_dropout['points'], gt_boxes=frustum_dropout['gt_boxes']\n",
    "    )\n",
    "\n",
    "mlab.savefig('figs/frustum_dropout_bottom.png')\n",
    "mlab.savefig('3d_renders/frustum_dropout_bottom.obj')\n",
    "currfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 207 µs, sys: 2.05 ms, total: 2.26 ms\n",
      "Wall time: 1.51 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "frustum_dropout = deepcopy(scene)\n",
    "frustum_dropout['gt_boxes'], frustum_dropout['points'] = global_frustum_dropout_left(frustum_dropout['gt_boxes'][0], frustum_dropout['points'][:, 1:], [0.4, 0.5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "34c4ea0cffcb41538b7831de3a756caf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x02X\\x00\\x00\\x02X\\x08\\x02\\x00\\x00\\x001\\x04\\x0f\\x8b\\x…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "V.draw_scenes(\n",
    "        points=frustum_dropout['points'], gt_boxes=frustum_dropout['gt_boxes']\n",
    "    )\n",
    "\n",
    "mlab.savefig('figs/frustum_dropout_left.png')\n",
    "mlab.savefig('figs/frustum_dropout_bottom.obj')\n",
    "currfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 915 µs, sys: 2.46 ms, total: 3.38 ms\n",
      "Wall time: 2.34 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "frustum_dropout = deepcopy(scene)\n",
    "frustum_dropout['gt_boxes'], frustum_dropout['points'] = global_frustum_dropout_right(frustum_dropout['gt_boxes'][0], frustum_dropout['points'][:, 1:], [0.4, 0.5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "71c8cbe62a1340a4a3bccf1aea4f4137",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x02X\\x00\\x00\\x02X\\x08\\x02\\x00\\x00\\x001\\x04\\x0f\\x8b\\x…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "V.draw_scenes(\n",
    "        points=frustum_dropout['points'], gt_boxes=frustum_dropout['gt_boxes']\n",
    "    )\n",
    "\n",
    "mlab.savefig('figs/frustum_dropout_right.png')\n",
    "mlab.savefig('figs/frustum_dropout_right.obj')\n",
    "currfig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Aux function to find which points are inside a bounding box"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_points_in_box(points, gt_box):\n",
    "    x, y, z = points[:,0], points[:,1], points[:,2]\n",
    "    cx, cy, cz = gt_box[0], gt_box[1], gt_box[2]\n",
    "    dx, dy, dz, rz = gt_box[3], gt_box[4], gt_box[5], gt_box[6]\n",
    "    shift_x, shift_y, shift_z = x - cx, y - cy, z - cz\n",
    "\n",
    "    MARGIN = 1e-2\n",
    "    cosa, sina = math.cos(-rz), math.sin(-rz)\n",
    "    local_x = shift_x * cosa + shift_y * (-sina)\n",
    "    local_y = shift_x * sina + shift_y * cosa\n",
    "\n",
    "    mask = np.logical_and(abs(shift_z) <= dz / 2.0, \\\n",
    "             np.logical_and(abs(local_x) <= dx / 2.0 + MARGIN, \\\n",
    "                 abs(local_y) <= dy / 2.0 + MARGIN ))\n",
    "\n",
    "    points = points[mask]\n",
    "\n",
    "    return points, mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 136,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.41 ms, sys: 361 µs, total: 2.77 ms\n",
      "Wall time: 4.81 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "points_to_filter = deepcopy(scene)\n",
    "for box in points_to_filter['gt_boxes'][0]:\n",
    "    filtered_points, mask = get_points_in_box(points_to_filter['points'][:,1:], points_to_filter['gt_boxes'][0][1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 137,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "V.draw_scenes(\n",
    "        points=points_to_filter['points'][:,1:][mask], gt_boxes=points_to_filter['gt_boxes'][0,1:]\n",
    "    )\n",
    "\n",
    "mlab.savefig('figs/points_in_box.png')\n",
    "mlab.savefig('3d_renders/points_in_box.obj')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[46.039,  7.479, -1.012,  0.9  ],\n",
       "       [46.826,  8.116, -1.317,  0.   ],\n",
       "       [46.142,  7.547, -0.775,  0.   ],\n",
       "       [46.065,  7.187, -1.011,  0.16 ],\n",
       "       [49.276,  6.82 , -1.093,  0.   ],\n",
       "       [46.023,  7.625, -1.012,  0.89 ],\n",
       "       [46.088,  6.894, -1.011,  0.12 ],\n",
       "       [46.098,  8.159, -1.016,  0.   ],\n",
       "       [46.773,  6.772, -1.028,  0.   ],\n",
       "       [46.283,  7.944, -0.78 ,  0.12 ],\n",
       "       [46.054,  7.333, -1.011,  0.9  ],\n",
       "       [46.125,  7.099, -0.773,  0.   ],\n",
       "       [46.043,  8.001, -1.014,  0.   ],\n",
       "       [46.363,  6.912, -0.778,  0.25 ],\n",
       "       [46.01 ,  7.846, -1.012,  0.54 ],\n",
       "       [46.354,  6.985, -0.778,  0.53 ],\n",
       "       [46.224,  8.083, -0.779,  0.   ],\n",
       "       [48.758,  7.272, -1.373,  0.   ],\n",
       "       [46.034,  7.034, -1.01 ,  0.25 ],\n",
       "       [46.169,  7.254, -0.775,  0.   ]], dtype=float32)"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "filtered_points"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Local translation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {},
   "outputs": [],
   "source": [
    "def local_translation_along_x(gt_boxes, points, offset_range):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]\n",
    "        points: (M, 3 + C),\n",
    "        offset_range: [min max]]\n",
    "    Returns:\n",
    "    \"\"\"\n",
    "    for idx, box in enumerate(gt_boxes):\n",
    "        offset = np.random.uniform(offset_range[0], offset_range[1])\n",
    "        points_in_box, mask = get_points_in_box(points, box)\n",
    "        points[mask, 0] += offset\n",
    "\n",
    "        gt_boxes[idx, 0] += offset\n",
    "\n",
    "        if gt_boxes.shape[1] > 7:\n",
    "            gt_boxes[idx, 7] += offset\n",
    "\n",
    "    return gt_boxes, points\n",
    "\n",
    "\n",
    "def local_translation_along_y(gt_boxes, points, offset_range):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]\n",
    "        points: (M, 3 + C),\n",
    "        offset_range: [min max]]\n",
    "    Returns:\n",
    "    \"\"\"\n",
    "    for idx, box in enumerate(gt_boxes):\n",
    "        offset = np.random.uniform(offset_range[0], offset_range[1])\n",
    "        points_in_box, mask = get_points_in_box(points, box)\n",
    "        points[mask, 1] += offset\n",
    "\n",
    "        gt_boxes[idx, 1] += offset\n",
    "\n",
    "        if gt_boxes.shape[1] > 8:\n",
    "            gt_boxes[idx, 8] += offset\n",
    "\n",
    "    return gt_boxes, points\n",
    "\n",
    "\n",
    "def local_translation_along_z(gt_boxes, points, offset_range):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]\n",
    "        points: (M, 3 + C),\n",
    "        offset_range: [min max]]\n",
    "    Returns:\n",
    "    \"\"\"\n",
    "    for idx, box in enumerate(gt_boxes):\n",
    "        offset = np.random.uniform(offset_range[0], offset_range[1])\n",
    "        points_in_box, mask = get_points_in_box(points, box)\n",
    "        points[mask, 2] += offset\n",
    "\n",
    "        gt_boxes[idx, 2] += offset\n",
    "\n",
    "    return gt_boxes, points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.11 ms, sys: 2.21 ms, total: 4.32 ms\n",
      "Wall time: 3.31 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "local_translation_x = deepcopy(scene)\n",
    "local_translation_x['gt_boxes'][0], local_translation_x['points'][:,1:] = local_translation_along_x(local_translation_x['gt_boxes'][0], local_translation_x['points'][:, 1:], [-5, 5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f1afa0d88504484aa4e765644bb09738",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x02X\\x00\\x00\\x02X\\x08\\x02\\x00\\x00\\x001\\x04\\x0f\\x8b\\x…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "currfig = show_pc(local_translation_x)\n",
    "mlab.savefig('figs/translation_local_x.png', size=(1280,720))\n",
    "mlab.savefig('3d_renders/translation_local_x.obj')\n",
    "currfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 0 ns, sys: 4.3 ms, total: 4.3 ms\n",
      "Wall time: 3.39 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "local_translation_y = deepcopy(scene)\n",
    "local_translation_y['gt_boxes'][0], local_translation_y['points'][:,1:] = local_translation_along_y(local_translation_y['gt_boxes'][0], local_translation_y['points'][:, 1:], [-5, 10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b538429a2346434da74d86f33cfce00a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x02X\\x00\\x00\\x02X\\x08\\x02\\x00\\x00\\x001\\x04\\x0f\\x8b\\x…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "currfig = show_pc(local_translation_y)\n",
    "mlab.savefig('figs/translation_local_y.png', size=(1280,720))\n",
    "mlab.savefig('3d_renders/translation_local_y.obj')\n",
    "currfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 0 ns, sys: 4.9 ms, total: 4.9 ms\n",
      "Wall time: 5.88 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "local_translation_z = deepcopy(scene)\n",
    "local_translation_z['gt_boxes'][0], local_translation_z['points'][:,1:] = local_translation_along_z(local_translation_z['gt_boxes'][0], local_translation_z['points'][:, 1:], [-5, 5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bbd7186eec594921ab29b86ec2e21651",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x02X\\x00\\x00\\x02X\\x08\\x02\\x00\\x00\\x001\\x04\\x0f\\x8b\\x…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "currfig = show_pc(local_translation_z)\n",
    "mlab.savefig('figs/translation_local_z.png', size=(1280,720))\n",
    "mlab.savefig('3d_renders/translation_local_z.obj')\n",
    "currfig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Local scaling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "metadata": {},
   "outputs": [],
   "source": [
    "def local_scaling(gt_boxes, points, scale_range):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading]\n",
    "        points: (M, 3 + C),\n",
    "        scale_range: [min, max]\n",
    "    Returns:\n",
    "    \"\"\"\n",
    "    if scale_range[1] - scale_range[0] < 1e-3:\n",
    "        return gt_boxes, points\n",
    "\n",
    "    for idx, box in enumerate(gt_boxes):\n",
    "        noise_scale = np.random.uniform(scale_range[0], scale_range[1])\n",
    "        points_in_box, mask = get_points_in_box(points, box)\n",
    "        \n",
    "        # tranlation to axis center\n",
    "        points[mask, 0] -= box[0]\n",
    "        points[mask, 1] -= box[1]\n",
    "        points[mask, 2] -= box[2]\n",
    "\n",
    "        # apply scaling\n",
    "        points[mask, :3] *= noise_scale\n",
    "\n",
    "        # tranlation back to original position\n",
    "        points[mask, 0] += box[0]\n",
    "        points[mask, 1] += box[1]\n",
    "        points[mask, 2] += box[2]\n",
    "\n",
    "        gt_boxes[idx, 3:6] *= noise_scale\n",
    "    return gt_boxes, points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 116,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 7.9 ms, sys: 0 ns, total: 7.9 ms\n",
      "Wall time: 7.09 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "scaled_local = deepcopy(scene)\n",
    "scaled_local['gt_boxes'][0], scaled_local['points'][:,1:] = local_scaling(scaled_local['gt_boxes'][0], scaled_local['points'][:, 1:], [0.4, 0.5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e5d0a9ccb6e44b51be5414254b84d116",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x02X\\x00\\x00\\x02X\\x08\\x02\\x00\\x00\\x001\\x04\\x0f\\x8b\\x…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "currfig = show_pc(scaled_local)\n",
    "mlab.savefig('figs/scaled_local.png', size=(1280,720))\n",
    "mlab.savefig('3d_renders/scaled_local.obj')\n",
    "currfig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Local rotation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "metadata": {},
   "outputs": [],
   "source": [
    "def local_rotation(gt_boxes, points, rot_range):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]\n",
    "        points: (M, 3 + C),\n",
    "        rot_range: [min, max]\n",
    "    Returns:\n",
    "    \"\"\"\n",
    "    \n",
    "    for idx, box in enumerate(gt_boxes):\n",
    "        noise_rotation = np.random.uniform(rot_range[0], rot_range[1])\n",
    "        points_in_box, mask = get_points_in_box(points, box)\n",
    "        \n",
    "        centroid_x = box[0]\n",
    "        centroid_y = box[1]\n",
    "        centroid_z = box[2]\n",
    "\n",
    "        # tranlation to axis center\n",
    "        points[mask, 0] -= centroid_x\n",
    "        points[mask, 1] -= centroid_y\n",
    "        points[mask, 2] -= centroid_z\n",
    "        box[0] -= centroid_x\n",
    "        box[1] -= centroid_y\n",
    "        box[2] -= centroid_z\n",
    "\n",
    "        # apply rotation\n",
    "        points[mask, :] = common_utils.rotate_points_along_z(points[np.newaxis, mask, :], np.array([noise_rotation]))[0]\n",
    "        box[0:3] = common_utils.rotate_points_along_z(box[np.newaxis, np.newaxis, 0:3], np.array([noise_rotation]))[0][0]\n",
    "\n",
    "        # tranlation back to original position\n",
    "        points[mask, 0] += centroid_x\n",
    "        points[mask, 1] += centroid_y\n",
    "        points[mask, 2] += centroid_z\n",
    "        box[0] += centroid_x\n",
    "        box[1] += centroid_y\n",
    "        box[2] += centroid_z\n",
    "\n",
    "        gt_boxes[idx, 6] += noise_rotation\n",
    "        if gt_boxes.shape[1] > 8:\n",
    "            gt_boxes[idx, 7:9] = common_utils.rotate_points_along_z(\n",
    "                np.hstack((gt_boxes[idx, 7:9], np.zeros((gt_boxes.shape[0], 1))))[np.newaxis, :, :],\n",
    "                np.array([noise_rotation])\n",
    "            )[0][:, 0:2]\n",
    "\n",
    "    return gt_boxes, points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 18 ms, sys: 0 ns, total: 18 ms\n",
      "Wall time: 16.3 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "rotated_local = deepcopy(scene)\n",
    "rotated_local['gt_boxes'][0], rotated_local['points'][:,1:] = local_rotation(rotated_local['gt_boxes'][0],rotated_local['points'][:, 1:], [-0.78539815, 0.78539816])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9d8c6dc4ae904ed38317f0b53f912e9d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x02X\\x00\\x00\\x02X\\x08\\x02\\x00\\x00\\x001\\x04\\x0f\\x8b\\x…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "currfig = show_pc(rotated_local)\n",
    "mlab.savefig('figs/rotated_local.png', size=(1280,720))\n",
    "mlab.savefig('3d_renders/rotated_local.obj')\n",
    "currfig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Local frustum dropout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 181,
   "metadata": {},
   "outputs": [],
   "source": [
    "def local_frustum_dropout_top(gt_boxes, points, intensity_range):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]],\n",
    "        points: (M, 3 + C),\n",
    "        intensity: [min, max]\n",
    "    Returns:\n",
    "    \"\"\"\n",
    "    for idx, box in enumerate(gt_boxes):\n",
    "        x, y, z, dx, dy, dz = box[0], box[1], box[2], box[3], box[4], box[5]\n",
    "\n",
    "        intensity = np.random.uniform(intensity_range[0], intensity_range[1])\n",
    "        points_in_box, mask = get_points_in_box(points, box)\n",
    "        threshold = (z + dz/2) - intensity * dz\n",
    "\n",
    "        points = points[np.logical_not(np.logical_and(mask, points[:,2] >= threshold))]\n",
    "\n",
    "    return gt_boxes, points\n",
    "\n",
    "def local_frustum_dropout_bottom(gt_boxes, points, intensity_range):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]],\n",
    "        points: (M, 3 + C),\n",
    "        intensity: [min, max]\n",
    "    Returns:\n",
    "    \"\"\"\n",
    "    for idx, box in enumerate(gt_boxes):\n",
    "        x, y, z, dx, dy, dz = box[0], box[1], box[2], box[3], box[4], box[5]\n",
    "\n",
    "        intensity = np.random.uniform(intensity_range[0], intensity_range[1])\n",
    "        points_in_box, mask = get_points_in_box(points, box)\n",
    "        threshold = (z - dz/2) + intensity * dz\n",
    "\n",
    "        points = points[np.logical_not(np.logical_and(mask, points[:,2] <= threshold))]\n",
    "\n",
    "    return gt_boxes, points\n",
    "\n",
    "def local_frustum_dropout_left(gt_boxes, points, intensity_range):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]],\n",
    "        points: (M, 3 + C),\n",
    "        intensity: [min, max]\n",
    "    Returns:\n",
    "    \"\"\"\n",
    "    for idx, box in enumerate(gt_boxes):\n",
    "        x, y, z, dx, dy, dz = box[0], box[1], box[2], box[3], box[4], box[5]\n",
    "\n",
    "        intensity = np.random.uniform(intensity_range[0], intensity_range[1])\n",
    "        points_in_box, mask = get_points_in_box(points, box)\n",
    "        threshold = (y + dy/2) - intensity * dy\n",
    "\n",
    "        points = points[np.logical_not(np.logical_and(mask, points[:,1] >= threshold))]\n",
    "\n",
    "    return gt_boxes, points\n",
    "\n",
    "def local_frustum_dropout_right(gt_boxes, points, intensity_range):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]],\n",
    "        points: (M, 3 + C),\n",
    "        intensity: [min, max]\n",
    "    Returns:\n",
    "    \"\"\"\n",
    "    for idx, box in enumerate(gt_boxes):\n",
    "        x, y, z, dx, dy, dz = box[0], box[1], box[2], box[3], box[4], box[5]\n",
    "\n",
    "        intensity = np.random.uniform(intensity_range[0], intensity_range[1])\n",
    "        points_in_box, mask = get_points_in_box(points, box)\n",
    "        threshold = (y - dy/2) + intensity * dy\n",
    "\n",
    "        points = points[np.logical_not(np.logical_and(mask, points[:,1] <= threshold))]\n",
    "\n",
    "    return gt_boxes, points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 4.3 ms, sys: 2.94 ms, total: 7.24 ms\n",
      "Wall time: 6.18 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "frustum_dropout = deepcopy(scene)\n",
    "frustum_dropout['gt_boxes'], frustum_dropout['points'] = local_frustum_dropout_top(frustum_dropout['gt_boxes'][0], frustum_dropout['points'][:, 1:], [0.5, 0.5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 177,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7929409b4ee441e8812957f3384d0c69",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x02X\\x00\\x00\\x02X\\x08\\x02\\x00\\x00\\x001\\x04\\x0f\\x8b\\x…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "currfig = mlab.figure(size=(1280,720))\n",
    "\n",
    "V.draw_scenes(\n",
    "        points=frustum_dropout['points'], gt_boxes=frustum_dropout['gt_boxes']\n",
    "    )\n",
    "\n",
    "mlab.savefig('figs/local_frustum_dropout_top.png')\n",
    "mlab.savefig('3d_renders/local_frustum_dropout_top.obj')\n",
    "currfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 174,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 5.07 ms, sys: 143 µs, total: 5.21 ms\n",
      "Wall time: 4.29 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "frustum_dropout = deepcopy(scene)\n",
    "frustum_dropout['gt_boxes'], frustum_dropout['points'] = local_frustum_dropout_bottom(frustum_dropout['gt_boxes'][0], frustum_dropout['points'][:, 1:], [0.3, 0.4])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 175,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a793c888f81b4a8a84630e85482cafb2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x02X\\x00\\x00\\x02X\\x08\\x02\\x00\\x00\\x001\\x04\\x0f\\x8b\\x…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "V.draw_scenes(\n",
    "        points=frustum_dropout['points'], gt_boxes=frustum_dropout['gt_boxes']\n",
    "    )\n",
    "\n",
    "mlab.savefig('figs/local_frustum_dropout_bottom.png')\n",
    "mlab.savefig('3d_renders/local_frustum_dropout_bottom.obj')\n",
    "currfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 179,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 8.67 ms, sys: 77 µs, total: 8.75 ms\n",
      "Wall time: 7.42 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "frustum_dropout = deepcopy(scene)\n",
    "frustum_dropout['gt_boxes'], frustum_dropout['points'] = local_frustum_dropout_left(frustum_dropout['gt_boxes'][0], frustum_dropout['points'][:, 1:], [0.3, 0.4])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 180,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c0086b860b214f2e9618e81b7e2880fe",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x02X\\x00\\x00\\x02X\\x08\\x02\\x00\\x00\\x001\\x04\\x0f\\x8b\\x…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "V.draw_scenes(\n",
    "        points=frustum_dropout['points'], gt_boxes=frustum_dropout['gt_boxes']\n",
    "    )\n",
    "\n",
    "mlab.savefig('figs/local_frustum_dropout_left.png')\n",
    "mlab.savefig('3d_renders/local_frustum_dropout_left.obj')\n",
    "currfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 182,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 6 ms, sys: 255 µs, total: 6.25 ms\n",
      "Wall time: 5.15 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "frustum_dropout = deepcopy(scene)\n",
    "frustum_dropout['gt_boxes'], frustum_dropout['points'] = local_frustum_dropout_right(frustum_dropout['gt_boxes'][0], frustum_dropout['points'][:, 1:], [0.3, 0.4])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 183,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f39c369dcdad43419beb218b6d418808",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x02X\\x00\\x00\\x02X\\x08\\x02\\x00\\x00\\x001\\x04\\x0f\\x8b\\x…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "V.draw_scenes(\n",
    "        points=frustum_dropout['points'], gt_boxes=frustum_dropout['gt_boxes']\n",
    "    )\n",
    "\n",
    "mlab.savefig('figs/local_frustum_dropout_right.png')\n",
    "mlab.savefig('3d_renders/local_frustum_dropout_right.obj')\n",
    "currfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "e7370f93d1d0cde622a1f8e1c04877d8463912d04d973331ad4851f04de6915a"
  },
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit",
   "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.8.10"
  },
  "metadata": {
   "interpreter": {
    "hash": "e7370f93d1d0cde622a1f8e1c04877d8463912d04d973331ad4851f04de6915a"
   }
  },
  "orig_nbformat": 2
 },
 "nbformat": 4,
 "nbformat_minor": 2
}