data_preprocessing_final.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "======\n",
      "Loading NuScenes tables for version v1.0-trainval...\n",
      "23 category,\n",
      "8 attribute,\n",
      "4 visibility,\n",
      "64386 instance,\n",
      "12 sensor,\n",
      "10200 calibrated_sensor,\n",
      "2631083 ego_pose,\n",
      "68 log,\n",
      "850 scene,\n",
      "34149 sample,\n",
      "2631083 sample_data,\n",
      "1166187 sample_annotation,\n",
      "4 map,\n",
      "Done loading in 44.0 seconds.\n",
      "======\n",
      "Reverse indexing ...\n",
      "Done reverse indexing in 15.2 seconds.\n",
      "======\n"
     ]
    }
   ],
   "source": [
    "%matplotlib inline\n",
    "from nuscenes.nuscenes import NuScenes\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import copy\n",
    "import PIL\n",
    "import os\n",
    "from skimage import color\n",
    "from nuscenes.utils.geometry_utils import view_points, box_in_image, BoxVisibility, transform_matrix\n",
    "from nuscenes.utils.data_classes import LidarPointCloud\n",
    "from pyquaternion import Quaternion\n",
    "#nusc = NuScenes(version='v1.0-mini', dataroot='data/sets/nuscenes', verbose=True)\n",
    "nusc = NuScenes(version='v1.0-trainval', dataroot='full_data', verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os.path as osp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# def project_points_to_BEV(pc, lidar_data, pose_record):\n",
    "#     cs_record = nusc.get('calibrated_sensor', lidar_data['calibrated_sensor_token'])\n",
    "    \n",
    "#     ref_to_ego = transform_matrix(translation=cs_record['translation'],\n",
    "#                                   rotation=Quaternion(cs_record[\"rotation\"]))\n",
    "\n",
    "#     # Compute rotation between 3D vehicle pose and \"flat\" vehicle pose (parallel to global z plane).\n",
    "#     ego_yaw = Quaternion(pose_record['rotation']).yaw_pitch_roll[0]\n",
    "#     rotation_vehicle_flat_from_vehicle = np.dot(\n",
    "#         Quaternion(scalar=np.cos(ego_yaw / 2), vector=[0, 0, np.sin(ego_yaw / 2)]).rotation_matrix,\n",
    "#         Quaternion(pose_record['rotation']).inverse.rotation_matrix)\n",
    "#     vehicle_flat_from_vehicle = np.eye(4)\n",
    "#     vehicle_flat_from_vehicle[:3, :3] = rotation_vehicle_flat_from_vehicle\n",
    "#     viewpoint = np.dot(vehicle_flat_from_vehicle, ref_to_ego) \n",
    "#     points = view_points(pc.points[:3, :], viewpoint, normalize=False)\n",
    "#     return points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "palette = {}\n",
    "\n",
    "palette['Bird'] = [165, 42, 42]\n",
    "palette['Ground Animal'] = [0, 192, 0]\n",
    "palette['Curb'] = [196, 196, 196]\n",
    "palette['Fence'] = [190, 153, 153]\n",
    "palette['Guard Rail'] = [180, 165, 180]\n",
    "palette['Barrier'] = [90, 120, 150]\n",
    "palette['Wall'] = [102, 102, 156]\n",
    "palette['Bile Lane'] = [128, 64, 255]\n",
    "palette['Crosswalk - Plain'] = [140, 140, 200]\n",
    "palette['Curb Cut'] = [170, 170, 170]\n",
    "palette['Parking'] = [250, 170, 160]\n",
    "palette['Pedestrian Area'] = [96, 96, 96]\n",
    "palette['Rail Track'] = [230, 150, 140]\n",
    "palette['Road'] = [128, 64, 128]\n",
    "palette['Service Lane'] = [110, 110, 110]\n",
    "palette['Sidewalk'] = [244, 35, 232]\n",
    "palette['Bridge'] = [150, 100, 100]\n",
    "palette['Building'] = [70, 70, 70]\n",
    "palette['Tunnel'] = [150, 120, 90]\n",
    "palette['Person'] = [220, 20, 60]\n",
    "palette['Bicyclist'] = [255, 0, 0]\n",
    "palette['Motorcyclist'] = [255, 0, 100]\n",
    "palette['Other Rider'] = [255, 0, 200]\n",
    "palette['Lane Marking - Crosswalk'] = [200, 128, 128]\n",
    "palette['Lane Marking - General'] = [255, 255, 255]\n",
    "palette['Mountain'] = [64, 170, 64]\n",
    "palette['Sand'] = [230, 160, 50]\n",
    "palette['Sky'] = [70, 130, 180]\n",
    "palette['Snow'] = [190, 255, 255]\n",
    "palette['Terrain'] = [152, 251, 152]\n",
    "palette['Vegetation'] = [107, 142, 35]\n",
    "palette['Water'] = [0, 170, 30]\n",
    "palette['Banner'] = [255, 255, 128]\n",
    "palette['Bench'] = [250, 0, 30]\n",
    "palette['Bike Rack'] = [100, 140, 180]\n",
    "palette['Billboard'] = [220, 220, 220]\n",
    "palette['Catch Basin'] = [220, 128, 128]\n",
    "palette['CCTV Camera'] = [222, 40, 40]\n",
    "palette['Fire Hydrant'] = [100, 170, 30]\n",
    "palette['Junction Box'] = [40, 40, 40]\n",
    "palette['Mailbox'] = [33, 33, 33]\n",
    "palette['Manhole'] = [100, 128, 160]\n",
    "palette['Phone Booth'] = [142, 0, 0]\n",
    "palette['Pothole'] = [70, 100, 150]\n",
    "palette['Street Light'] = [210, 170, 100]\n",
    "palette['Pole'] = [153, 153, 153]\n",
    "palette['Traffic Sign Frame'] = [128, 128, 128]\n",
    "palette['Utility Pole'] = [0, 0, 80]\n",
    "palette['Traffic Light'] = [250, 170, 30]\n",
    "palette['Traffic Sign (Back)'] = [192, 192, 192]\n",
    "palette['Traffic Sign (Front)'] = [220, 220, 0]\n",
    "palette['Trash Can'] = [140, 140, 20]\n",
    "palette['Bicycle'] = [119, 11, 32]\n",
    "palette['Boat'] = [150, 0, 255]\n",
    "palette['Bus'] = [0, 60, 100]\n",
    "palette['Car'] = [0, 0, 142]\n",
    "palette['Caravan'] = [0, 0, 90]\n",
    "palette['Motorcycle'] = [0, 0, 230]\n",
    "palette['On Rails'] = [0, 80, 100]\n",
    "palette['Other Vehicle'] = [128, 64, 64]\n",
    "palette['Trailer'] = [0, 0, 110]\n",
    "palette['Truck'] = [0, 0, 70]\n",
    "palette['Wheeled Slow'] = [0, 0, 192]\n",
    "palette['Car Mount'] = [32, 32, 32]\n",
    "palette['Ego Vehicle'] = [120, 10, 10]\n",
    "\n",
    "for k in palette.keys():\n",
    "    palette[k].append(255)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "pixel_to_classidx = {}\n",
    "class_to_idx = {}\n",
    "count = 0 \n",
    "for k in palette: \n",
    "    pixel_to_classidx[tuple(palette[k])] = (k, count)\n",
    "    class_to_idx[k] = count\n",
    "    count+=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{0: 'Bird',\n",
       " 1: 'Ground Animal',\n",
       " 2: 'Curb',\n",
       " 3: 'Fence',\n",
       " 4: 'Guard Rail',\n",
       " 5: 'Barrier',\n",
       " 6: 'Wall',\n",
       " 7: 'Bile Lane',\n",
       " 8: 'Crosswalk - Plain',\n",
       " 9: 'Curb Cut',\n",
       " 10: 'Parking',\n",
       " 11: 'Pedestrian Area',\n",
       " 12: 'Rail Track',\n",
       " 13: 'Road',\n",
       " 14: 'Service Lane',\n",
       " 15: 'Sidewalk',\n",
       " 16: 'Bridge',\n",
       " 17: 'Building',\n",
       " 18: 'Tunnel',\n",
       " 19: 'Person',\n",
       " 20: 'Bicyclist',\n",
       " 21: 'Motorcyclist',\n",
       " 22: 'Other Rider',\n",
       " 23: 'Lane Marking - Crosswalk',\n",
       " 24: 'Lane Marking - General',\n",
       " 25: 'Mountain',\n",
       " 26: 'Sand',\n",
       " 27: 'Sky',\n",
       " 28: 'Snow',\n",
       " 29: 'Terrain',\n",
       " 30: 'Vegetation',\n",
       " 31: 'Water',\n",
       " 32: 'Banner',\n",
       " 33: 'Bench',\n",
       " 34: 'Bike Rack',\n",
       " 35: 'Billboard',\n",
       " 36: 'Catch Basin',\n",
       " 37: 'CCTV Camera',\n",
       " 38: 'Fire Hydrant',\n",
       " 39: 'Junction Box',\n",
       " 40: 'Mailbox',\n",
       " 41: 'Manhole',\n",
       " 42: 'Phone Booth',\n",
       " 43: 'Pothole',\n",
       " 44: 'Street Light',\n",
       " 45: 'Pole',\n",
       " 46: 'Traffic Sign Frame',\n",
       " 47: 'Utility Pole',\n",
       " 48: 'Traffic Light',\n",
       " 49: 'Traffic Sign (Back)',\n",
       " 50: 'Traffic Sign (Front)',\n",
       " 51: 'Trash Can',\n",
       " 52: 'Bicycle',\n",
       " 53: 'Boat',\n",
       " 54: 'Bus',\n",
       " 55: 'Car',\n",
       " 56: 'Caravan',\n",
       " 57: 'Motorcycle',\n",
       " 58: 'On Rails',\n",
       " 59: 'Other Vehicle',\n",
       " 60: 'Trailer',\n",
       " 61: 'Truck',\n",
       " 62: 'Wheeled Slow',\n",
       " 63: 'Car Mount',\n",
       " 64: 'Ego Vehicle'}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class_to_idx\n",
    "idx_to_class = {class_to_idx[k] : k for k in class_to_idx}\n",
    "idx_to_class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_seg(og_seg, plot_images=False, colortype='rgb'):\n",
    "    '''\n",
    "    og_seg - image from which intermediate representations will be extracted from (array)\n",
    "    plot_images - plots images if True, does not plot if False\n",
    "    colortype - 'binary': plots 1 or 0 (used for occupancy grid), 'grayscale': plots in grayscale for lidar mapping,\n",
    "                'rgb' or any other string: plots in original palette colors\n",
    "    returns ret - list of intermediate representations(road, lane, and obstacle in that order so far)\n",
    "    '''\n",
    "    if(plot_images):\n",
    "        plt.figure()\n",
    "        plt.imshow(og_seg)\n",
    "    ret = []\n",
    "    #road segmentation\n",
    "    inter_seg_road = copy.deepcopy(og_seg)\n",
    "    inter_seg_road[(og_seg != palette['Road']).any(axis=2)] = [0,0,0,255]  \n",
    "    ret.append(inter_seg_road)\n",
    "    \n",
    "    #lane segmentation\n",
    "    inter_seg_lane = copy.deepcopy(og_seg)\n",
    "    crosswalk = (og_seg != palette['Lane Marking - Crosswalk']).any(axis=2)\n",
    "    general = (og_seg != palette['Lane Marking - General']).any(axis=2) \n",
    "    inter_seg_lane[np.logical_and(crosswalk, general)] = [0,0,0,255] \n",
    "    ret.append(inter_seg_lane)\n",
    "    \n",
    "    #obstacle segmentation (did not include Curb Cut as obstacle but did include Curb)\n",
    "    inter_seg_obstacle = copy.deepcopy(og_seg)\n",
    "    building = (og_seg != palette['Building']).any(axis=2)\n",
    "    curb = (og_seg != palette['Curb']).any(axis=2)\n",
    "    vegetation = (og_seg != palette['Vegetation']).any(axis=2)\n",
    "    inter_seg_obstacle[np.logical_and(np.logical_and(building,curb),vegetation)] = [0,0,0,255]\n",
    "    ret.append(inter_seg_obstacle)\n",
    "    \n",
    "    if(colortype == 'grayscale'):\n",
    "        for i in range(0, len(ret)):\n",
    "            temp = color.rgb2gray(ret[i])\n",
    "            ret[i] = temp\n",
    "            \n",
    "    elif(colortype == 'binary'):\n",
    "        for i in range(0, len(ret)):\n",
    "            temp = color.rgb2gray(ret[i])\n",
    "            temp[temp > 0] = 1\n",
    "            ret[i] = temp\n",
    "    \n",
    "    \n",
    "    if(plot_images):\n",
    "        for i in ret:\n",
    "            plt.figure()\n",
    "            if(colortype =='grayscale'):\n",
    "                plt.imshow(i, cmap='gray')\n",
    "            else:\n",
    "                plt.imshow(i)\n",
    "    return ret\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_semantic_class(points, cam_img):\n",
    "    num_points = points.shape[1]\n",
    "    \n",
    "    class_vec = [-1] * num_points\n",
    "    num_points_classified = 0\n",
    "    for i in range(num_points):\n",
    "        current_point = np.round(points[:,i][:2])\n",
    "        r = int(current_point[0])\n",
    "        c = int(current_point[1])\n",
    "        object_type, class_id = pixel_to_classidx[(tuple(cam_img[c,r]))]\n",
    "        class_vec[i] = class_id\n",
    "        num_points_classified+=1\n",
    "#     print(str(num_points_classified) + '/' + str(len(class_vec)))\n",
    "            \n",
    "    return class_vec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def project_points_to_image(current_pc, pointsensor, cam):\n",
    "    \n",
    "    pc = copy.deepcopy(current_pc)\n",
    "    # Points live in the point sensor frame. So they need to be transformed via global to the image plane.\n",
    "    # First step: transform the point-cloud to the ego vehicle frame for the timestamp of the sweep.\n",
    "    cs_record = nusc.get('calibrated_sensor', pointsensor['calibrated_sensor_token'])\n",
    "    pc.rotate(Quaternion(cs_record['rotation']).rotation_matrix)\n",
    "    pc.translate(np.array(cs_record['translation']))\n",
    "    \n",
    "#     import pdb; pdb.set_trace()\n",
    "    # Second step: transform to the global frame.\n",
    "    poserecord = nusc.get('ego_pose', pointsensor['ego_pose_token'])\n",
    "    pc.rotate(Quaternion(poserecord['rotation']).rotation_matrix)\n",
    "    pc.translate(np.array(poserecord['translation']))\n",
    "    \n",
    "    t2 = transform_matrix(translation=poserecord['translation'],rotation=Quaternion(poserecord['rotation']),inverse=True)\n",
    "#     import pdb; pdb.set_trace()\n",
    "    global_frame_pc = copy.deepcopy(pc)\n",
    "\n",
    "    # Third step: transform into the ego vehicle frame for the timestamp of the image.\n",
    "    poserecord = nusc.get('ego_pose', cam['ego_pose_token'])\n",
    "    pc.translate(-np.array(poserecord['translation']))\n",
    "    pc.rotate(Quaternion(poserecord['rotation']).rotation_matrix.T)\n",
    "\n",
    "    # Fourth step: transform into the camera.\n",
    "    cs_record = nusc.get('calibrated_sensor', cam['calibrated_sensor_token'])\n",
    "    pc.translate(-np.array(cs_record['translation']))\n",
    "    pc.rotate(Quaternion(cs_record['rotation']).rotation_matrix.T)\n",
    "    \n",
    "    return pc, global_frame_pc\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mask_img_points(points, depths, im, min_dist=1.0):\n",
    "    mask = np.ones(depths.shape[0], dtype=bool)\n",
    "    mask = np.logical_and(mask, depths > min_dist)\n",
    "    mask = np.logical_and(mask, points[0, :] > 1)\n",
    "    mask = np.logical_and(mask, points[0, :] < im.size[0] - 1)\n",
    "    mask = np.logical_and(mask, points[1, :] > 1)\n",
    "    mask = np.logical_and(mask, points[1, :] < im.size[1] - 1)\n",
    "    return mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] No such file or directory: 'full_data/sets/nuscenes/samples/LIDAR_TOP/n008-2018-08-01-15-16-36-0400__LIDAR_TOP__1533151603547590.pcd.bin'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-12-dc92e84c1ea5>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     29\u001b[0m     \u001b[0mcam_back\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnusc\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'sample_data'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcamera_token_BACK\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     30\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 31\u001b[0;31m     \u001b[0morig_pc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtimes\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mLidarPointCloud\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfrom_file_multisweep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnusc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msample_rec\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'LIDAR_TOP'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'LIDAR_TOP'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnsweeps\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\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     32\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     33\u001b[0m     \u001b[0;31m#get ego vehicle pose\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/nuscenes/utils/data_classes.py\u001b[0m in \u001b[0;36mfrom_file_multisweep\u001b[0;34m(cls, nusc, sample_rec, chan, ref_chan, nsweeps, min_distance)\u001b[0m\n\u001b[1;32m     97\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0m_\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnsweeps\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     98\u001b[0m             \u001b[0;31m# Load up the pointcloud and remove points close to the sensor.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 99\u001b[0;31m             \u001b[0mcurrent_pc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcls\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfrom_file\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mosp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjoin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnusc\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdataroot\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcurrent_sd_rec\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'filename'\u001b[0m\u001b[0;34m]\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    100\u001b[0m             \u001b[0mcurrent_pc\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mremove_close\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmin_distance\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    101\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/nuscenes/utils/data_classes.py\u001b[0m in \u001b[0;36mfrom_file\u001b[0;34m(cls, file_name)\u001b[0m\n\u001b[1;32m    252\u001b[0m         \u001b[0;32massert\u001b[0m \u001b[0mfile_name\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mendswith\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'.bin'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'Unsupported filetype {}'\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mformat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfile_name\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    253\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 254\u001b[0;31m         \u001b[0mscan\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfromfile\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfile_name\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfloat32\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    255\u001b[0m         \u001b[0mpoints\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mscan\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreshape\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m5\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m:\u001b[0m\u001b[0mcls\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnbr_dims\u001b[0m\u001b[0;34m(\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    256\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mcls\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpoints\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mT\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory: 'full_data/sets/nuscenes/samples/LIDAR_TOP/n008-2018-08-01-15-16-36-0400__LIDAR_TOP__1533151603547590.pcd.bin'"
     ]
    }
   ],
   "source": [
    "\n",
    "\n",
    "my_scene_token = nusc.field2token('scene', 'name', 'scene-0103')[0]\n",
    "scene_rec = nusc.get('scene', my_scene_token)\n",
    "\n",
    "current_token = scene_rec['first_sample_token']\n",
    "has_more = True \n",
    "scene_rec\n",
    "count = 0 \n",
    "history = []\n",
    "plot_result = True\n",
    "for i in range(scene_rec['nbr_samples']): \n",
    "    \n",
    "    #get current sample data\n",
    "    sample_rec = nusc.get('sample', current_token)\n",
    "\n",
    "    #get camera tokens\n",
    "    camera_token = sample_rec['data']['CAM_FRONT']\n",
    "    camera_token_FRONT_LEFT = sample_rec['data']['CAM_FRONT_LEFT']\n",
    "    camera_token_FRONT_RIGHT = sample_rec['data']['CAM_FRONT_RIGHT']\n",
    "    camera_token_BACK = sample_rec['data']['CAM_BACK']\n",
    "\n",
    "    #get lidar data info\n",
    "    pointsensor_token = sample_rec['data']['LIDAR_TOP']\n",
    "    pointsensor = nusc.get('sample_data', pointsensor_token)\n",
    "\n",
    "    #get camera info\n",
    "    cam = nusc.get('sample_data', camera_token)\n",
    "    cam_front_left = nusc.get('sample_data', camera_token_FRONT_LEFT)\n",
    "    cam_front_right = nusc.get('sample_data', camera_token_FRONT_RIGHT)\n",
    "    cam_back = nusc.get('sample_data', camera_token_BACK)\n",
    "    \n",
    "    orig_pc, times = LidarPointCloud.from_file_multisweep(nusc, sample_rec, 'LIDAR_TOP', 'LIDAR_TOP', nsweeps=1)\n",
    "    \n",
    "    #get ego vehicle pose\n",
    "    pose_record = nusc.get('ego_pose', pointsensor['ego_pose_token'])\n",
    "    \n",
    "    #get point cloud in camera frame prior to putting inside image plane\n",
    "    pc, global_frame_pc = project_points_to_image(orig_pc, pointsensor, cam) \n",
    "    pc_front_left, _ = project_points_to_image(orig_pc, pointsensor, cam_front_left)\n",
    "    pc_front_right, _ = project_points_to_image(orig_pc, pointsensor, cam_front_right)\n",
    "    pc_back, _ = project_points_to_image(orig_pc, pointsensor, cam_back)\n",
    "  \n",
    "    #get image representation of camera data\n",
    "    im = PIL.Image.open(osp.join(nusc.dataroot, cam['filename']))\n",
    "    im_fl = PIL.Image.open(osp.join(nusc.dataroot, cam_front_left['filename']))\n",
    "    im_fr = PIL.Image.open(osp.join(nusc.dataroot, cam_front_right['filename']))\n",
    "    im_b = PIL.Image.open(osp.join(nusc.dataroot, cam_back['filename']))\n",
    "\n",
    "     # Grab the depths (camera frame z axis points away from the camera).\n",
    "    depths = pc.points[2, :]\n",
    "    depths_front_left = pc_front_left.points[2,:]\n",
    "    depths_front_right = pc_front_right.points[2,:]\n",
    "    depths_back = pc_back.points[2,:]\n",
    "    \n",
    "#     print(osp.join(nusc.dataroot, cam['filename']))\n",
    "#     print(osp.join(nusc.dataroot, cam_front_left['filename']))\n",
    "#     print(osp.join(nusc.dataroot, cam_front_right['filename']))\n",
    "#     print(osp.join(nusc.dataroot, cam_back['filename']))\n",
    "\n",
    "    # Take the actual picture (matrix multiplication with camera-matrix + renormalization).\n",
    "    cs_record = nusc.get('calibrated_sensor', cam['calibrated_sensor_token'])\n",
    "    cs_record_front_left = nusc.get('calibrated_sensor', cam_front_left['calibrated_sensor_token'])\n",
    "    cs_record_front_right = nusc.get('calibrated_sensor', cam_front_right['calibrated_sensor_token'])\n",
    "    cs_record_back = nusc.get('calibrated_sensor', cam_back['calibrated_sensor_token'])\n",
    "    \n",
    "    #get matrix representation of camera image\n",
    "    cam_data_arr_front = plt.imread(osp.join(nusc.dataroot, cam['filename']))\n",
    "    cam_data_arr_front_left = plt.imread(osp.join(nusc.dataroot, cam_front_left['filename']))\n",
    "    cam_data_arr_front_right = plt.imread(osp.join(nusc.dataroot, cam_front_right['filename']))\n",
    "    cam_data_arr_back = plt.imread(osp.join(nusc.dataroot, cam_back['filename']))\n",
    "    \n",
    "    #get point cloud data in the image plane across all cameras\n",
    "    points = view_points(pc.points[:3, :], np.array(cs_record['camera_intrinsic']), normalize=True)\n",
    "    points_front_left = view_points(pc_front_left.points[:3, :], np.array(cs_record_front_left['camera_intrinsic']), normalize=True)\n",
    "    points_front_right = view_points(pc_front_right.points[:3, :], np.array(cs_record_front_right['camera_intrinsic']), normalize=True)\n",
    "    points_back = view_points(pc_back.points[:3, :], np.array(cs_record_back['camera_intrinsic']), normalize=True)\n",
    "\n",
    "    #get points that are actually inside the image plane\n",
    "    mask_front = mask_img_points(points, depths, im)\n",
    "    mask_front_left = mask_img_points(points_front_left, depths_front_left, im_fl)\n",
    "    mask_front_right = mask_img_points(points_front_right, depths_front_right, im_fr)\n",
    "    mask_back = mask_img_points(points_back, depths_back, im_b)\n",
    "    \n",
    "    # get points inside the image\n",
    "    valid_img_points_front = points[:, mask_front]\n",
    "    valid_img_points_front_left = points_front_left[:, mask_front_left]\n",
    "    valid_img_points_front_right = points_front_right[:, mask_front_right]\n",
    "    valid_img_points_back = points_back[:, mask_back]\n",
    "    \n",
    "    #perform point painting and get class per associated valid lidar point\n",
    "    semantic_class_points_front = get_semantic_class(valid_img_points_front, cam_data_arr_front)   \n",
    "    \n",
    "    semantic_class_points_front_left = get_semantic_class(valid_img_points_front_left, cam_data_arr_front_left)   \n",
    "\n",
    "    semantic_class_points_front_right = get_semantic_class(valid_img_points_front_right, cam_data_arr_front_right)   \n",
    "    semantic_class_points_back = get_semantic_class(valid_img_points_back, cam_data_arr_back)   \n",
    "    \n",
    "    masks = [mask_front, mask_front_left, mask_front_right, mask_back]\n",
    "    semantic_class_points = [semantic_class_points_front,\\\n",
    "                             semantic_class_points_front_left,\\\n",
    "                             semantic_class_points_front_right,\\\n",
    "                             semantic_class_points_back]\n",
    "    \n",
    "    \n",
    "    history.append((pc, global_frame_pc, pointsensor, pose_record, masks, semantic_class_points))\n",
    "    #can take history and predict\n",
    "    if i >= 2*past_seconds and i <= scene_rec['nbr_samples'] - 2*future_seconds: \n",
    "        needed_history = history[-(2*past_seconds+1):-1]\n",
    "        needed_ego_pose = history[-1][3]\n",
    "        \n",
    "        res = 0.75\n",
    "        offset = 64\n",
    "        \n",
    "        \n",
    "        yaw = Quaternion(needed_ego_pose['rotation']).yaw_pitch_roll[0]\n",
    "        rotation_matrix = np.array([[np.cos(yaw), -np.sin(yaw)],[np.sin(yaw), np.cos(yaw)]]).T\n",
    "                \n",
    "        count = i - 2*past_seconds\n",
    "        #with ego pose, transform the corresponding point clouds in history to bev \n",
    "        for pc_h, gpc_h, ps_h, ep_h, m, s_c_p in needed_history:\n",
    "            #generate target rep\n",
    "            shifted_ep_h = np.expand_dims(np.array(ep_h['translation']),1) - np.expand_dims(needed_ego_pose['translation'], 1)\n",
    "            transformed_ep_h = np.dot(rotation_matrix, shifted_ep_h[:2,:])\n",
    "            x = math.floor((int(transformed_ep_h[0]) + offset)/res)\n",
    "            y = math.floor((int(transformed_ep_h[1]) + offset)/res)\n",
    "            target_rep = np.zeros((256,256))\n",
    "            target_rep[y][x] += 1\n",
    "            \n",
    "            \n",
    "            #generate all reps except target\n",
    "            \n",
    "            gpc_h.points[:3,:] -= np.expand_dims(needed_ego_pose['translation'], 1)\n",
    "            gpc_h.points[:2,:] = np.dot(rotation_matrix, gpc_h.points[:2,:])\n",
    "\n",
    "            #ensures we don't get lidar points that are too far away\n",
    "            point_filter_1_x = gpc_h.points[0,:] >= -50\n",
    "            point_filter_1_y = gpc_h.points[1,:] >= -50\n",
    "            point_filter_1 = np.logical_and(point_filter_1_x, point_filter_1_y)\n",
    "            \n",
    "            point_filter_2_x = gpc_h.points[0,:] <= 50\n",
    "            point_filter_2_y = gpc_h.points[1,:] <= 50\n",
    "            point_filter_2 = np.logical_and(point_filter_2_x, point_filter_2_y)\n",
    "\n",
    "            point_filter = np.logical_and(point_filter_1, point_filter_2)\n",
    "        \n",
    "            m_f, m_f_l, m_f_r, m_b = m[0], m[1], m[2], m[3]\n",
    "            s_c_p_f, s_c_p_f_l, s_c_p_f_r, s_c_p_b = s_c_p[0], s_c_p[1], s_c_p[2], s_c_p[3]\n",
    "            overall_mask  = np.logical_or(m_f, m_f_l)\n",
    "            overall_mask  = np.logical_or(overall_mask, m_f_r)            \n",
    "            overall_mask  = np.logical_or(overall_mask, m_b)\n",
    "            \n",
    "            \n",
    "            all_colored = np.array([-1] * gpc_h.points.shape[1])\n",
    "            all_colored[m_f] = s_c_p[0]\n",
    "            all_colored[m_f_l] =  s_c_p[1]\n",
    "            all_colored[m_f_r] = s_c_p[2]\n",
    "            all_colored[m_b] = s_c_p[3]\n",
    "            \n",
    "            m_fl, m_fr = np.logical_and(point_filter, m_f_l),np.logical_and(point_filter, m_f_r)\n",
    "            m_front, m_back = np.logical_and(point_filter, m_f),np.logical_and(point_filter, m_b)\n",
    "            \n",
    "            s_c_p_f = all_colored[m_front]\n",
    "            s_c_p_f_l = all_colored[m_fl]\n",
    "            s_c_p_f_r = all_colored[m_fr]\n",
    "            s_c_p_b = all_colored[m_back]\n",
    "        \n",
    "            road_rep_fl, lane_rep_fl, obstacle_rep_fl, vehicle_rep_fl = get_intermediate_rep(gpc_h.points, m_fl, s_c_p_f_l)\n",
    "            road_rep_fr, lane_rep_fr, obstacle_rep_fr, vehicle_rep_fr = get_intermediate_rep(gpc_h.points, m_fr, s_c_p_f_r)\n",
    "            \n",
    "            road_rep_f, lane_rep_f, obstacle_rep_f, vehicle_rep_f = get_intermediate_rep(gpc_h.points, m_front, s_c_p_f)\n",
    "            road_rep_b, lane_rep_b, obstacle_rep_b, vehicle_rep_b = get_intermediate_rep(gpc_h.points, m_back, s_c_p_b)\n",
    "            \n",
    "            road_rep = road_rep_fl + road_rep_fr + road_rep_f + road_rep_b\n",
    "            lane_rep = lane_rep_fl + lane_rep_fr + lane_rep_f + lane_rep_b\n",
    "            \n",
    "            obstacle_rep = obstacle_rep_fl + obstacle_rep_fr + obstacle_rep_f + obstacle_rep_b\n",
    "            vehicle_rep = vehicle_rep_fl + vehicle_rep_fr + vehicle_rep_f + vehicle_rep_b\n",
    "            \n",
    "            \n",
    "            \n",
    "            \n",
    "#             if  plot_result: \n",
    "#                 #back\n",
    "#                 plt.axis((-100,100,-100,100))\n",
    "#                 plt.xlabel('x coordinate w.r.t to current ego pose')\n",
    "#                 plt.ylabel('y coordinate w.r.t to current ego pose')\n",
    "#                 plt.title(\"vehicle pose w.r.t to current pose\")\n",
    "#                 plt.scatter(gpc_h.points[:,m_b][0,:], gpc_h.points[:,m_b][1,:], c=s_c_p_b, s = 0.2)\n",
    "#                 plt.savefig(\"painted_history_{}_back.png\".format(count))\n",
    "#                 plt.clf()\n",
    "                \n",
    "#                 #front \n",
    "#                 plt.axis((-100,100,-100,100))\n",
    "#                 plt.xlabel('x coordinate w.r.t to current ego pose')\n",
    "#                 plt.ylabel('y coordinate w.r.t to current ego pose')\n",
    "#                 plt.title(\"vehicle pose w.r.t to current pose\")\n",
    "#                 plt.scatter(gpc_h.points[:,m_f][0,:], gpc_h.points[:,m_f][1,:],c=s_c_p_f, s = 0.2)\n",
    "#                 plt.savefig(\"painted_history_{}_front.png\".format(count))\n",
    "#                 plt.clf()\n",
    "                \n",
    "                \n",
    "#                 #front_left \n",
    "#                 plt.axis((-100,100,-100,100))\n",
    "#                 plt.xlabel('x coordinate w.r.t to current ego pose')\n",
    "#                 plt.ylabel('y coordinate w.r.t to current ego pose')\n",
    "#                 plt.title(\"vehicle pose w.r.t to current pose\")\n",
    "#                 plt.scatter(gpc_h.points[:,m_f_l][0,:], gpc_h.points[:,m_f_l][1,:],c=s_c_p_f_l, s = 0.2)\n",
    "#                 plt.savefig(\"painted_history_{}_front_left.png\".format(count))\n",
    "#                 plt.clf()\n",
    "                \n",
    "                \n",
    "#                 #front_right \n",
    "#                 plt.axis((-100,100,-100,100))\n",
    "#                 plt.xlabel('x coordinate w.r.t to current ego pose')\n",
    "#                 plt.ylabel('y coordinate w.r.t to current ego pose')\n",
    "#                 plt.title(\"vehicle pose w.r.t to current pose\")\n",
    "#                 plt.scatter(gpc_h.points[:,m_f_r][0,:], gpc_h.points[:,m_f_r][1,:],c=s_c_p_f_r, s = 0.2)\n",
    "#                 plt.savefig(\"painted_history_{}_front_right.png\".format(count))\n",
    "#                 plt.clf()\n",
    "                \n",
    "                \n",
    "#                 #combined\n",
    "#                 plt.axis((-100,100,-100,100))\n",
    "#                 plt.xlabel('x coordinate w.r.t to current ego pose')\n",
    "#                 plt.ylabel('y coordinate w.r.t to current ego pose')\n",
    "#                 plt.title(\"vehicle pose w.r.t to current pose\")\n",
    "#                 plt.scatter(gpc_h.points[:,overall_mask][0,:],gpc_h.points[:,overall_mask][1,:], c = all_colored[all_colored != -1], s = 0.2)\n",
    "#                 plt.savefig(\"painted_history_{}_all.png\".format(count))\n",
    "#                 plt.clf()\n",
    "\n",
    "                \n",
    "                \n",
    "#             count+=1\n",
    "        global_frame_pc.points[:3,:] = global_frame_pc.points[:3,:] - np.expand_dims(needed_ego_pose['translation'], 1)\n",
    "        global_frame_pc.points[:2,:] = np.dot(rotation_matrix, global_frame_pc.points[:2,:])\n",
    "        \n",
    "#         if plot_result: \n",
    "#                 #back\n",
    "#                 import pdb; pdb.set_trace()\n",
    "#                 plt.axis((-100,100,-100,100))\n",
    "#                 plt.xlabel('x coordinate w.r.t to current ego pose')\n",
    "#                 plt.ylabel('y coordinate w.r.t to current ego pose')\n",
    "#                 plt.title(\"vehicle pose w.r.t to current pose\")\n",
    "#                 plt.scatter(global_frame_pc.points[:,mask_back][0,:], global_frame_pc.points[:,mask_back][1,:], c=semantic_class_points_back, s = 0.2)\n",
    "#                 plt.savefig(\"painted_history_{}_back.png\".format(count))\n",
    "#                 plt.clf()\n",
    "                \n",
    "#                 #front \n",
    "#                 plt.axis((-100,100,-100,100))\n",
    "#                 plt.xlabel('x coordinate w.r.t to current ego pose')\n",
    "#                 plt.ylabel('y coordinate w.r.t to current ego pose')\n",
    "#                 plt.title(\"vehicle pose w.r.t to current pose\")\n",
    "#                 plt.scatter(global_frame_pc.points[:,mask_front][0,:], global_frame_pc.points[:,mask_front][1,:],c=semantic_class_points_front, s = 0.2)\n",
    "#                 plt.savefig(\"painted_history_{}_front.png\".format(count))\n",
    "#                 plt.clf()\n",
    "                \n",
    "                \n",
    "#                 #front_left \n",
    "#                 plt.axis((-100,100,-100,100))\n",
    "#                 plt.xlabel('x coordinate w.r.t to current ego pose')\n",
    "#                 plt.ylabel('y coordinate w.r.t to current ego pose')\n",
    "#                 plt.title(\"vehicle pose w.r.t to current pose\")\n",
    "#                 plt.scatter(global_frame_pc.points[:,mask_front_left][0,:], global_frame_pc.points[:,mask_front_left][1,:],c=semantic_class_points_front_left, s = 0.2)\n",
    "#                 plt.savefig(\"painted_history_{}_front_left.png\".format(count))\n",
    "#                 plt.clf()\n",
    "                \n",
    "                \n",
    "#                 #front_right \n",
    "#                 plt.axis((-100,100,-100,100))\n",
    "#                 plt.xlabel('x coordinate w.r.t to current ego pose')\n",
    "#                 plt.ylabel('y coordinate w.r.t to current ego pose')\n",
    "#                 plt.title(\"vehicle pose w.r.t to current pose\")\n",
    "#                 plt.scatter(global_frame_pc.points[:,mask_front_right][0,:], global_frame_pc.points[:,mask_front_right][1,:],c=semantic_class_points_front_right, s = 0.2)\n",
    "#                 plt.savefig(\"painted_history_{}_front_right.png\".format(count))\n",
    "#                 plt.clf()\n",
    "#                 break\n",
    "        \n",
    "    if not sample_rec['next'] == \"\":\n",
    "        current_token = sample_rec['next']\n",
    "    else:\n",
    "        has_more = False\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_stats(s_c_p):\n",
    "    \n",
    "    objects_found_front = {}\n",
    "\n",
    "    for point in s_c_p[0]:\n",
    "        if idx_to_class[point] not in objects_found_front:\n",
    "            objects_found_front[idx_to_class[point]] = 1\n",
    "        else:\n",
    "             objects_found_front[idx_to_class[point]] += 1\n",
    "\n",
    "    objects_found_front_left = {}\n",
    "\n",
    "    for point in s_c_p[1]:\n",
    "        if idx_to_class[point] not in objects_found_front_left:\n",
    "            objects_found_front_left[idx_to_class[point]] = 1\n",
    "        else:\n",
    "             objects_found_front_left[idx_to_class[point]] += 1\n",
    "\n",
    "    objects_found_front_right = {} \n",
    "\n",
    "    for point in s_c_p[2]:\n",
    "        if idx_to_class[point] not in objects_found_front_right:\n",
    "            objects_found_front_right[idx_to_class[point]] = 1\n",
    "        else:\n",
    "             objects_found_front_right[idx_to_class[point]] += 1\n",
    "\n",
    "    objects_found_back = {} \n",
    "\n",
    "    for point in s_c_p[3]:\n",
    "        if idx_to_class[point] not in objects_found_back:\n",
    "            objects_found_back[idx_to_class[point]] = 1\n",
    "        else:\n",
    "             objects_found_back[idx_to_class[point]] += 1\n",
    "                \n",
    "    print(\" Found in front camera\")\n",
    "    \n",
    "    for class_point in objects_found_front: \n",
    "        print(\"{} found {} times in front camera\".format(class_point, objects_found_front[class_point]))\n",
    "        \n",
    "    print(\" Found in front left camera\")\n",
    "    \n",
    "    for class_point in objects_found_front_left: \n",
    "        print(\"{} found {} times in front camera\".format(class_point, objects_found_front_left[class_point]))  \n",
    "    \n",
    "    print(\" Found in front right camera\")\n",
    "    \n",
    "    for class_point in objects_found_front_right: \n",
    "        print(\"{} found {} times in front camera\".format(class_point, objects_found_front_right[class_point]))  \n",
    "\n",
    "    print(\" Found in back camera\")\n",
    "    \n",
    "    for class_point in objects_found_back: \n",
    "        print(\"{} found {} times in front camera\".format(class_point, objects_found_back[class_point]))  \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "from nuscenes.eval.common.utils import quaternion_yaw\n",
    "from nuscenes.prediction import PredictHelper\n",
    "from nuscenes.prediction.input_representation.utils import get_crops, get_rotation_matrix, convert_to_pixel_coords"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_intermediate_rep(BEV_points, mask, semantic_class_points, center_coordinates):\n",
    "    corresponding_BEV_points= BEV_points[:,mask] \n",
    "    \n",
    "    num_points = corresponding_BEV_points.shape[1]\n",
    "    lane_rep =  np.zeros((256,256))\n",
    "    road_rep = np.zeros((256,256))\n",
    "    obstacle_rep = np.zeros((256,256))\n",
    "    res = 0.5\n",
    "    for i in range(num_points):        \n",
    "        row_pixel, column_pixel = convert_to_pixel_coords(corresponding_BEV_points[:,i][:2], \\\n",
    "                                                          center_coordinates, \\\n",
    "                                                          (128, 128), res)\n",
    "        if row_pixel < 0 or column_pixel < 0 or row_pixel >= 256 or column_pixel >= 256: \n",
    "            import pdb; pdb.set_trace()\n",
    "        semantic_class = semantic_class_points[i]\n",
    "        if semantic_class == class_to_idx['Road']:\n",
    "            road_rep[row_pixel][column_pixel] += 1\n",
    "            \n",
    "        if semantic_class == class_to_idx['Lane Marking - General'] or semantic_class == class_to_idx['Lane Marking - Crosswalk']:\n",
    "            lane_rep[row_pixel][column_pixel] += 1\n",
    "            \n",
    "        if semantic_class == class_to_idx['Building'] or semantic_class == class_to_idx['Curb'] or semantic_class == class_to_idx['Vegetation']:\n",
    "            obstacle_rep[row_pixel][column_pixel] += 1\n",
    "        \n",
    "        \n",
    "    \n",
    "    return road_rep, lane_rep, obstacle_rep\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "944.3590389713584 554.6627969832381 1072.3590389713584 682.6627969832381\n",
      "944.3590389713584 554.6627969832381 1072.3590389713584 682.6627969832381\n",
      "944.3590389713584 554.6627969832381 1072.3590389713584 682.6627969832381\n",
      "944.3590389713584 554.6627969832381 1072.3590389713584 682.6627969832381\n",
      "944.3590389713584 554.6627969832381 1072.3590389713584 682.6627969832381\n",
      "944.3590389713584 554.6627969832381 1072.3590389713584 682.6627969832381\n",
      "944.3590389713584 554.6627969832381 1072.3590389713584 682.6627969832381\n",
      "944.3590389713584 554.6627969832381 1072.3590389713584 682.6627969832381\n",
      "944.3590389713584 554.6627969832381 1072.3590389713584 682.6627969832381\n",
      "944.3590389713584 554.6627969832381 1072.3590389713584 682.6627969832381\n",
      "944.3590389713584 554.6627969832381 1072.3590389713584 682.6627969832381\n",
      "944.3590389713584 554.6627969832381 1072.3590389713584 682.6627969832381\n",
      "> <ipython-input-29-e7561d9bdc01>(284)<module>()\n",
      "-> if not sample_rec['next'] == \"\":\n",
      "(Pdb) q\n"
     ]
    },
    {
     "ename": "BdbQuit",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mBdbQuit\u001b[0m                                   Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-29-e7561d9bdc01>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m    282\u001b[0m             \u001b[0;32mimport\u001b[0m \u001b[0mpdb\u001b[0m\u001b[0;34m;\u001b[0m \u001b[0mpdb\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mset_trace\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    283\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 284\u001b[0;31m         \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0msample_rec\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'next'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m==\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    285\u001b[0m             \u001b[0mcurrent_token\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msample_rec\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'next'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    286\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-29-e7561d9bdc01>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m    282\u001b[0m             \u001b[0;32mimport\u001b[0m \u001b[0mpdb\u001b[0m\u001b[0;34m;\u001b[0m \u001b[0mpdb\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mset_trace\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    283\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 284\u001b[0;31m         \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0msample_rec\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'next'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m==\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    285\u001b[0m             \u001b[0mcurrent_token\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msample_rec\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'next'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    286\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/opt/conda/lib/python3.7/bdb.py\u001b[0m in \u001b[0;36mtrace_dispatch\u001b[0;34m(self, frame, event, arg)\u001b[0m\n\u001b[1;32m     86\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0;31m# None\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     87\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mevent\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'line'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 88\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdispatch_line\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mframe\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     89\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mevent\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'call'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     90\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdispatch_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mframe\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0marg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/opt/conda/lib/python3.7/bdb.py\u001b[0m in \u001b[0;36mdispatch_line\u001b[0;34m(self, frame)\u001b[0m\n\u001b[1;32m    111\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstop_here\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mframe\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbreak_here\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mframe\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    112\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muser_line\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mframe\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 113\u001b[0;31m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mquitting\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;32mraise\u001b[0m \u001b[0mBdbQuit\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    114\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrace_dispatch\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    115\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mBdbQuit\u001b[0m: "
     ]
    }
   ],
   "source": [
    "past_seconds = 2\n",
    "future_seconds = 4\n",
    "helper = PredictHelper(nusc)\n",
    "\n",
    "for j in range(len(nusc.scene)//10):\n",
    "\n",
    "    #every scene\n",
    "    os.makedirs(\"nuScenes_project_dataset/scene_\" + str(j))\n",
    "    \n",
    "    print(j)\n",
    "    my_scene_token = nusc.scene[j]['token']\n",
    "    scene_rec = nusc.get('scene', my_scene_token)\n",
    "    #print(len(scene_rec))\n",
    "\n",
    "    current_token = scene_rec['first_sample_token']\n",
    "    current_seq = []\n",
    "    \n",
    "    sequence_count = 0\n",
    "    \n",
    "    for i in range(scene_rec['nbr_samples']): \n",
    "        #get current sample data\n",
    "        sample_rec = nusc.get('sample', current_token)\n",
    "        \n",
    "        annotation_tokens = sample_rec['anns']\n",
    "\n",
    "        #get camera tokens\n",
    "        camera_token = sample_rec['data']['CAM_FRONT']\n",
    "        camera_token_FRONT_LEFT = sample_rec['data']['CAM_FRONT_LEFT']\n",
    "        camera_token_FRONT_RIGHT = sample_rec['data']['CAM_FRONT_RIGHT']\n",
    "        camera_token_BACK = sample_rec['data']['CAM_BACK']\n",
    "\n",
    "        pointsensor_token = sample_rec['data']['LIDAR_TOP']\n",
    "        pointsensor = nusc.get('sample_data', pointsensor_token)\n",
    "\n",
    "        #get camera info\n",
    "        cam = nusc.get('sample_data', camera_token)\n",
    "        cam_front_left = nusc.get('sample_data', camera_token_FRONT_LEFT)\n",
    "        cam_front_right = nusc.get('sample_data', camera_token_FRONT_RIGHT)\n",
    "        cam_back = nusc.get('sample_data', camera_token_BACK)\n",
    "\n",
    "        orig_pc, times = LidarPointCloud.from_file_multisweep(nusc, sample_rec, 'LIDAR_TOP', 'LIDAR_TOP', nsweeps=10)\n",
    "\n",
    "        #get ego vehicle pose\n",
    "        pose_record = nusc.get('ego_pose', pointsensor['ego_pose_token'])\n",
    "\n",
    "        #get point cloud in camera frame prior to putting inside image plane\n",
    "        pc, global_frame_pc = project_points_to_image(orig_pc, pointsensor, cam) \n",
    "        pc_front_left, _ = project_points_to_image(orig_pc, pointsensor, cam_front_left)\n",
    "        pc_front_right, _ = project_points_to_image(orig_pc, pointsensor, cam_front_right)\n",
    "        pc_back, _ = project_points_to_image(orig_pc, pointsensor, cam_back)\n",
    "        \n",
    "#         if 'sweeps' in cam['filename'] or 'sweeps' in cam_front_left['filename'] or 'sweeps' in cam_front_right['filename'] or 'sweeps' in cam_back['filename']:\n",
    "#             print(\"found in sweeps\")\n",
    "\n",
    "        #get image representation of camera data\n",
    "        im = PIL.Image.open(osp.join(nusc.dataroot, cam['filename']))\n",
    "        im_fl = PIL.Image.open(osp.join(nusc.dataroot, cam_front_left['filename']))\n",
    "        im_fr = PIL.Image.open(osp.join(nusc.dataroot, cam_front_right['filename']))\n",
    "        im_b = PIL.Image.open(osp.join(nusc.dataroot, cam_back['filename']))\n",
    "        \n",
    "         # Grab the depths (camera frame z axis points away from the camera).\n",
    "        depths = pc.points[2, :]\n",
    "        depths_front_left = pc_front_left.points[2,:]\n",
    "        depths_front_right = pc_front_right.points[2,:]\n",
    "        depths_back = pc_back.points[2,:]\n",
    "\n",
    "        # Take the actual picture (matrix multiplication with camera-matrix + renormalization).\n",
    "        cs_record = nusc.get('calibrated_sensor', cam['calibrated_sensor_token'])\n",
    "        cs_record_front_left = nusc.get('calibrated_sensor', cam_front_left['calibrated_sensor_token'])\n",
    "        cs_record_front_right = nusc.get('calibrated_sensor', cam_front_right['calibrated_sensor_token'])\n",
    "        cs_record_back = nusc.get('calibrated_sensor', cam_back['calibrated_sensor_token'])\n",
    "\n",
    "        #get matrix representation of camera image\n",
    "        cam_data_arr_front = plt.imread(osp.join(nusc.dataroot, cam['filename']))\n",
    "        cam_data_arr_front_left = plt.imread(osp.join(nusc.dataroot, cam_front_left['filename']))\n",
    "        cam_data_arr_front_right = plt.imread(osp.join(nusc.dataroot, cam_front_right['filename']))\n",
    "        cam_data_arr_back = plt.imread(osp.join(nusc.dataroot, cam_back['filename']))\n",
    "\n",
    "        #get point cloud data in the image plane across all cameras\n",
    "        points = view_points(pc.points[:3, :], np.array(cs_record['camera_intrinsic']), normalize=True)\n",
    "        points_front_left = view_points(pc_front_left.points[:3, :], np.array(cs_record_front_left['camera_intrinsic']), normalize=True)\n",
    "        points_front_right = view_points(pc_front_right.points[:3, :], np.array(cs_record_front_right['camera_intrinsic']), normalize=True)\n",
    "        points_back = view_points(pc_back.points[:3, :], np.array(cs_record_back['camera_intrinsic']), normalize=True)\n",
    "\n",
    "        #get points that are actually inside the image plane\n",
    "        mask_front = mask_img_points(points, depths, im)\n",
    "        mask_front_left = mask_img_points(points_front_left, depths_front_left, im_fl)\n",
    "        mask_front_right = mask_img_points(points_front_right, depths_front_right, im_fr)\n",
    "        mask_back = mask_img_points(points_back, depths_back, im_b)\n",
    "\n",
    "        # get points inside the image\n",
    "        valid_img_points_front = points[:, mask_front]\n",
    "        valid_img_points_front_left = points_front_left[:, mask_front_left]\n",
    "        valid_img_points_front_right = points_front_right[:, mask_front_right]\n",
    "        valid_img_points_back = points_back[:, mask_back]\n",
    "\n",
    "#         import pdb; pdb.set_trace()\n",
    "        \n",
    "        #perform point painting and get class per associated valid lidar point\n",
    "        semantic_class_points_front = get_semantic_class(valid_img_points_front, cam_data_arr_front)   \n",
    "\n",
    "        semantic_class_points_front_left = get_semantic_class(valid_img_points_front_left, cam_data_arr_front_left)   \n",
    "\n",
    "        semantic_class_points_front_right = get_semantic_class(valid_img_points_front_right, cam_data_arr_front_right)   \n",
    "        semantic_class_points_back = get_semantic_class(valid_img_points_back, cam_data_arr_back)   \n",
    "\n",
    "        masks = [mask_front, mask_front_left, mask_front_right, mask_back]\n",
    "        semantic_class_points = [semantic_class_points_front,\\\n",
    "                                 semantic_class_points_front_left,\\\n",
    "                                 semantic_class_points_front_right,\\\n",
    "                                 semantic_class_points_back]\n",
    "        images = [im, im_fl, im_fr, im_b]\n",
    "\n",
    "        current_seq.append((pc, global_frame_pc, pointsensor, pose_record, masks, semantic_class_points, images, current_token))\n",
    "        \n",
    "        #create sequences\n",
    "        if len(current_seq) >= 13:\n",
    "            #every sequence\n",
    "            os.makedirs(\"nuScenes_project_dataset/scene_\" + str(j) + \"/sequence_\" + str(sequence_count))\n",
    "            \n",
    "            intermediate_reps = []\n",
    "            \n",
    "            current_frame = i - 2*future_seconds\n",
    "            \n",
    "            time_frame = current_seq[current_frame-2*past_seconds:-1]\n",
    "            needed_ego_pose = current_seq[current_frame][3]\n",
    "\n",
    "            res = 0.5\n",
    "            offset = 64\n",
    "\n",
    "            yaw = Quaternion(needed_ego_pose['rotation']).yaw_pitch_roll[0]\n",
    "            rotation_matrix = np.array([[np.cos(yaw), -np.sin(yaw)],[np.sin(yaw), np.cos(yaw)]]).T\n",
    "            ref_ego_x, ref_ego_y = needed_ego_pose['translation'][:2]\n",
    "            \n",
    "            \n",
    "#             import pdb; pdb.set_trace()\n",
    "            #with ego pose, transform the corresponding point clouds in history to bev \n",
    "            #count = 0\n",
    "        \n",
    "            for pc_h, gpc_orig, ps_h, ep_h, m, s_c_p, ims, token in time_frame:\n",
    "            \n",
    "                gpc_h = copy.deepcopy(gpc_orig)\n",
    "                \n",
    "#                 import pdb; pdb.set_trace()\n",
    "\n",
    "                #get target vehicle representation\n",
    "                location = ep_h['translation'][:2]\n",
    "                row_pixel, column_pixel = convert_to_pixel_coords(location, \\\n",
    "                                                                  (ref_ego_x, ref_ego_y), \\\n",
    "                                                                  (128, 128), res)\n",
    "                target_rep = np.zeros((256,256))\n",
    "                width_target = 1\n",
    "                length_target = 1\n",
    "                box = pixels_to_box_corners(row_pixel, column_pixel, length_target, width_target, yaw)\n",
    "                cv2.fillPoly(target_rep, pts=[np.int0(box)], color=1)\n",
    "                rotation_mat = get_rotation_matrix(target_rep.shape, yaw+np.pi/2)\n",
    "                target_rep = cv2.warpAffine(target_rep, rotation_mat, target_rep.shape)\n",
    "\n",
    "                \n",
    "                \n",
    "                #get other vehicle representation \n",
    "                history = helper.get_past_for_sample(token,\n",
    "                                      0,\n",
    "                                      in_agent_frame=False,\n",
    "                                      just_xy=False)\n",
    "                history = reverse_history(history)\n",
    "                \n",
    "#                 import pdb; pdb.set_trace()\n",
    "\n",
    "                present_time = helper.get_annotations_for_sample(token)\n",
    "\n",
    "                history = add_present_time_to_history(present_time, history)\n",
    "                vehicle_rep = np.zeros((256, 256))\n",
    "                draw_other_vehicle_boxes(needed_ego_pose, (128,128),\n",
    "                         history, vehicle_rep, resolution=res)\n",
    "                vehicle_rep = cv2.warpAffine(vehicle_rep, rotation_mat, vehicle_rep.shape)\n",
    "                \n",
    "    \n",
    "                #all the representations have same shape of (256, 256)\n",
    "                                            \n",
    "                \n",
    "#                 target_rep[row_pixel][column_pixel] += 1\n",
    "#                 target_rep = cv2.warpAffine(target_rep, rotation_mat, target_rep.shape)\n",
    "            \n",
    "\n",
    "                #generate all reps except target\n",
    "\n",
    "#                 gpc_h.points[:3,:] -= np.expand_dims(needed_ego_pose['translation'], 1)\n",
    "#                 gpc_h.points[:2,:] = np.dot(rotation_matrix, gpc_h.points[:2,:])\n",
    "                \n",
    "#                 import pdb; pdb.set_trace()\n",
    "                \n",
    "                min_coor_x = ref_ego_x - offset\n",
    "                min_coor_y = ref_ego_y - offset\n",
    "                \n",
    "                max_coor_x = ref_ego_x + offset\n",
    "                max_coor_y = ref_ego_y + offset\n",
    "                \n",
    "                print(min_coor_x, min_coor_y, max_coor_x, max_coor_y)\n",
    "                \n",
    "                #ensures we don't get lidar points that are too far away\n",
    "                point_filter_1_x = gpc_h.points[0,:] > min_coor_x\n",
    "                point_filter_1_y = gpc_h.points[1,:] > min_coor_y\n",
    "                point_filter_1 = np.logical_and(point_filter_1_x, point_filter_1_y)\n",
    "\n",
    "                point_filter_2_x = gpc_h.points[0,:] < max_coor_x\n",
    "                point_filter_2_y = gpc_h.points[1,:] < max_coor_y\n",
    "                point_filter_2 = np.logical_and(point_filter_2_x, point_filter_2_y)\n",
    "\n",
    "                point_filter = np.logical_and(point_filter_1, point_filter_2)\n",
    "                \n",
    "                m_f, m_f_l, m_f_r, m_b = m[0], m[1], m[2], m[3]\n",
    "                s_c_p_f, s_c_p_f_l, s_c_p_f_r, s_c_p_b = s_c_p[0], s_c_p[1], s_c_p[2], s_c_p[3]\n",
    "\n",
    "\n",
    "                all_colored = np.array([-1] * gpc_h.points.shape[1])\n",
    "                all_colored[m_f] = s_c_p[0]\n",
    "                all_colored[m_f_l] =  s_c_p[1]\n",
    "                all_colored[m_f_r] = s_c_p[2]\n",
    "                all_colored[m_b] = s_c_p[3]\n",
    "\n",
    "                m_f_l, m_f_r = np.logical_and(point_filter, m_f_l),np.logical_and(point_filter, m_f_r)\n",
    "                m_f, m_b = np.logical_and(point_filter, m_f),np.logical_and(point_filter, m_b)\n",
    "\n",
    "                s_c_p_f = all_colored[m_f]\n",
    "                s_c_p_f_l = all_colored[m_f_l]\n",
    "                s_c_p_f_r = all_colored[m_f_r]\n",
    "                s_c_p_b = all_colored[m_b]\n",
    "                \n",
    "                road_rep_fl, lane_rep_fl, obstacle_rep_fl = get_intermediate_rep(gpc_h.points, m_f_l, s_c_p_f_l, (ref_ego_x, ref_ego_y))\n",
    "                road_rep_fr, lane_rep_fr, obstacle_rep_fr = get_intermediate_rep(gpc_h.points, m_f_r, s_c_p_f_r, (ref_ego_x, ref_ego_y))\n",
    "                road_rep_f, lane_rep_f, obstacle_rep_f = get_intermediate_rep(gpc_h.points, m_f, s_c_p_f, (ref_ego_x, ref_ego_y))\n",
    "                road_rep_b, lane_rep_b, obstacle_rep_b = get_intermediate_rep(gpc_h.points, m_b, s_c_p_b, (ref_ego_x, ref_ego_y))\n",
    "\n",
    "                road_rep = road_rep_fl + road_rep_fr + road_rep_f + road_rep_b\n",
    "                lane_rep = lane_rep_fl + lane_rep_fr + lane_rep_f + lane_rep_b\n",
    "\n",
    "                obstacle_rep = obstacle_rep_fl + obstacle_rep_fr + obstacle_rep_f + obstacle_rep_b\n",
    "                \n",
    "                #rotate image and perform the dilation for the road, lane, and obstacle representations\n",
    "                road_rep = cv2.warpAffine(road_rep, rotation_mat, road_rep.shape)\n",
    "                kernel = np.ones((5,5), np.uint8)\n",
    "                road_rep = cv2.dilate(road_rep,kernel,iterations = 1)\n",
    "                \n",
    "                lane_rep = cv2.warpAffine(lane_rep, rotation_mat, lane_rep.shape)\n",
    "                kernel = np.ones((2,2), np.uint8)\n",
    "                lane_rep = cv2.dilate(lane_rep,kernel,iterations = 1)\n",
    "                \n",
    "                obstacle_rep = cv2.warpAffine(obstacle_rep, rotation_mat, obstacle_rep.shape)\n",
    "                kernel = np.ones((5,5), np.uint8)\n",
    "                obstacle_rep = cv2.dilate(obstacle_rep,kernel,iterations = 1)                \n",
    "                \n",
    "                intermediate_reps.append([target_rep, lane_rep, obstacle_rep, road_rep, vehicle_rep]) \n",
    "\n",
    "            frame_count = 0\n",
    "            for k,data in enumerate(intermediate_reps):\n",
    "                if k < len(intermediate_reps) - 1:\n",
    "                    data.append(intermediate_reps[k+1][4])\n",
    "                else:\n",
    "                    ego_pose_current = copy.deepcopy(current_seq[-1][3])\n",
    "                    #generate target rep\n",
    "                    location = ego_pose_current['translation'][:2]\n",
    "                    row_pixel, column_pixel = convert_to_pixel_coords(location, \\\n",
    "                                                                  (ref_ego_x, ref_ego_y), \\\n",
    "                                                                  (128, 128), res)\n",
    "                    target_rep = np.zeros((256,256))\n",
    "                    width_target = 1\n",
    "                    length_target = 1\n",
    "                    box = pixels_to_box_corners(row_pixel, column_pixel, length_target, width_target, yaw)\n",
    "                    cv2.fillPoly(target_rep, pts=[np.int0(box)], color=1)\n",
    "                    rotation_mat = get_rotation_matrix(target_rep.shape, yaw+np.pi/2)\n",
    "                    target_rep = cv2.warpAffine(target_rep, rotation_mat, target_rep.shape)\n",
    "                    \n",
    "                    data.append(target_rep)\n",
    "                #every frame\n",
    "                np.save(\"nuScenes_project_dataset/scene_\" + str(j) + \"/sequence_\" + str(sequence_count) + \"/frame_\" + str(frame_count), np.array(data))\n",
    "                frame_count += 1\n",
    "                #import pdb; pdb.set_trace()\n",
    "            \n",
    "            #print(data)\n",
    "            sequence_count += 1\n",
    "            import pdb; pdb.set_trace()\n",
    "            \n",
    "        if not sample_rec['next'] == \"\":\n",
    "            current_token = sample_rec['next']\n",
    "        else:\n",
    "            has_more = False\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "#os.makedirs(\"nuScenes_project_dataset\")\n",
    "a = 1\n",
    "b = 16\n",
    "c = 10\n",
    "#print(\"nuScenes_project_dataset/sequence_\" + str(a))\n",
    "#os.makedirs(\"nuScenes_project_dataset/scene_\" + str(a) + \"/sequence_\" + str(b) + \"/frame_\" + str(c))\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### NEW STRUCTURE\n",
    "#for every scene (85)\n",
    "    #for every sequence (start of history to end of future ground truth) around 28 or 29\n",
    "        #store intermediate reps with ground truths\n",
    "        \n",
    "#talk to Nachiket about ground truths\n",
    "#create data directory above\n",
    "#work on training script \n",
    "#work on validation script - mimic what authors did"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Any, Dict, List, Tuple, Callable\n",
    "import numpy as np\n",
    "History = Dict[str, List[Dict[str, Any]]]\n",
    "\n",
    "def reverse_history(history: History) -> History:\n",
    "    \"\"\"\n",
    "    Reverse history so that most distant observations are first.\n",
    "    We do this because we want to draw more recent bounding boxes on top of older ones.\n",
    "    :param history: result of get_past_for_sample PredictHelper method.\n",
    "    :return: History with the values reversed.\n",
    "    \"\"\"\n",
    "    return {token: anns[::-1] for token, anns in history.items()}\n",
    "\n",
    "\n",
    "def add_present_time_to_history(current_time: List[Dict[str, Any]],\n",
    "                                history: History) -> History:\n",
    "    \"\"\"\n",
    "    Adds the sample annotation records from the current time to the\n",
    "    history object.\n",
    "    :param current_time: List of sample annotation records from the\n",
    "        current time. Result of get_annotations_for_sample method of\n",
    "        PredictHelper.\n",
    "    :param history: Result of get_past_for_sample method of PredictHelper.\n",
    "    :return: History with values from current_time appended.\n",
    "    \"\"\"\n",
    "\n",
    "    for annotation in current_time:\n",
    "        token = annotation['instance_token']\n",
    "\n",
    "        if token in history:\n",
    "\n",
    "            # We append because we've reversed the history\n",
    "            history[token].append(annotation)\n",
    "\n",
    "        else:\n",
    "            history[token] = [annotation]\n",
    "\n",
    "    return history\n",
    "\n",
    "def pixels_to_box_corners(row_pixel: int,\n",
    "                          column_pixel: int,\n",
    "                          length_in_pixels: float,\n",
    "                          width_in_pixels: float,\n",
    "                          yaw_in_radians: float) -> np.ndarray:\n",
    "    \"\"\"\n",
    "    Computes four corners of 2d bounding box for agent.\n",
    "    The coordinates of the box are in pixels.\n",
    "    :param row_pixel: Row pixel of the agent.\n",
    "    :param column_pixel: Column pixel of the agent.\n",
    "    :param length_in_pixels: Length of the agent.\n",
    "    :param width_in_pixels: Width of the agent.\n",
    "    :param yaw_in_radians: Yaw of the agent (global coordinates).\n",
    "    :return: numpy array representing the four corners of the agent.\n",
    "    \"\"\"\n",
    "\n",
    "    # cv2 has the convention where they flip rows and columns so it matches\n",
    "    # the convention of x and y on a coordinate plane\n",
    "    # Also, a positive angle is a clockwise rotation as opposed to counterclockwise\n",
    "    # so that is why we negate the rotation angle\n",
    "    coord_tuple = ((column_pixel, row_pixel), (length_in_pixels, width_in_pixels), -yaw_in_radians * 180 / np.pi)\n",
    "\n",
    "    box = cv2.boxPoints(coord_tuple)\n",
    "\n",
    "    return box\n",
    "\n",
    "\n",
    "def get_track_box(annotation: Dict[str, Any],\n",
    "                  center_coordinates: Tuple[float, float],\n",
    "                  center_pixels: Tuple[float, float],\n",
    "                  resolution: float = 0.1) -> np.ndarray:\n",
    "    \"\"\"\n",
    "    Get four corners of bounding box for agent in pixels.\n",
    "    :param annotation: The annotation record of the agent.\n",
    "    :param center_coordinates: (x, y) coordinates in global frame\n",
    "        of the center of the image.\n",
    "    :param center_pixels: (row_index, column_index) location of the center\n",
    "        of the image in pixel coordinates.\n",
    "    :param resolution: Resolution pixels/meter of the image.\n",
    "    \"\"\"\n",
    "\n",
    "    assert resolution > 0\n",
    "\n",
    "    location = annotation['translation'][:2]\n",
    "    yaw_in_radians = quaternion_yaw(Quaternion(annotation['rotation']))\n",
    "\n",
    "    #print('yaw_in_radians', yaw_in_radians)\n",
    "    row_pixel, column_pixel = convert_to_pixel_coords(location,\n",
    "                                                      center_coordinates,\n",
    "                                                      center_pixels, resolution)\n",
    "    #print('row_pixel, column_pixel', row_pixel, column_pixel)\n",
    "    #print('center_pixels', center_pixels)\n",
    "\n",
    "    width = annotation['size'][0] / resolution\n",
    "    length = annotation['size'][1] / resolution\n",
    "\n",
    "    # Width and length are switched here so that we can draw them along the x-axis as\n",
    "    # opposed to the y. This makes rotation easier.\n",
    "    return pixels_to_box_corners(row_pixel, column_pixel, length, width, yaw_in_radians)\n",
    "\n",
    "\n",
    "def draw_other_vehicle_boxes(ref_ego_pose: Dict[str, Any],\n",
    "                        center_agent_pixels: Tuple[float, float],\n",
    "                        vehicle_history: History,\n",
    "                        base_image: np.ndarray,\n",
    "                        resolution: float = 0.1) -> None:\n",
    "    color = (1, 1, 1)\n",
    "    ref_loc_x, ref_loc_y = ref_ego_pose['translation'][:2]\n",
    "\n",
    "    for instance_token, annotations in vehicle_history.items():\n",
    "\n",
    "        num_points = len(annotations)\n",
    "\n",
    "        for i, annotation in enumerate(annotations):\n",
    "\n",
    "            box = get_track_box(annotation, (ref_loc_x, ref_loc_y), center_agent_pixels, resolution)\n",
    "\n",
    "            if 'object' in annotation['category_name']:\n",
    "                continue\n",
    "            \n",
    "            cv2.fillPoly(base_image, pts=[np.int0(box)], color=color)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for file in os.listdir(\"nuScenes_project_dataset/scene_0/sequence_0\"):\n",
    "    print(file)\n",
    "    file_path = \"nuScenes_project_dataset/scene_0/sequence_0/\" + file\n",
    "    frame = np.load(file_path)\n",
    "    print(frame.shape)\n",
    "    plt.imshow(frame[2,:,:])\n",
    "    plt.show()\n",
    "    import pdb; pdb.set_trace()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#main dataset\n",
    "    #scene\n",
    "        #sequence\n",
    "            # x number of frames"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}