{"cells":[{"metadata":{},"cell_type":"markdown","source":"Main goal here is to learn to work with the data as-is (because not all needed features are presnet in the SDK), join all lidars into one point cloud and apply annotations to make sure we're not missing anything.\n\n**EDIT** as @alexamadori pointed out in comments, most scenes have data only from one lidar, so merging all 3 lidars might not make much sense. Still writing this kernel was useful to me to understand the data and coordinate systems better.\n\nOverview of what we're doing below:\n\n- we load lidar data (3d points), these points are in coordinate system of the lidar, rotated and translated relative to the car\n- using sensor information of the lidar, we translate points from lidar coordinate frame to car coordinate frame, this allows us to merge data from all 3 lidars\n- annotations and submission are in global coordinates. We translate annotations into the car coordinates, which allows to have both the data and annotations in the same (car) coordinate frame, which can then be used for training\n\nWe learn:\n\n- how the raw data looks like, so that we understand how SDK works better and can extend it if needed (also if you see any missing features, file issues at https://github.com/lyft/nuscenes-devkit/issues/)\n- how to translate objects between various coordinate frames\n\nIn terms of custom packages, we'll use only pyquaternion which is also used by Lyft SDK."},{"metadata":{"trusted":true},"cell_type":"code","source":"%matplotlib inline\nimport json\nimport os.path\n\nfrom matplotlib import pyplot as plt\nimport numpy as np\nimport pandas as pd\nfrom pyquaternion import Quaternion","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Some small helpers for loading the json data. We won't load all data here."},{"metadata":{"trusted":true},"cell_type":"code","source":"class Table:\n    def __init__(self, data):\n        self.data = data\n        self.index = {x['token']: x for x in data}\n\n\nDATA_ROOT = '/kaggle/input/3d-object-detection-for-autonomous-vehicles/'\n\n\ndef load_table(name, root=os.path.join(DATA_ROOT, 'train_data')):\n    with open(os.path.join(root, name), 'rb') as f:\n        return Table(json.load(f))\n\n    \nscene = load_table('scene.json')\nsample = load_table('sample.json')\nsample_data = load_table('sample_data.json')\nego_pose = load_table('ego_pose.json')\ncalibrated_sensor = load_table('calibrated_sensor.json')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Like in official tutorial, let's check the first scene"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"my_scene = scene.data[0]\nmy_scene","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"And the first sample from that scene"},{"metadata":{"trusted":true},"cell_type":"code","source":"sample.index[my_scene['first_sample_token']]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We'll use ``sample_data`` to fetch lidar images related to this sample. First check what is inside ``sample_data``"},{"metadata":{"trusted":true},"cell_type":"code","source":"sample_data.data[0]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Now fetch lidar images related to the sample (note that to make it efficient you'll want to add an index like in official SDK)"},{"metadata":{"trusted":true},"cell_type":"code","source":"lidars = []\nfor x in sample_data.data:\n    if x['sample_token'] == my_scene['first_sample_token'] and 'lidar' in x['filename']:\n        lidars.append(x)\nlidars","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"All lidars happen to have the same ego_pose (because they are on the same car?)"},{"metadata":{"trusted":true},"cell_type":"code","source":"{x['ego_pose_token'] for x in lidars}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Now let's load lidar's point data, we'll keep only first 3 columns (point coordinates)"},{"metadata":{"trusted":true},"cell_type":"code","source":"lidars_data = [\n    # here, sorry\n    np.fromfile(os.path.join(DATA_ROOT, x['filename']).replace('/lidar/', '/train_lidar/'), dtype=np.float32)\n    .reshape(-1, 5)[:, :3] for x in lidars]\nlidars_data[0].shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"lidars_data[0]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"lidars_data[0].min(axis=0), lidars_data[0].max(axis=0)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Apparently, we're mostly underground but that's ok.\n\nNow to the interesting stuff - translating all lidars into car coordinate system. For that we need sensor info for each lidar"},{"metadata":{"trusted":true},"cell_type":"code","source":"[calibrated_sensor.index[x['calibrated_sensor_token']] for x in lidars]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Some helpers"},{"metadata":{"trusted":true},"cell_type":"code","source":"def rotate_points(points, rotation, inverse=False):\n    assert points.shape[1] == 3\n    q = Quaternion(rotation)\n    if inverse:\n        q = q.inverse\n    return np.dot(q.rotation_matrix, points.T).T\n    \ndef apply_pose(points, cs):\n    \"\"\" Translate (lidar) points to vehicle coordinates, given a calibrated sensor.\n    \"\"\"\n    points = rotate_points(points, cs['rotation'])\n    points = points + np.array(cs['translation'])\n    return points\n\ndef inverse_apply_pose(points, cs):\n    \"\"\" Reverse of apply_pose (we'll need it later).\n    \"\"\"\n    points = points - np.array(cs['translation'])\n    points = rotate_points(points, np.array(cs['rotation']), inverse=True)\n    return points","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"And let's vizualize all lidars. First let's see what happens if we don't use poses from the lidars"},{"metadata":{"trusted":true},"cell_type":"code","source":"def viz_all_lidars(lidars, lidars_data, clip=50, skip_apply_pose=False):\n    all_points = []\n    all_colors = []\n    for color, points, lidar in zip([[1, 0, 0, 0.5], [0, 1, 0, 0.5], [0, 0, 1, 0.5]], lidars_data, lidars):\n        cs = calibrated_sensor.index[lidar['calibrated_sensor_token']]\n        if not skip_apply_pose:\n            points = apply_pose(points, cs)\n        all_points.append(points)\n        all_colors.append(np.array([color] * len(points)))\n    all_points = np.concatenate(all_points)\n    all_colors = np.concatenate(all_colors)\n    perm = np.random.permutation(len(all_points))\n    all_points = all_points[perm]\n    all_colors = all_colors[perm]\n\n    plt.figure(figsize=(12, 12))\n    plt.axis('equal')\n    plt.grid()\n    plt.scatter(np.clip(all_points[:, 0], -clip, clip), np.clip(all_points[:, 1], -clip, clip), s=1, c=all_colors)\n\nviz_all_lidars(lidars, lidars_data, clip=20, skip_apply_pose=True)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Each lidar has different color, and we permute points to make them more visible (there is probably a better way to do this). Also note that we clip all points which fall outside of a small area, for training you most likely won't do this.\n\nAbove looks wrong, doesn't it? This is because each lidar is rotated and also slightly translated relative to the car, and lidar points are in lidar coordinates. Now let's enable translation to the car coorinatates:"},{"metadata":{"trusted":true},"cell_type":"code","source":"viz_all_lidars(lidars, lidars_data, clip=20)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We see that data from different lidars agrees, so all translations were right, yasss. Note that curvy lines come from the ground and don't need to aling, we look mostly at cars to check alignment.\n\nLet's zoom out a bit:"},{"metadata":{"trusted":true},"cell_type":"code","source":"viz_all_lidars(lidars, lidars_data, clip=50)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Next we load annotated data. We'll  only use annotations available from train.csv because they are more compact and faster to load. They are indexed by sample token."},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(os.path.join(DATA_ROOT, 'train.csv')).set_index('Id')\ntrain_df.loc[my_scene['first_sample_token']]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"A little helper to convert annotation from string"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_annotations(token):\n    annotations = np.array(train_df.loc[token].PredictionString.split()).reshape(-1, 8)\n    return {\n        'point': annotations[:, :3].astype(np.float32),\n        'wlh': annotations[:, 3:6].astype(np.float32),\n        'rotation': annotations[:, 6].astype(np.float32),\n        'cls': np.array(annotations[:, 7]),\n    }\n\nget_annotations(my_scene['first_sample_token']).keys()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Now, this annotations are in global coordinates. Let's translate them into the frame of the car, for that we'll use the ego_pose of the lidars (they are all the same)"},{"metadata":{"trusted":true},"cell_type":"code","source":"ego_pose.index[lidars[0]['ego_pose_token']]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We'll use inverse_apply_pose which we defined above, will show only the car class, and will deal only with point centers"},{"metadata":{"trusted":true},"cell_type":"code","source":"def viz_annotation_centers(token, lidars, clip=50):\n    # translate annotation points to the car frame\n    ego_pose_token, = {x['ego_pose_token'] for x in lidars}\n    ep = ego_pose.index[ego_pose_token]\n    annotations = get_annotations(token)\n    car_points = annotations['point'][annotations['cls'] == 'car']\n    car_points = inverse_apply_pose(car_points, ep)\n    \n    plt.scatter(np.clip(car_points[:, 0], -clip, clip),\n                np.clip(car_points[:, 1], -clip, clip),\n                s=30,\n                color='black')\n    \nviz_all_lidars(lidars, lidars_data, clip=50)\nviz_annotation_centers(my_scene['first_sample_token'], lidars, clip=50)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Car markers look aligned to some blips in the lidar data, so probably above is correct. Note that here again we clipped the data which you'd need to remove for actual training.\n\nLet's zoom in a bit"},{"metadata":{"trusted":true},"cell_type":"code","source":"viz_all_lidars(lidars, lidars_data, clip=20)\nviz_annotation_centers(my_scene['first_sample_token'], lidars, clip=20)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Thanks for reading, if you liked it, please give an upvote."}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":1}