Mayavi_visualization.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "enabling-texas",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mayavi import mlab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "insured-gabriel",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from collections import Counter\n",
    "import imageio\n",
    "import cv2\n",
    "import sys\n",
    "\n",
    "from IPython import display\n",
    "def draw_lidar(pointcloud, color=None, fig=None, bg_color=(0, 0, 0), points_scale=0.3, points_mode=\"sphere\",\n",
    "               points_color=None, color_by_intensity=False, pointcloud_label=False):\n",
    "    points_mode = \"point\"\n",
    "    if fig is None:\n",
    "        fig = mlab.figure(figure=None, bgcolor=bg_color, fgcolor=None, engine=None, size=(1500, 1000))\n",
    "    if color is None:\n",
    "        color = pointcloud[:, 2]\n",
    "    if pointcloud_label:\n",
    "        color = pointcloud[:, 4]\n",
    "    if color_by_intensity:\n",
    "        color = pointcloud[:, 2]\n",
    "    mlab.points3d(pointcloud[:, 0], pointcloud[:, 1], pointcloud[:, 2],\n",
    "                  color, color=points_color, mode=points_mode, colormap=\"gnuplot\",\n",
    "                  scale_factor=points_scale, figure=fig)  # points\n",
    "    mlab.points3d(0, 0, 0, color=(1, 1, 1), mode=\"sphere\", scale_factor=0.2)  # origin\n",
    "    axes = np.array([[2.0, 0.0, 0.0, 0.0], [0.0, 2.0, 0.0, 0.0], [0.0, 0.0, 2.0, 0.0]], dtype=np.float64)\n",
    "    mlab.plot3d([0, axes[0, 0]], [0, axes[0, 1]], [0, axes[0, 2]],\n",
    "                color=(1, 0, 0), tube_radius=None, figure=fig)\n",
    "    mlab.plot3d([0, axes[1, 0]], [0, axes[1, 1]], [0, axes[1, 2]],\n",
    "                color=(0, 1, 0), tube_radius=None, figure=fig)\n",
    "    mlab.plot3d([0, axes[2, 0]], [0, axes[2, 1]], [0, axes[2, 2]],\n",
    "                color=(0, 0, 1), tube_radius=None, figure=fig)\n",
    "\n",
    "    fov = np.array([[20.0, 20.0, 0.0, 0.0], [20.0, -20.0, 0.0, 0.0]], dtype=np.float64)\n",
    "    mlab.plot3d([0, fov[0, 0]], [0, fov[0, 1]], [0, fov[0, 2]],\n",
    "                color=(1, 1, 1), tube_radius=None, line_width=1, figure=fig)\n",
    "    mlab.plot3d([0, fov[1, 0]], [0, fov[1, 1]], [0, fov[1, 2]],\n",
    "                color=(1, 1, 1), tube_radius=None, line_width=1, figure=fig)\n",
    "\n",
    "    x1, y1 = 0, -20\n",
    "    x2, y2 = 40, 20\n",
    "    \n",
    "    mlab.plot3d([x1, x1], [y1, y2], [0, 0],\n",
    "                color=(0.5, 0.5, 0.5), tube_radius=0.1, line_width=1, figure=fig)\n",
    "    mlab.plot3d([x2, x2], [y1, y2], [0, 0],\n",
    "                color=(0.5, 0.5, 0.5), tube_radius=0.1, line_width=1, figure=fig)\n",
    "    mlab.plot3d([x1, x2], [y1, y1], [0, 0],\n",
    "                color=(0.5, 0.5, 0.5), tube_radius=0.1, line_width=1, figure=fig)\n",
    "    mlab.plot3d([x1, x2], [y2, y2], [0, 0],\n",
    "                color=(0.5, 0.5, 0.5), tube_radius=0.1, line_width=1, figure=fig)\n",
    "\n",
    "    mlab.view(azimuth=180, elevation=70, focalpoint=[12.0909996, -1.04700089, -2.03249991],\n",
    "              distance=62.0, figure=fig)\n",
    "    return fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "infinite-prerequisite",
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_gt_boxes3d(gt_boxes3d, fig, color=(1, 1, 1)):\n",
    "    \"\"\"\n",
    "    Draw 3D bounding boxes\n",
    "    Args:\n",
    "        gt_boxes3d: numpy array (3,8) for XYZs of the box corners\n",
    "        fig: figure handler\n",
    "        color: RGB value tuple in range (0,1), box line color\n",
    "    \"\"\"\n",
    "    for k in range(0, 4):\n",
    "        i, j = k, (k + 1) % 4\n",
    "        mlab.plot3d([gt_boxes3d[0, i], gt_boxes3d[0, j]], [gt_boxes3d[1, i], gt_boxes3d[1, j]],\n",
    "                    [gt_boxes3d[2, i], gt_boxes3d[2, j]], tube_radius=None, line_width=2, color=color, figure=fig)\n",
    "\n",
    "        i, j = k + 4, (k + 1) % 4 + 4\n",
    "        mlab.plot3d([gt_boxes3d[0, i], gt_boxes3d[0, j]], [gt_boxes3d[1, i], gt_boxes3d[1, j]],\n",
    "                    [gt_boxes3d[2, i], gt_boxes3d[2, j]], tube_radius=None, line_width=2, color=color, figure=fig)\n",
    "\n",
    "        i, j = k, k + 4\n",
    "        mlab.plot3d([gt_boxes3d[0, i], gt_boxes3d[0, j]], [gt_boxes3d[1, i], gt_boxes3d[1, j]],\n",
    "                    [gt_boxes3d[2, i], gt_boxes3d[2, j]], tube_radius=None, line_width=2, color=color, figure=fig)\n",
    "    return fig "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "entitled-there",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Notebook initialized with ipy backend.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d39d7de5d74f42a2b62a2a51ac9b64d7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x05\\xdc\\x00\\x00\\x03\\xe8\\x08\\x02\\x00\\x00\\x00Kl\\xab\\xb…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mlab.init_notebook('ipy')\n",
    "def display_frame_statistics():\n",
    "    points  = np.fromfile(str('/home/sathish/Documents/dataset/Kitti/training/velodyne/000008.bin'),dtype=np.float32).reshape([-1, 4])\n",
    "    fig=draw_lidar(points, pointcloud_label=False)\n",
    "\n",
    "    display.display(fig)\n",
    "display_frame_statistics()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "anticipated-angle",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cv2\n",
    "import os\n",
    "import mayavi\n",
    "import matplotlib.pyplot as plt\n",
    "import loi as xmlParser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "hungry-announcement",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_tracklets_for_frames(n_frames, xml_path):\n",
    "    \"\"\"\n",
    "    Loads dataset labels also referred to as tracklets, saving them individually for each frame.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    n_frames    : Number of frames in the dataset.\n",
    "    xml_path    : Path to the tracklets XML.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    Tuple of dictionaries with integer keys corresponding to absolute frame numbers and arrays as values. First array\n",
    "    contains coordinates of bounding box vertices for each object in the frame, and the second array contains objects\n",
    "    types as strings.\n",
    "    \"\"\"\n",
    "    tracklets = xmlParser.parseXML(xml_path)    \n",
    "    frame_tracklets = {}\n",
    "    frame_tracklets_types = {}\n",
    "    for i in range(n_frames):\n",
    "        frame_tracklets[i] = []\n",
    "        frame_tracklets_types[i] = []\n",
    "\n",
    "    # loop over tracklets\n",
    "    for i, tracklet in enumerate(tracklets):\n",
    "        # this part is inspired by kitti object development kit matlab code: computeBox3D\n",
    "        h, w, l = tracklet.size\n",
    "        # in velodyne coordinates around zero point and without orientation yet\n",
    "        trackletBox = np.array([\n",
    "            [-l / 2, -l / 2, l / 2, l / 2, -l / 2, -l / 2, l / 2, l / 2],\n",
    "            [w / 2, -w / 2, -w / 2, w / 2, w / 2, -w / 2, -w / 2, w / 2],\n",
    "            [0.0, 0.0, 0.0, 0.0, h, h, h, h]\n",
    "        ])\n",
    "        # loop over all data in tracklet\n",
    "        for translation, rotation, state, occlusion, truncation, amtOcclusion, amtBorders, absoluteFrameNumber in tracklet:\n",
    "            # determine if object is in the image; otherwise continue\n",
    "            if truncation not in (xmlParser.TRUNC_IN_IMAGE, xmlParser.TRUNC_TRUNCATED):\n",
    "                continue\n",
    "            # re-create 3D bounding box in velodyne coordinate system\n",
    "            yaw = rotation[2]  # other rotations are supposedly 0\n",
    "            assert np.abs(rotation[:2]).sum() == 0, 'object rotations other than yaw given!'\n",
    "            rotMat = np.array([\n",
    "                [np.cos(yaw), -np.sin(yaw), 0.0],\n",
    "                [np.sin(yaw), np.cos(yaw), 0.0],\n",
    "                [0.0, 0.0, 1.0]\n",
    "            ])\n",
    "            cornerPosInVelo = np.dot(rotMat, trackletBox) + np.tile(translation, (8, 1)).T\n",
    "            frame_tracklets[absoluteFrameNumber] = frame_tracklets[absoluteFrameNumber] + [cornerPosInVelo]\n",
    "            frame_tracklets_types[absoluteFrameNumber] = frame_tracklets_types[absoluteFrameNumber] + [\n",
    "                tracklet.objectType]\n",
    "\n",
    "    return (frame_tracklets, frame_tracklets_types)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "retired-kazakhstan",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Notebook initialized with ipy backend.\n"
     ]
    }
   ],
   "source": [
    "mlab.init_notebook('ipy')\n",
    "class Object3D:\n",
    "    def __init__(self, label):\n",
    "        data = label.split(' ')\n",
    "\n",
    "        self.classification = data[0]\n",
    "        self.truncation = float(data[1])\n",
    "        self.occlusion = int(data[2])\n",
    "        self.alpha_observation_angle = float(data[3])\n",
    "\n",
    "        self.xmin_left = float(data[4])\n",
    "        self.ymin_top = float(data[5])\n",
    "        self.xmax_right = float(data[6])\n",
    "        self.ymax_bottom = float(data[7])\n",
    "\n",
    "        self.bbox2d = np.array([self.xmin_left, self.ymin_top, self.xmax_right, self.ymax_bottom])\n",
    "\n",
    "        self.height = float(data[8])\n",
    "        self.width = float(data[9])\n",
    "        self.length = float(data[10])\n",
    "\n",
    "        self.obj3d_location_camera_coords = (float(data[11]), float(data[12]), float(data[13]))\n",
    "\n",
    "        self.yaw_rotation_y = float(data[14])\n",
    "\n",
    "\n",
    "class Calibration:\n",
    "    def __init__(self, calib_filename, from_video=False):\n",
    "        if not from_video:\n",
    "            calibration = self.read_calibration_file(calib_filename)\n",
    "        else:\n",
    "            calibration = self.read_calibration_from_video(calib_filename)\n",
    "        self.P = calibration[\"P2\"]\n",
    "        self.P = np.reshape(self.P, [3, 4])\n",
    "        self.V2C = calibration[\"Tr_velo_to_cam\"]\n",
    "        self.V2C = np.reshape(self.V2C, [3, 4])\n",
    "        self.C2V = inverse_rigid_transform(self.V2C)\n",
    "        self.R0 = calibration[\"R0_rect\"]\n",
    "        self.R0 = np.reshape(self.R0, [3, 3])\n",
    "\n",
    "        self.c_u = self.P[0, 2]\n",
    "        self.c_v = self.P[1, 2]\n",
    "        self.f_u = self.P[0, 0]\n",
    "        self.f_v = self.P[1, 1]\n",
    "        self.b_x = self.P[0, 3] / (-self.f_u)\n",
    "        self.b_y = self.P[1, 3] / (-self.f_v)\n",
    "\n",
    "    def read_calibration_file(self, filename):\n",
    "        data = {}\n",
    "        with open(filename, \"r\") as file:\n",
    "            for line in file.readlines():\n",
    "                line = line.rstrip()\n",
    "                if len(line) == 0:\n",
    "                    continue\n",
    "                key, value = line.split(\":\", 1)\n",
    "                try:\n",
    "                    data[key] = np.array([float(x) for x in value.split()])\n",
    "                except ValueError:\n",
    "                    pass\n",
    "        return data\n",
    "\n",
    "    def read_calibration_from_video(self, calib_dir):\n",
    "        data = {}\n",
    "        cam2cam = self.read_calibration_file(os.path.join(calib_dir, \"calib_cam_to_cam.txt\"))\n",
    "        velo2cam = self.read_calibration_file(os.path.join(calib_dir, \"calib_velo_to_cam.txt\"))\n",
    "\n",
    "        Tr_velo_to_cam = np.zeros((3, 4))\n",
    "        Tr_velo_to_cam[0:3, 0:3] = np.reshape(velo2cam[\"R\"], [3, 3])\n",
    "        Tr_velo_to_cam[:, 3] = velo2cam[\"T\"]\n",
    "\n",
    "        data[\"Tr_velo_to_cam\"] = np.reshape(Tr_velo_to_cam, [12])\n",
    "        data[\"R0_rect\"] = cam2cam[\"R_rect_00\"]\n",
    "        data[\"P2\"] = cam2cam[\"P_rect_02\"]\n",
    "\n",
    "        return data\n",
    "\n",
    "    def project_lidar_to_image(self, points3d_lidar):\n",
    "        points3d_rect = self.project_lidar_to_rect(points3d_lidar)\n",
    "        return self.project_rect_to_image(points3d_rect)\n",
    "\n",
    "    def project_lidar_to_rect(self, points3d_lidar):\n",
    "        points3d_ref = self.project_lidar_to_ref(points3d_lidar)\n",
    "        return self.project_ref_to_rect(points3d_ref)\n",
    "\n",
    "    def project_lidar_to_ref(self, points3d_lidar):\n",
    "        points3d_lidar = self.cart2hom(points3d_lidar)\n",
    "        return np.dot(points3d_lidar, np.transpose(self.V2C))\n",
    "\n",
    "    def project_ref_to_rect(self, points3d_ref):\n",
    "        return np.transpose(np.dot(self.R0, np.transpose(points3d_ref)))\n",
    "\n",
    "    def cart2hom(self, points3d):\n",
    "        \"\"\"In: nx3 - Points in Cartesian\n",
    "           Out: nx4 - Points in Homogeneous\n",
    "        \"\"\"\n",
    "        num_points = points3d.shape[0]\n",
    "        points3d_hom = np.hstack((points3d, np.ones((num_points, 1))))\n",
    "        return points3d_hom\n",
    "\n",
    "    def project_rect_to_image(self, points3d_rect):\n",
    "        points3d_rect = self.cart2hom(points3d_rect)\n",
    "        points2d = np.dot(points3d_rect, np.transpose(self.P))\n",
    "        points2d[:, 0] /= points2d[:, 2]\n",
    "        points2d[:, 1] /= points2d[:, 2]\n",
    "        return points2d[:, 0:2]\n",
    "\n",
    "    def project_rect_to_lidar(self, points3d_rect):\n",
    "        points3d_ref = self.project_rect_to_ref(points3d_rect)\n",
    "        return self.project_ref_to_lidar(points3d_ref)\n",
    "\n",
    "    def project_rect_to_ref(self, points3d_rect):\n",
    "        return np.transpose(np.dot(np.linalg.inv(self.R0), np.transpose(points3d_rect)))\n",
    "\n",
    "    def project_ref_to_lidar(self, points3d_ref):\n",
    "        points3d_ref = self.cart2hom(points3d_ref)\n",
    "        return np.dot(points3d_ref, np.transpose(self.C2V))\n",
    "\n",
    "\n",
    "class KITTIVideo:\n",
    "    def __init__(self, img_dir, lidar_dir, calib_dir):\n",
    "        self.calib = Calibration(calib_dir, from_video=True)\n",
    "        self.img_dir = img_dir\n",
    "        self.lidar_dir = lidar_dir\n",
    "\n",
    "        self.img_filenames = sorted([os.path.join(img_dir, filename) for filename in os.listdir(img_dir)])\n",
    "        self.lidar_filenames = sorted([os.path.join(lidar_dir, filename) for filename in os.listdir(lidar_dir)])\n",
    "\n",
    "        self.num_samples = len(self.img_filenames)\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.num_samples\n",
    "\n",
    "    def get_image(self, idx):\n",
    "        return load_image(self.img_filenames[idx])\n",
    "\n",
    "    def get_lidar(self, idx):\n",
    "        return get_velodyne_scan_points(self.lidar_filenames[idx])\n",
    "\n",
    "    def get_calibration(self):\n",
    "        return self.calib\n",
    "\n",
    "\n",
    "def inverse_rigid_transform(tr):\n",
    "    inv_tr = np.zeros_like(tr)\n",
    "    inv_tr[0:3, 0:3] = np.transpose(tr[0:3, 0:3])\n",
    "    inv_tr[0:3, 3] = np.dot(-np.transpose(tr[0:3, 0:3]), tr[0:3, 3])\n",
    "    return inv_tr\n",
    "\n",
    "\n",
    "def load_image(image_filename):\n",
    "    return cv2.imread(image_filename)\n",
    "\n",
    "\n",
    "def read_label(label_file_name):\n",
    "    objects = []\n",
    "    with open(label_file_name, 'r') as lf:\n",
    "        for line in lf.readlines():\n",
    "            objects.append(Object3D(line.rstrip()))\n",
    "    return objects\n",
    "\n",
    "\n",
    "def rotate_y(p):\n",
    "    cp = np.cos(p)\n",
    "    sp = np.sin(p)\n",
    "\n",
    "    return np.array([[cp, 0, sp], [0, 1, 0], [-sp, 0, cp]])\n",
    "\n",
    "\n",
    "def project_3d_to_image2d(points3d, P):\n",
    "    n = points3d.shape[0]\n",
    "    points3d_extended = np.hstack((points3d, np.ones((n, 1))))\n",
    "    points2d = np.dot(points3d_extended, np.transpose(P))\n",
    "    points2d[:, 0] /= points2d[:, 2]\n",
    "    points2d[:, 1] /= points2d[:, 2]\n",
    "    return points2d[:, 0:2]\n",
    "\n",
    "\n",
    "def draw_projected_bbox3d(image, vertices, color, thickness):\n",
    "    vertices = vertices.astype(np.int32)\n",
    "    for idx in range(0, 4):\n",
    "        i, j = idx, (idx + 1) % 4\n",
    "        cv2.line(image, (vertices[i, 0], vertices[i, 1]),\n",
    "                 (vertices[j, 0], vertices[j, 1]),\n",
    "                 color, thickness)\n",
    "        i, j = idx + 4, (idx + 1) % 4 + 4\n",
    "        cv2.line(image, (vertices[i, 0], vertices[i, 1]),\n",
    "                 (vertices[j, 0], vertices[j, 1]),\n",
    "                 color, thickness)\n",
    "        i, j = idx, idx + 4\n",
    "        cv2.line(image, (vertices[i, 0], vertices[i, 1]),\n",
    "                 (vertices[j, 0], vertices[j, 1]),\n",
    "                 color, thickness)\n",
    "    return image\n",
    "\n",
    "\n",
    "def compute_bbox3d(obj, P):\n",
    "    R = rotate_y(obj.yaw_rotation_y)\n",
    "\n",
    "    l = obj.length\n",
    "    w = obj.width\n",
    "    h = obj.height\n",
    "\n",
    "    # 3D BBox Corners\n",
    "    x_corners = [l / 2, l / 2, -l / 2, -l / 2, l / 2, l / 2, -l / 2, -l / 2]\n",
    "    y_corners = [0, 0, 0, 0, -h, -h, -h, -h]\n",
    "    z_corners = [w / 2, -w / 2, -w / 2, w / 2, w / 2, -w / 2, -w / 2, w / 2]\n",
    "\n",
    "    # rotate and translate 3D BBox\n",
    "    corners3d = np.dot(R, np.vstack([x_corners, y_corners, z_corners]))\n",
    "\n",
    "    corners3d[0, :] = corners3d[0, :] + obj.obj3d_location_camera_coords[0]\n",
    "    corners3d[1, :] = corners3d[1, :] + obj.obj3d_location_camera_coords[1]\n",
    "    corners3d[2, :] = corners3d[2, :] + obj.obj3d_location_camera_coords[2]\n",
    "\n",
    "    if np.any(corners3d[2, :] < 0.1): # Draw 3D BBoxes for objects which are infront of Camera\n",
    "        corners2d = None\n",
    "        return corners2d, np.transpose(corners3d)\n",
    "\n",
    "    # Project 3D BBox onto Image plane\n",
    "    corners2d = project_3d_to_image2d(np.transpose(corners3d), P)\n",
    "    return corners2d, np.transpose(corners3d)\n",
    "\n",
    "\n",
    "def get_velodyne_scan_points(lidar_filename, n_vec=4):\n",
    "    lidar_scanpoints = np.fromfile(lidar_filename, np.float32)\n",
    "    lidar_scanpoints = lidar_scanpoints.reshape((-1, n_vec))\n",
    "    return lidar_scanpoints\n",
    "def draw_groundtruth_3dbboxes(groundtruth_3dbboxes, fig, color=(1, 1, 1), line_width=1, draw_text=True,\n",
    "                              text_scale=(1, 1, 1), color_list=None, label=\"\"):\n",
    "    num_bboxes = len(groundtruth_3dbboxes)\n",
    "    for idx in range(num_bboxes):\n",
    "        bbox = groundtruth_3dbboxes[idx]\n",
    "        if color_list is not None:\n",
    "            color = color_list[idx]\n",
    "        if draw_text:\n",
    "            mlab.text3d(bbox[4, 0], bbox[4, 1], bbox[4, 2], label, scale=text_scale, color=color, figure=fig)\n",
    "        for k in range(0, 4):\n",
    "            i, j = k, (k + 1) % 4\n",
    "            mlab.plot3d([bbox[i, 0], bbox[j, 0]],\n",
    "                        [bbox[i, 1], bbox[j, 1]],\n",
    "                        [bbox[i, 2], bbox[j, 2]],\n",
    "                        color=color,\n",
    "                        tube_radius=None,\n",
    "                        line_width=line_width,\n",
    "                        figure=fig)\n",
    "            i, j = k + 4, (k + 1) % 4 + 4\n",
    "            mlab.plot3d([bbox[i, 0], bbox[j, 0]],\n",
    "                        [bbox[i, 1], bbox[j, 1]],\n",
    "                        [bbox[i, 2], bbox[j, 2]],\n",
    "                        color=color,\n",
    "                        tube_radius=None,\n",
    "                        line_width=line_width,\n",
    "                        figure=fig)\n",
    "            i, j = k, k + 4\n",
    "            mlab.plot3d([bbox[i, 0], bbox[j, 0]],\n",
    "                        [bbox[i, 1], bbox[j, 1]],\n",
    "                        [bbox[i, 2], bbox[j, 2]],\n",
    "                        color=color,\n",
    "                        tube_radius=None,\n",
    "                        line_width=line_width,\n",
    "                        figure=fig)\n",
    "    return fig\n",
    "\n",
    "\n",
    "def draw_lidar_simple(pointcloud, color=None):\n",
    "    figure = mlab.figure(figure=None, bgcolor=(0, 0, 0), fgcolor=None, engine=None, size=(1500, 1000))\n",
    "    if color is None:\n",
    "        color = pointcloud[:, 2]\n",
    "    mlab.points3d(pointcloud[:, 0], pointcloud[:, 1], pointcloud[:, 2],\n",
    "                  color, color=None, mode=\"point\", colormap=\"gnuplot\", scale_factor=1, figure=figure)  # points\n",
    "    mlab.points3d(0, 0, 0, color=(1, 1, 1), mode=\"sphere\", scale_factor=0.2)  # origin\n",
    "    axes = np.array([[2.0, 0.0, 0.0, 0.0], [0.0, 2.0, 0.0, 0.0], [0.0, 0.0, 2.0, 0.0]], dtype=np.float64)\n",
    "    mlab.plot3d([0, axes[0, 0]], [0, axes[0, 1]], [0, axes[0, 2]],\n",
    "                color=(1, 0, 0), tube_radius=None, figure=figure)\n",
    "    mlab.plot3d([0, axes[1, 0]], [0, axes[1, 1]], [0, axes[1, 2]],\n",
    "                color=(0, 1, 0), tube_radius=None, figure=figure)\n",
    "    mlab.plot3d([0, axes[2, 0]], [0, axes[2, 1]], [0, axes[2, 2]],\n",
    "                color=(0, 0, 1), tube_radius=None, figure=figure)\n",
    "    mlab.view(azimuth=180, elevation=70, focalpoint=[12.0909996, -1.04700089, -2.03249991],\n",
    "              distance=62.0, figure=figure)\n",
    "    return figure\n",
    "\n",
    "\n",
    "def draw_lidar(pointcloud, color=None, fig=None, bg_color=(0, 0, 0), points_scale=0.3, points_mode=\"sphere\",\n",
    "               points_color=None, color_by_intensity=False, pointcloud_label=False):\n",
    "    points_mode = \"point\"\n",
    "    if fig is None:\n",
    "        fig = mlab.figure(figure=None, bgcolor=bg_color, fgcolor=None, engine=None, size=(1500, 1000))\n",
    "    if color is None:\n",
    "        color = pointcloud[:, 2]\n",
    "    if pointcloud_label:\n",
    "        color = pointcloud[:, 4]\n",
    "    if color_by_intensity:\n",
    "        color = pointcloud[:, 2]\n",
    "    mlab.points3d(pointcloud[:, 0], pointcloud[:, 1], pointcloud[:, 2],\n",
    "                  color, color=points_color, mode=points_mode, colormap=\"gnuplot\",\n",
    "                  scale_factor=points_scale, figure=fig)  # points\n",
    "    mlab.points3d(0, 0, 0, color=(1, 1, 1), mode=\"sphere\", scale_factor=0.2)  # origin\n",
    "    axes = np.array([[2.0, 0.0, 0.0, 0.0], [0.0, 2.0, 0.0, 0.0], [0.0, 0.0, 2.0, 0.0]], dtype=np.float64)\n",
    "    mlab.plot3d([0, axes[0, 0]], [0, axes[0, 1]], [0, axes[0, 2]],\n",
    "                color=(1, 0, 0), tube_radius=None, figure=fig)\n",
    "    mlab.plot3d([0, axes[1, 0]], [0, axes[1, 1]], [0, axes[1, 2]],\n",
    "                color=(0, 1, 0), tube_radius=None, figure=fig)\n",
    "    mlab.plot3d([0, axes[2, 0]], [0, axes[2, 1]], [0, axes[2, 2]],\n",
    "                color=(0, 0, 1), tube_radius=None, figure=fig)\n",
    "\n",
    "    fov = np.array([[20.0, 20.0, 0.0, 0.0], [20.0, -20.0, 0.0, 0.0]], dtype=np.float64)\n",
    "    mlab.plot3d([0, fov[0, 0]], [0, fov[0, 1]], [0, fov[0, 2]],\n",
    "                color=(1, 1, 1), tube_radius=None, line_width=1, figure=fig)\n",
    "    mlab.plot3d([0, fov[1, 0]], [0, fov[1, 1]], [0, fov[1, 2]],\n",
    "                color=(1, 1, 1), tube_radius=None, line_width=1, figure=fig)\n",
    "\n",
    "    x1, y1 = 0, -20\n",
    "    x2, y2 = 40, 20\n",
    "    \n",
    "    mlab.plot3d([x1, x1], [y1, y2], [0, 0],\n",
    "                color=(0.5, 0.5, 0.5), tube_radius=0.1, line_width=1, figure=fig)\n",
    "    mlab.plot3d([x2, x2], [y1, y2], [0, 0],\n",
    "                color=(0.5, 0.5, 0.5), tube_radius=0.1, line_width=1, figure=fig)\n",
    "    mlab.plot3d([x1, x2], [y1, y1], [0, 0],\n",
    "                color=(0.5, 0.5, 0.5), tube_radius=0.1, line_width=1, figure=fig)\n",
    "    mlab.plot3d([x1, x2], [y2, y2], [0, 0],\n",
    "                color=(0.5, 0.5, 0.5), tube_radius=0.1, line_width=1, figure=fig)\n",
    "\n",
    "    mlab.view(azimuth=180, elevation=70, focalpoint=[12.0909996, -1.04700089, -2.03249991],\n",
    "              distance=62.0, figure=fig)\n",
    "    return fig\n",
    "\n",
    "\n",
    "\n",
    "class KITTIObject:\n",
    "    def __init__(self, root_dir, dataset_type, ):\n",
    "        self.root_dir = root_dir\n",
    "        self.dataset_type = dataset_type\n",
    "\n",
    "        self.dataset_dir = os.path.join(self.root_dir, self.dataset_type)\n",
    "\n",
    "        self.num_samples = len(os.listdir(self.dataset_dir))\n",
    "\n",
    "        self.image_dir = os.path.join(self.dataset_dir, \"image_2\")\n",
    "        self.label_dir = os.path.join(self.dataset_dir, \"label_2\")\n",
    "        self.calib_dir = os.path.join(self.dataset_dir, \"calib\")\n",
    "        self.lidar_dir = os.path.join(self.dataset_dir, \"velodyne\")\n",
    "\n",
    "    def get_label_objects(self, idx):\n",
    "        if idx > self.num_samples:\n",
    "            return None\n",
    "        label_filename = os.path.join(self.label_dir, \"%06d.txt\" % idx)\n",
    "        return read_label(label_filename)\n",
    "\n",
    "    def get_lidar_data(self, idx, n_vec=4):\n",
    "        lidar_filename = os.path.join(self.lidar_dir, \"%06d.bin\" % idx)\n",
    "        return get_velodyne_scan_points(lidar_filename, n_vec)\n",
    "\n",
    "    def get_calibration(self, idx):\n",
    "        calib_filename = os.path.join(self.calib_dir, \"%06d.txt\" % idx)\n",
    "        return Calibration(calib_filename)\n",
    "\n",
    "    def get_image(self, idx):\n",
    "        image_filename = os.path.join(self.image_dir, \"%06d.png\" % idx)\n",
    "        return load_image(image_filename)\n",
    "\n",
    "\n",
    "def get_lidar_index_in_image_fov(pcd, calib, xmin, ymin, xmax, ymax, clip_distance=2.0):\n",
    "    points2d = calib.project_lidar_to_image(pcd)\n",
    "    fov_idxs = ((points2d[:, 0] < xmax) & (points2d[:, 0] >= xmin) &\n",
    "                (points2d[:, 1] < ymax) & (points2d[:, 1] >= ymin))\n",
    "    fov_idxs = fov_idxs & (pcd[:, 0] > clip_distance)\n",
    "    return fov_idxs\n",
    "\n",
    "\n",
    "def show_lidar_data(point_cloud_data, objects, calibration, figure,\n",
    "                    img_fov=False, img_width=None, img_height=None, cam_img=None, pc_label=False, save=False):\n",
    "    # print((\"All point num: \",\n",
    "    point_cloud_data.shape[0]\n",
    "    if img_fov:\n",
    "        pcd_index = get_lidar_index_in_image_fov(point_cloud_data[:, :3], calibration, 0, 0, img_width, img_height)\n",
    "        point_cloud_data = point_cloud_data[pcd_index, :]\n",
    "        # print((\"FOV point num: \", point_cloud_data.shape))\n",
    "    draw_lidar(point_cloud_data, fig=figure, pointcloud_label=pc_label)\n",
    "\n",
    "    color = (0, 1, 0)\n",
    "\n",
    "    if objects is not None:\n",
    "        for obj in objects:\n",
    "            if obj.classification == \"DontCare\":\n",
    "                continue\n",
    "            _, bbox3d = compute_bbox3d(obj, calibration.P)\n",
    "            bbox3d_velodyne = calibration.project_rect_to_lidar(bbox3d)\n",
    "            draw_groundtruth_3dbboxes([bbox3d_velodyne], fig=figure, color=color, label=obj.classification)\n",
    "        mlab.show(1)\n",
    "\n",
    "\n",
    "def get_lidar_data_in_image_fov(pcd, calib, xmin, ymin, xmax, ymax, clip_distance=2.0):\n",
    "    points2d = calib.project_lidar_to_image(pcd)\n",
    "    fov_idxs = ((points2d[:, 0] < xmax) & (points2d[:, 0] >= xmin) &\n",
    "                (points2d[:, 1] < ymax) & (points2d[:, 1] >= ymin))\n",
    "    fov_idxs = fov_idxs & (pcd[:, 0] > clip_distance)\n",
    "    imgfov_pcd = pcd[fov_idxs, :]\n",
    "\n",
    "    return imgfov_pcd, points2d, fov_idxs\n",
    "\n",
    "\n",
    "def show_lidar_data_on_image(point_cloud_data, img, calib, img_width, img_height):\n",
    "    img = np.copy(img)\n",
    "    imgfov_pcd, points2d, fov_indices = \\\n",
    "        get_lidar_data_in_image_fov(point_cloud_data, calib, 0, 0, img_width, img_height)\n",
    "    imgfov_points2d = points2d[fov_indices, :]\n",
    "    imgfov_pc_rect = calib.project_lidar_to_rect(imgfov_pcd)\n",
    "\n",
    "    cmap = plt.cm.get_cmap(\"hsv\", 256)\n",
    "    cmap = np.array([cmap(i) for i in range(256)])[:, :3] * 255\n",
    "\n",
    "    for i in range(imgfov_points2d.shape[0]):\n",
    "        depth = imgfov_pc_rect[i, 2]\n",
    "        color = cmap[int(640.0 / depth), :]\n",
    "        cv2.circle(img, (int(np.round(imgfov_points2d[i, 0])), int(np.round(imgfov_points2d[i, 1]))),\n",
    "                   2, color=tuple(color), thickness=-1)\n",
    "    return img\n",
    "\n",
    "\n",
    "def show_bboxes_on_image(image, objects, calib):\n",
    "    img_2d = image.copy()\n",
    "    img_3d = image.copy()\n",
    "\n",
    "    if objects is not None:\n",
    "        for obj in objects:\n",
    "            if obj.classification == \"DontCare\":\n",
    "                continue\n",
    "            cv2.rectangle(img_2d, (int(obj.xmin_left), int(obj.ymin_top)),\n",
    "                          (int(obj.xmax_right), int(obj.ymax_bottom)),\n",
    "                          (0, 255, 0), 3)\n",
    "            bbox3d_points2d, _ = compute_bbox3d(obj, calib.P)\n",
    "            img_3d = draw_projected_bbox3d(img_3d, bbox3d_points2d, (0, 255, 0), 3)\n",
    "    return img_2d, img_3d\n",
    "\n",
    "\n",
    "def visualize_video(img_dir, lidar_dir, calib_dir, output_dir):\n",
    "    dataset = KITTIVideo(img_dir, lidar_dir, calib_dir)\n",
    "\n",
    "    fig = None\n",
    "\n",
    "    for i in range(len(dataset)):\n",
    "        img = dataset.get_image(i)\n",
    "        pcd = dataset.get_lidar(i)\n",
    "        # cv2.imshow(\"Video\", img)\n",
    "        fig = draw_lidar(pcd, pointcloud_label=False)\n",
    "\n",
    "        mlab.savefig(os.path.join(output_dir, \"P%s.png\" % i))\n",
    "\n",
    "    return fig  # return last figure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "signed-saver",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "root_path = r\"/home/sathish/Downloads/2011_09_26/2011_09_26/2011_09_26\"\n",
    "\n",
    "lidar_path = r\"/home/sathish/Downloads/2011_09_26/2011_09_26/2011_09_26/2011_09_26_drive_0001_sync/velodyne_points/data\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "executive-skill",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d9e812a8d6e44ba29d9ba4bdc7e6faaa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x03 \\x00\\x00\\x01\\xf4\\x08\\x02\\x00\\x00\\x00J\"\\xde/\\x00\\…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython import display\n",
    "dataset = KITTIObject(r'/home/sathish/Documents/dataset/Kitti', 'training')\n",
    "data_idx = 1\n",
    "idx=1\n",
    "objects = dataset.get_label_objects(idx)\n",
    "pcd = dataset.get_lidar_data(idx)\n",
    "calib = dataset.get_calibration(idx)\n",
    "img = dataset.get_image(idx)\n",
    "img_height, img_width, _ = img.shape\n",
    "fig_3d = mlab.figure(bgcolor=(0, 0, 0), size=(800, 500))\n",
    "show_lidar_data(pcd, objects, calib, fig_3d, False, img_width, img_height,save=True)\n",
    "display.display(fig_3d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eleven-hypothesis",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_dir = r\"/home/sathish/Downloads/2011_09_26/2011_09_26/2011_09_26/2011_09_26_drive_0001_sync/image_02/data\"\n",
    "lidar_dir = r\"/home/sathish/Downloads/2011_09_26/2011_09_26/2011_09_26/2011_09_26_drive_0001_sync/velodyne_points/data\"\n",
    "calib_dir = r\"/home/sathish/Downloads/2011_09_26/2011_09_26/2011_09_26\"\n",
    "output_path = r\"/home/sathish/Output\"\n",
    "fig_3d = visualize_video(image_dir, lidar_dir, calib_dir, output_path)\n",
    "display.display(fig_3d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "narrow-profit",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "metric-wiring",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "outdoor-philip",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "alien-shelter",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "looking-universal",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parsing tracklet file /home/sathish/Downloads/2011_09_26/2011_09_26/2011_09_26_drive_0048_sync/tracklet_labels.xml\n",
      "File contains 8 tracklets\n",
      "Loaded 8 tracklets.\n"
     ]
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "theoretical-truth",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([[ 3.61872747,  3.57328205,  7.34074598,  7.3861914 ,  3.61872747,\n",
       "          3.57328205,  7.34074598,  7.3861914 ],\n",
       "        [-1.78385112, -3.4259711 , -3.5302351 , -1.88811512, -1.78385112,\n",
       "         -3.4259711 , -3.5302351 , -1.88811512],\n",
       "        [-1.68103695, -1.68103695, -1.68103695, -1.68103695, -0.26733065,\n",
       "         -0.26733065, -0.26733065, -0.26733065]]),\n",
       " array([[32.64796338, 32.66172576, 28.68707403, 28.67331166, 32.64796338,\n",
       "         32.66172576, 28.68707403, 28.67331166],\n",
       "        [ 3.64524225,  5.40046889,  5.43163333,  3.67640668,  3.64524225,\n",
       "          5.40046889,  5.43163333,  3.67640668],\n",
       "        [-1.65444398, -1.65444398, -1.65444398, -1.65444398, -0.22131778,\n",
       "         -0.22131778, -0.22131778, -0.22131778]]),\n",
       " array([[12.00426498, 12.04482511,  8.10180148,  8.06124135, 12.00426498,\n",
       "         12.04482511,  8.10180148,  8.06124135],\n",
       "        [ 4.04690329,  5.65158057,  5.7512452 ,  4.14656792,  4.04690329,\n",
       "          5.65158057,  5.7512452 ,  4.14656792],\n",
       "        [-1.8505518 , -1.8505518 , -1.8505518 , -1.8505518 , -0.3641955 ,\n",
       "         -0.3641955 , -0.3641955 , -0.3641955 ]]),\n",
       " array([[16.95901834, 17.02631093, 13.13312039, 13.0658278 , 16.95901834,\n",
       "         17.02631093, 13.13312039, 13.0658278 ],\n",
       "        [ 3.91088232,  5.57293131,  5.73055775,  4.06850875,  3.91088232,\n",
       "          5.57293131,  5.73055775,  4.06850875],\n",
       "        [-1.796495  , -1.796495  , -1.796495  , -1.796495  , -0.3151032 ,\n",
       "         -0.3151032 , -0.3151032 , -0.3151032 ]])]"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tracklet_rects[21]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "located-switch",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}