extraction.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "from azure.storage.blob import ContainerClient\n",
    "import io\n",
    "import time\n",
    "from PIL import Image\n",
    "import datetime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "523\n"
     ]
    }
   ],
   "source": [
    "file_path_color_left = \"data/00\\image_left/\"\n",
    "file_path_color_right = \"data/00\\image_right/\"\n",
    "\n",
    "left_images = os.listdir(file_path_color_left)\n",
    "right_images = os.listdir(file_path_color_right)\n",
    "\n",
    "print(len(left_images))\n",
    "# plt.figure(figsize=(12,4))\n",
    "# plt.imshow(cv2.imread(file_path_color_left + left_images[165],0)) \n",
    "#turning it into greyscale because we don't really need color unless we're interested in matching moving objects/features. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pixel_coord_np(height,width):\n",
    "    \"\"\"\n",
    "    Pixel in homogenous coordinate\n",
    "    Returns:\n",
    "        Pixel coordinate:       [3, width * height]\n",
    "    \"\"\"\n",
    "    x = np.linspace(0, width - 1, width).astype(np.float32)\n",
    "    y = np.linspace(0, height - 1, height).astype(np.float32)\n",
    "    [x, y] = np.meshgrid(x, y)\n",
    "    return np.vstack((x.flatten(), y.flatten(), np.ones_like(x.flatten())))\n",
    "\n",
    "\n",
    "def intrinsic_from_fov(height, width, fov=90):\n",
    "    \"\"\"\n",
    "    Basic Pinhole Camera Model\n",
    "    intrinsic params from fov and sensor width and height in pixels\n",
    "    Returns:\n",
    "        K:      [4, 4]\n",
    "    \"\"\"\n",
    "    px, py = (width / 2, height / 2)\n",
    "    hfov = fov / 360. * 2. * np.pi\n",
    "    fx = width / (2. * np.tan(hfov / 2.))\n",
    "\n",
    "    vfov = 2. * np.arctan(np.tan(hfov / 2) * height / width)\n",
    "    fy = height / (2. * np.tan(vfov / 2.))\n",
    "\n",
    "    return np.array([[fx, 0, px, 0.],\n",
    "                     [0, fy, py, 0.],\n",
    "                     [0, 0, 1., 0.],\n",
    "                     [0., 0., 0., 1.]])\n",
    "\n",
    "def quaternion_rotation_matrix(Q):\n",
    "    \"\"\"\n",
    "    Covert a quaternion into a full projection matrix.\n",
    " \n",
    "    Input\n",
    "    :param Q: A 7 element array representing translation and the quaternion (q0,q1,q2,q3) \n",
    " \n",
    "    Output\n",
    "    :return: A 3x4 element matrix representing the full projection matrix. \n",
    "             This projection matrix converts a point in the local reference \n",
    "             frame to a point in the global reference frame.\n",
    "    \"\"\"\n",
    "    # Extract the values from Q\n",
    "    t0 = Q[0]\n",
    "    t1 = Q[1]\n",
    "    t2 = Q[2]\n",
    "    q0 = Q[3]\n",
    "    q1 = Q[4]\n",
    "    q2 = Q[5]\n",
    "    q3 = Q[6]\n",
    "     \n",
    "    # First row of the rotation matrix\n",
    "    r00 = 2 * (q0 * q0 + q1 * q1) - 1\n",
    "    r01 = 2 * (q1 * q2 - q0 * q3)\n",
    "    r02 = 2 * (q1 * q3 + q0 * q2)\n",
    "     \n",
    "    # Second row of the rotation matrix\n",
    "    r10 = 2 * (q1 * q2 + q0 * q3)\n",
    "    r11 = 2 * (q0 * q0 + q2 * q2) - 1\n",
    "    r12 = 2 * (q2 * q3 - q0 * q1)\n",
    "     \n",
    "    # Third row of the rotation matrix\n",
    "    r20 = 2 * (q1 * q3 - q0 * q2)\n",
    "    r21 = 2 * (q2 * q3 + q0 * q1)\n",
    "    r22 = 2 * (q0 * q0 + q3 * q3) - 1\n",
    "     \n",
    "    # 3x4 projection matrix\n",
    "    pro_matrix = np.array([[r00, r01, r02, t0],\n",
    "                           [r10, r11, r12, t1],\n",
    "                           [r20, r21, r22, t2]])\n",
    "                            \n",
    "    return pro_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "523\n"
     ]
    }
   ],
   "source": [
    "class Dataset_Handler():\n",
    "    def __init__(self, sequence, lidar=False, progress_bar=True, low_memory=True):\n",
    "        \n",
    "        \n",
    "        # This will tell our odometry function if handler contains lidar info\n",
    "        self.lidar = lidar\n",
    "        # This will tell odometry functin how to access data from this object\n",
    "        self.low_memory = low_memory\n",
    "        \n",
    "        # Set file paths and get ground truth poses\n",
    "        self.seq_dir = \"data\\{}/\".format(sequence)\n",
    "        self.poses_dir = \"data\\{}.csv\".format(sequence)\n",
    "        self.depth_dir = \"data\\{}/\".format(sequence)\n",
    "\n",
    "        poses = pd.read_csv(self.poses_dir, header=None)\n",
    "        \n",
    "        # Get names of files to iterate through\n",
    "        self.left_image_files = os.listdir(self.seq_dir + 'image_left')\n",
    "        self.right_image_files = os.listdir(self.seq_dir + 'image_right')\n",
    "        self.left_depth_files = os.listdir(self.depth_dir + 'depth_left')\n",
    "        self.right_depth_files = os.listdir(self.depth_dir + 'depth_right')\n",
    "        \n",
    "        # self.velodyne_files = os.listdir(self.seq_dir + 'flow')\n",
    "        self.num_frames = len(self.left_image_files)\n",
    "        print(self.num_frames)\n",
    "        # self.lidar_path = self.seq_dir + 'flow/'\n",
    "        self.first_image_left = cv2.imread(self.seq_dir + 'image_left/' \n",
    "                                               + self.left_image_files[0])\n",
    "        height = 480\n",
    "        width = 640\n",
    "        K = intrinsic_from_fov(height,width)\n",
    "        self.P0 = np.array(K[:3,:4])\n",
    "        self.P1 = np.array(K[:3,:4])\n",
    "        \n",
    "        \n",
    "        # Get calibration details for scene\n",
    "        # calib = pd.read_csv(self.seq_dir + 'calib.txt', delimiter=' ', header=None, index_col=0)\n",
    "        # self.P0 = np.array(calib.loc['P0:']).reshape((3,4))\n",
    "        # self.P1 = np.array(calib.loc['P1:']).reshape((3,4))\n",
    "        # self.P2 = np.array(calib.loc['P2:']).reshape((3,4)) #RGB cams\n",
    "        # self.P3 = np.array(calib.loc['P3:']).reshape((3,4)) #RGB cams\n",
    "        # This is the transformation matrix for LIDAR\n",
    "        # self.Tr = np.array(calib.loc['Tr:']).reshape((3,4))\n",
    "        \n",
    "        # Get times and ground truth poses\n",
    "        self.times = np.array(pd.read_csv(self.seq_dir + 'times.txt', \n",
    "                                          delimiter=' ', \n",
    "                                          header=None))\n",
    "        self.gt = np.zeros((len(poses), 3, 4))\n",
    "        for i in range(len(poses)):\n",
    "            self.gt[i] = np.array(quaternion_rotation_matrix(poses.iloc[i])).reshape((3, 4))\n",
    "        \n",
    "        # Get images and lidar loaded\n",
    "        if self.low_memory:\n",
    "            # Will use generators to provide data sequentially to save RAM\n",
    "            # Use class method to set up generators\n",
    "            self.reset_frames()\n",
    "            # Store original frame to memory for testing functions\n",
    "            self.first_image_left = cv2.imread(self.seq_dir + 'image_left/' \n",
    "                                               + self.left_image_files[0])\n",
    "            self.first_image_right = cv2.imread(self.seq_dir + 'image_right/' \n",
    "                                               + self.right_image_files[0])\n",
    "            self.second_image_left = cv2.imread(self.seq_dir + 'image_left/' \n",
    "                                               + self.left_image_files[1])\n",
    "            self.first_depth_left = np.load(self.depth_dir + 'depth_left/'\n",
    "                                                + self.left_depth_files[0])\n",
    "            self.first_depth_right = np.load(self.depth_dir + 'depth_right/'\n",
    "                                                + self.right_depth_files[0])\n",
    "            self.second_depth_left = np.load(self.depth_dir + 'depth_left/'\n",
    "                                                + self.left_depth_files[1])\n",
    "                                                \n",
    "            if self.lidar:\n",
    "                self.first_pointcloud = np.fromfile(self.lidar_path + self.velodyne_files[0],\n",
    "                                                    dtype=np.float32, \n",
    "                                                    count=-1).reshape((-1, 4))\n",
    "            self.imheight = height\n",
    "            self.imwidth = width\n",
    "            \n",
    "        else:\n",
    "            # If RAM is not a concern (>32GB), pass low_memory=False\n",
    "            if progress_bar:\n",
    "                import progressbar\n",
    "                bar = progressbar.ProgressBar(max_value=self.num_frames)\n",
    "            self.images_left = []\n",
    "            self.images_right = []\n",
    "            self.depths_left = []\n",
    "            self.depths_right =[]\n",
    "            self.pointclouds = []\n",
    "            for i, name_left in enumerate(self.left_image_files):\n",
    "                name_right = self.right_image_files[i]\n",
    "                d_left = self.left_depth_files\n",
    "                d_right = self.right_depth_files\n",
    "                self.images_left.append(cv2.imread(self.seq_dir + 'image_left/' + name_left))\n",
    "                self.images_right.append(cv2.imread(self.seq_dir + 'image_right/' + name_right))\n",
    "                self.depths_left.append(np.load(self.depth_dir + 'depth_left/' + d_left))\n",
    "                self.depths_right.append(np.load(self.depth_dir + 'depth_right/' + d_right))\n",
    "                if self.lidar:\n",
    "                    pointcloud = np.fromfile(self.lidar_path + self.velodyne_files[i], \n",
    "                                             dtype=np.float32, \n",
    "                                             count=-1).reshape([-1,4])\n",
    "                    self.pointclouds.append(pointcloud)\n",
    "                if progress_bar:\n",
    "                    bar.update(i+1)\n",
    "                \n",
    "            self.imheight = self.images_left[0].shape[0]\n",
    "            self.imwidth = self.images_left[0].shape[1]\n",
    "            # Keep consistent instance variable names as when using low_memory\n",
    "            self.first_image_left = self.images_left[0]\n",
    "            self.first_image_right = self.images_right[0]\n",
    "            self.second_image_left = self.images_left[1]\n",
    "            if self.lidar:\n",
    "                self.first_pointcloud = self.pointclouds[0]\n",
    "            \n",
    "    def reset_frames(self):\n",
    "        # Resets all generators to the first frame of the sequence\n",
    "        self.images_left = (cv2.imread(self.seq_dir + 'image_left/' + name_left)\n",
    "                            for name_left in self.left_image_files)\n",
    "        self.images_right = (cv2.imread(self.seq_dir + 'image_right/' + name_right)\n",
    "                            for name_right in self.right_image_files)\n",
    "        self.depths_left = (np.load(self.depth_dir+'depth_left/' + d_left)\n",
    "                            for d_left in self.left_depth_files)\n",
    "        self.depths_right = (np.load(self.depth_dir+'depth_right/' + d_right)\n",
    "                            for d_right in self.right_depth_files)                    \n",
    "        if self.lidar:\n",
    "            self.pointclouds = (np.fromfile(self.lidar_path + velodyne_file, \n",
    "                                            dtype=np.float32, \n",
    "                                            count=-1).reshape((-1, 4))\n",
    "                                for velodyne_file in self.velodyne_files)\n",
    "        pass\n",
    "handler = Dataset_Handler(\"00\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_left_disparity_map(left_image, right_image, matcher = 'bm', rgb = True, verbose = True):\n",
    "    '''\n",
    "    Takes a left and right pair of images to computes the disparity map for the left\n",
    "    image. \n",
    "    \n",
    "    Arguments:\n",
    "    img_left -- image from left camera\n",
    "    img_right -- image from right camera\n",
    "    \n",
    "    Optional Arguments:\n",
    "    matcher -- (str) can be 'bm' for StereoBM or 'sgbm' for StereoSGBM matching.\n",
    "    rgb -- (bool) set to True if passing RGB images as input. \n",
    "    verbose -- (bool) set to True to report matching type and time to compute\n",
    "    \n",
    "    Returns:\n",
    "    disp_left -- disparity map for the left camera image\n",
    "    \n",
    "    '''\n",
    "    sad_window = 6\n",
    "    num_disparities = sad_window * 16\n",
    "    block_size = 11\n",
    "    matcher_name = matcher\n",
    "\n",
    "    if matcher_name == 'bm':\n",
    "        matcher = cv2.StereoBM_create(numDisparities=num_disparities,\n",
    "                                        blockSize=block_size)\n",
    "                                        \n",
    "    elif matcher_name == 'sgbm':\n",
    "        matcher = cv2.StereoSGBM_create(numDisparities=num_disparities,\n",
    "                                        blockSize=block_size,\n",
    "                                        P1 = 8*3*sad_window**2,\n",
    "                                        P2 = 32*3*sad_window**2,\n",
    "                                        mode = cv2.STEREO_SGBM_MODE_SGBM_3WAY)\n",
    "    if rgb:\n",
    "        left_image = cv2.cvtColor(left_image, cv2.COLOR_BGR2GRAY)\n",
    "        right_image = cv2.cvtColor(right_image, cv2.COLOR_BGR2GRAY)\n",
    "    \n",
    "    start = datetime.datetime.now()\n",
    "    disp_left = matcher.compute(left_image, right_image).astype(np.float32)/16\n",
    "    end = datetime.datetime.now()\n",
    "\n",
    "    if verbose:\n",
    "        print(f'Time to compute disparity map using Stereo{matcher_name.upper()}:', end-start)\n",
    "    \n",
    "    return disp_left\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time to compute disparity map using StereoBM: 0:00:00.009997\n"
     ]
    }
   ],
   "source": [
    "first_left = cv2.imread(file_path_color_left + left_images[167])\n",
    "first_right = cv2.imread(file_path_color_right + right_images[167])\n",
    "\n",
    "disp = compute_left_disparity_map(left_image=first_left, right_image=first_right,\n",
    "                                    matcher = 'bm', verbose = True)\n",
    "\n",
    "plt.figure(figsize=(15,15))\n",
    "plt.imshow(disp);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "def decompose_projection_matrix(p):\n",
    "    '''\n",
    "    Shortcut to use cv2.decomposeProjectionMatrix(), which only returns k, r, t, and divides\n",
    "    t by the scale, then returns it as a vector with shape (3,) (non-homogeneous)\n",
    "    \n",
    "    Arguments:\n",
    "    p -- projection matrix to be decomposed\n",
    "    \n",
    "    Returns:\n",
    "    k, r, t -- intrinsic matrix, rotation matrix, and 3D translation vector\n",
    "    \n",
    "    '''\n",
    "    k, r, t, _, _, _, _ = cv2.decomposeProjectionMatrix(p)\n",
    "    t = (t / t[3])[:3]\n",
    "    \n",
    "    return k, r, t\n",
    "def calc_depth_map(disp_left, k_left, t_left, t_right, rectified=True):\n",
    "    '''\n",
    "    Assuming we don't have access to the depth map...\n",
    "    \n",
    "    Calculate depth map using a disparity map, intrinsic camera matrix, and translation vectors\n",
    "    from camera extrinsic matrices (to calculate baseline). Note that default behavior is for\n",
    "    rectified projection matrix for right camera. If using a regular projection matrix, pass\n",
    "    rectified=False to avoid issues.\n",
    "\n",
    "    \n",
    "    \n",
    "    Arguments:\n",
    "    disp_left -- disparity map of left camera\n",
    "    k_left -- intrinsic matrix for left camera\n",
    "    t_left -- translation vector for left camera\n",
    "    t_right -- translation vector for right camera\n",
    "    \n",
    "    Optional Arguments:\n",
    "    rectified -- (bool) set to False if t_right is not from rectified projection matrix\n",
    "    \n",
    "    Returns:\n",
    "    depth_map -- calculated depth map for left camera\n",
    "    \n",
    "    '''\n",
    "    # Get focal length of x axis for left camera\n",
    "    f = k_left[0][0]\n",
    "    \n",
    "    # Calculate baseline of stereo pair\n",
    "    if rectified:\n",
    "        b = t_right[0] - t_left[0] \n",
    "    else:\n",
    "        b = t_left[0] - t_right[0]\n",
    "        \n",
    "    # Avoid instability and division by zero\n",
    "    disp_left[disp_left == 0.0] = 0.1\n",
    "    disp_left[disp_left == -1.0] = 0.1\n",
    "    \n",
    "    # Make empty depth map then fill with depth\n",
    "    depth_map = np.ones(disp_left.shape)\n",
    "    depth_map = f * b / disp_left\n",
    "    \n",
    "    return depth_map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[3.1132812, 3.1132812, 3.1152344, ..., 3.9003906, 3.9023438,\n",
       "        3.9042969],\n",
       "       [3.1171875, 3.1191406, 3.1191406, ..., 3.9082031, 3.9101562,\n",
       "        3.9121094],\n",
       "       [3.1230469, 3.1230469, 3.125    , ..., 3.9160156, 3.9179688,\n",
       "        3.9199219],\n",
       "       ...,\n",
       "       [2.0664062, 2.0664062, 2.0683594, ..., 3.1738281, 3.2265625,\n",
       "        3.2285156],\n",
       "       [2.0625   , 2.0644531, 2.0644531, ..., 3.1679688, 3.21875  ,\n",
       "        3.2207031],\n",
       "       [2.0585938, 2.0605469, 2.0605469, ..., 3.2070312, 3.2089844,\n",
       "        3.2128906]], dtype=float32)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "handler.first_depth_left"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pointcloud2image(pointcloud, imheight, imwidth, Tr, P0):\n",
    "    '''\n",
    "    ...\n",
    "    Takes a pointcloud of shape Nx4 and projects it onto an image plane, first transforming\n",
    "    the X, Y, Z coordinates of points to the camera frame with tranformation matrix Tr, then\n",
    "    projecting them using camera projection matrix P0.\n",
    "    \n",
    "    Arguments:\n",
    "    pointcloud -- array of shape Nx4 containing (X, Y, Z, reflectivity)\n",
    "    imheight -- height (in pixels) of image plane\n",
    "    imwidth -- width (in pixels) of image plane\n",
    "    Tr -- 3x4 transformation matrix between lidar (X, Y, Z, 1) homogeneous and camera (X, Y, Z)\n",
    "    P0 -- projection matrix of camera (should have identity transformation if Tr used)\n",
    "    \n",
    "    Returns:\n",
    "    render -- a (imheight x imwidth) array containing depth (Z) information from lidar scan\n",
    "    \n",
    "    '''\n",
    "\n",
    "    pointcloud = pointcloud[pointcloud[:, 0] > 0]\n",
    "    \n",
    "\n",
    "    pointcloud = np.hstack([pointcloud[:, :3], np.ones(pointcloud.shape[0]).reshape((-1,1))])\n",
    "    \n",
    "    # Transform pointcloud into camera coordinate frame\n",
    "    cam_xyz = Tr.dot(pointcloud.T)\n",
    "    \n",
    "    # Ignore any points behind the camera (probably redundant but just in case)\n",
    "    cam_xyz = cam_xyz[:, cam_xyz[2] > 0]\n",
    "    \n",
    "    # Extract the Z row which is the depth from camera\n",
    "    depth = cam_xyz[2].copy()\n",
    "    \n",
    "    # Project coordinates in camera frame to flat plane at Z=1 by dividing by Z\n",
    "    cam_xyz /= cam_xyz[2]\n",
    "    \n",
    "    # Add row of ones to make our 3D coordinates on plane homogeneous for dotting with P0\n",
    "    cam_xyz = np.vstack([cam_xyz, np.ones(cam_xyz.shape[1])])\n",
    "    \n",
    "    # Get pixel coordinates of X, Y, Z points in camera coordinate frame\n",
    "    projection = P0.dot(cam_xyz)\n",
    "    #projection = (projection / projection[2])\n",
    "    \n",
    "    # Turn pixels into integers for indexing\n",
    "    pixel_coordinates = np.round(projection.T, 0)[:, :2].astype('int')\n",
    "    #pixel_coordinates = np.array(pixel_coordinates)\n",
    "    \n",
    "    # Limit pixel coordinates considered to those that fit on the image plane\n",
    "    indices = np.where((pixel_coordinates[:, 0] < imwidth)\n",
    "                       & (pixel_coordinates[:, 0] >= 0)\n",
    "                       & (pixel_coordinates[:, 1] < imheight)\n",
    "                       & (pixel_coordinates[:, 1] >= 0)\n",
    "                      )\n",
    "    pixel_coordinates = pixel_coordinates[indices]\n",
    "    depth = depth[indices]\n",
    "    \n",
    "    # Establish empty render image, then fill with the depths of each point\n",
    "    render = np.zeros((imheight, imwidth))\n",
    "    for j, (u, v) in enumerate(pixel_coordinates):\n",
    "        if u >= imwidth or u < 0:\n",
    "            continue\n",
    "        if v >= imheight or v < 0:\n",
    "            continue\n",
    "        render[v, u] = depth[j]\n",
    "    # Fill zero values with large distance so they will be ignored. (Using same max value)\n",
    "\n",
    "\n",
    "    \n",
    "    return render"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "def stereo_2_depth(img_left, img_right, P0, P1, matcher='bm', rgb=True, verbose=False, \n",
    "                   rectified=True):\n",
    "    '''\n",
    "    Takes stereo pair of images and returns a depth map for the left camera. If your projection\n",
    "    matrices are not rectified, set rectified=False.\n",
    "    \n",
    "    Arguments:\n",
    "    img_left -- image of left camera\n",
    "    img_right -- image of right camera\n",
    "    P0 -- Projection matrix for the left camera\n",
    "    P1 -- Projection matrix for the right camera\n",
    "    \n",
    "    Optional Arguments:\n",
    "    matcher -- (str) can be 'bm' for StereoBM or 'sgbm' for StereoSGBM\n",
    "    rgb -- (bool) set to True if images passed are RGB. Default is False\n",
    "    verbose -- (bool) set to True to report computation time and method\n",
    "    rectified -- (bool) set to False if P1 not rectified to P0. Default is True\n",
    "    \n",
    "    Returns:\n",
    "    depth -- depth map for left camera\n",
    "    \n",
    "    '''\n",
    "    # Compute disparity map\n",
    "    disp = compute_left_disparity_map(img_left, \n",
    "                                      img_right, \n",
    "                                      matcher=matcher, \n",
    "                                      rgb=rgb, \n",
    "                                      verbose=verbose)\n",
    "    # Decompose projection matrices\n",
    "    k_left, r_left, t_left = decompose_projection_matrix(P0)\n",
    "    k_right, r_right, t_right = decompose_projection_matrix(P1)\n",
    "    # Calculate depth map for left camera\n",
    "    depth = calc_depth_map(disp, k_left, t_left, t_right)\n",
    "    \n",
    "    return depth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_features(image, detector='sift', mask=None):\n",
    "    \"\"\"\n",
    "    Find keypoints and descriptors for the image\n",
    "\n",
    "    Arguments:\n",
    "    image -- a grayscale image\n",
    "\n",
    "    Returns:\n",
    "    kp -- list of the extracted keypoints (features) in an image\n",
    "    des -- list of the keypoint descriptors in an image\n",
    "    \"\"\"\n",
    "    if detector == 'sift':\n",
    "        det = cv2.SIFT_create()\n",
    "    elif detector == 'orb':\n",
    "        det = cv2.ORB_create()\n",
    "    elif detector == 'surf':\n",
    "        det = cv2.xfeatures2d.SURF_create()\n",
    "        \n",
    "    kp, des = det.detectAndCompute(image, mask)\n",
    "    \n",
    "    return kp, des"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "def match_features(des1, des2, matching='BF', detector='sift', sort=True, k=2):\n",
    "    \"\"\"\n",
    "    Match features from two images\n",
    "\n",
    "    Arguments:\n",
    "    des1 -- list of the keypoint descriptors in the first image\n",
    "    des2 -- list of the keypoint descriptors in the second image\n",
    "    matching -- (str) can be 'BF' for Brute Force or 'FLANN'\n",
    "    detector -- (str) can be 'sift or 'orb'. Default is 'sift'\n",
    "    sort -- (bool) whether to sort matches by distance. Default is True\n",
    "    k -- (int) number of neighbors to match to each feature.\n",
    "\n",
    "    Returns:\n",
    "    matches -- list of matched features from two images. Each match[i] is k or less matches for \n",
    "               the same query descriptor\n",
    "    \"\"\"\n",
    "    if matching == 'BF':\n",
    "        if detector == 'sift':\n",
    "            matcher = cv2.BFMatcher_create(cv2.NORM_L2, crossCheck=False)\n",
    "        elif detector == 'orb':\n",
    "            matcher = cv2.BFMatcher_create(cv2.NORM_HAMMING2, crossCheck=False)\n",
    "        matches = matcher.knnMatch(des1, des2, k=k)\n",
    "    elif matching == 'FLANN':\n",
    "        FLANN_INDEX_KDTREE = 1\n",
    "        index_params = dict(algorithm = FLANN_INDEX_KDTREE, trees=5)\n",
    "        search_params = dict(checks=50)\n",
    "        matcher = cv2.FlannBasedMatcher(index_params, search_params)\n",
    "        matches = matcher.knnMatch(des1, des2, k=k)\n",
    "    \n",
    "    if sort:\n",
    "        matches = sorted(matches, key = lambda x:x[0].distance)\n",
    "\n",
    "    return matches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_matches(image1, kp1, image2, kp2, match):\n",
    "    \"\"\"\n",
    "    Visualize corresponding matches in two images\n",
    "\n",
    "    Arguments:\n",
    "    image1 -- the first image in a matched image pair\n",
    "    kp1 -- list of the keypoints in the first image\n",
    "    image2 -- the second image in a matched image pair\n",
    "    kp2 -- list of the keypoints in the second image\n",
    "    match -- list of matched features from the pair of images\n",
    "\n",
    "    Returns:\n",
    "    image_matches -- an image showing the corresponding matches on both image1 and image2 or None if you don't use this function\n",
    "    \"\"\"\n",
    "    image_matches = cv2.drawMatches(image1, kp1, image2, kp2, match, None, flags=2)\n",
    "    plt.figure(figsize=(16, 6), dpi=100)\n",
    "    plt.imshow(image_matches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_matches_distance(matches, dist_threshold):\n",
    "    \"\"\"\n",
    "    Filter matched features from two images by distance between the best matches\n",
    "\n",
    "    Arguments:\n",
    "    match -- list of matched features from two images\n",
    "    dist_threshold -- maximum allowed relative distance between the best matches, (0.0, 1.0) \n",
    "\n",
    "    Returns:\n",
    "    filtered_match -- list of good matches, satisfying the distance threshold\n",
    "    \"\"\"\n",
    "    filtered_match = []\n",
    "    for m, n in matches:\n",
    "        if m.distance <= dist_threshold*n.distance:\n",
    "            filtered_match.append(m)\n",
    "\n",
    "    return filtered_match"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "def estimate_motion(match, kp1, kp2, k, depth1=None, max_depth=3000):\n",
    "    \"\"\"\n",
    "    Estimate camera motion from a pair of subsequent image frames\n",
    "\n",
    "    Arguments:\n",
    "    match -- list of matched features from the pair of images\n",
    "    kp1 -- list of the keypoints in the first image\n",
    "    kp2 -- list of the keypoints in the second image\n",
    "    k -- camera intrinsic calibration matrix \n",
    "    \n",
    "    Optional arguments:\n",
    "    depth1 -- Depth map of the first frame. Set to None to use Essential Matrix decomposition\n",
    "    max_depth -- Threshold of depth to ignore matched features. 3000 is default\n",
    "\n",
    "    Returns:\n",
    "    rmat -- estimated 3x3 rotation matrix\n",
    "    tvec -- estimated 3x1 translation vector\n",
    "    image1_points -- matched feature pixel coordinates in the first image. \n",
    "                     image1_points[i] = [u, v] -> pixel coordinates of i-th match\n",
    "    image2_points -- matched feature pixel coordinates in the second image. \n",
    "                     image2_points[i] = [u, v] -> pixel coordinates of i-th match\n",
    "               \n",
    "    \"\"\"\n",
    "    rmat = np.eye(3)\n",
    "    tvec = np.zeros((3, 1))\n",
    "    \n",
    "    image1_points = np.float32([kp1[m.queryIdx].pt for m in match])\n",
    "    image2_points = np.float32([kp2[m.trainIdx].pt for m in match])\n",
    "\n",
    "   \n",
    "    # if depth1 is not None:\n",
    "    cx = k[0, 2]\n",
    "    cy = k[1, 2]\n",
    "    fx = k[0, 0]\n",
    "    fy = k[1, 1]\n",
    "    object_points = np.zeros((0, 3))\n",
    "    delete = []\n",
    "    \n",
    "    # Extract depth information of query image at match points and build 3D positions\n",
    "    for i, (u, v) in enumerate(image1_points):\n",
    "        z = depth1[int(v), int(u)]\n",
    "        \n",
    "\n",
    "        if z > max_depth:\n",
    "            delete.append(i)\n",
    "            continue\n",
    "            \n",
    "        # Use arithmetic to extract x and y (faster than using inverse of k)\n",
    "        x = z*(u-cx)/fx\n",
    "        y = z*(v-cy)/fy\n",
    "\n",
    "        object_points = np.vstack([object_points, np.array([x, y, z])])\n",
    "\n",
    "        # Equivalent math with dot product w/ inverse of k matrix, but SLOWER (see Appendix A)\n",
    "        #object_points = np.vstack([object_points, np.linalg.inv(k).dot(z*np.array([u, v, 1]))])\n",
    "\n",
    "    image1_points = np.delete(image1_points, delete, 0)\n",
    "\n",
    "    image2_points = np.delete(image2_points, delete, 0)\n",
    "\n",
    "    \n",
    "    # Use PnP algorithm with RANSAC for robustness to outliers\n",
    "    \n",
    "    _,rvec, tvec, inliers = cv2.solvePnPRansac(object_points, image2_points, cameraMatrix=k, distCoeffs=None)\n",
    "\n",
    "    rmat = cv2.Rodrigues(rvec)[0]\n",
    "\n",
    "    \n",
    "    return rmat, tvec, image1_points, image2_points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of matches before filtering: 47\n",
      "Number of matches after filtering: 5\n"
     ]
    }
   ],
   "source": [
    "image_left11 = handler.first_image_left\n",
    "image_plus11 = handler.second_image_left\n",
    "depth_left11 = np.load(handler.depth_dir + 'depth_left/'\n",
    "                                                + handler.left_depth_files[133])\n",
    "image_right11 = handler.first_image_right\n",
    "depth_plus11 = np.load(handler.depth_dir + 'depth_left/'\n",
    "                                                + handler.left_depth_files[1])\n",
    "\n",
    "kp0, des0 = extract_features(image_left11, 'sift')\n",
    "kp1, des1 = extract_features(image_plus11, 'sift')\n",
    "matches = match_features(des0, des1, matching='BF', detector='sift', sort=True)\n",
    "print('Number of matches before filtering:', len(matches))\n",
    "matches = filter_matches_distance(matches, 0.45)\n",
    "print('Number of matches after filtering:', len(matches))\n",
    "visualize_matches(image_left11, kp0, image_plus11, kp1, matches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time to compute disparity map using StereoBM: 0:00:00.006027\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005975\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007059\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005005\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006033\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006032\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007002\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008971\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005008\n",
      "Time to compute disparity map using StereoBM: 0:00:00.010032\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008002\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006968\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006027\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005998\n",
      "Time to compute disparity map using StereoBM: 0:00:00.004999\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006000\n",
      "Time to compute disparity map using StereoBM: 0:00:00.009004\n",
      "Time to compute disparity map using StereoBM: 0:00:00.004030\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005006\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007015\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006002\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007987\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007984\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006001\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005040\n",
      "Time to compute disparity map using StereoBM: 0:00:00.009011\n",
      "Time to compute disparity map using StereoBM: 0:00:00.004999\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006999\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006990\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005026\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006000\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005007\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006002\n",
      "Time to compute disparity map using StereoBM: 0:00:00.010000\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008001\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007071\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006003\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007002\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006000\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005973\n",
      "Time to compute disparity map using StereoBM: 0:00:00.009014\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007000\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005959\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005027\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008971\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005963\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006999\n",
      "Time to compute disparity map using StereoBM: 0:00:00.004992\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008035\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005034\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008070\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005965\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007045\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005029\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005024\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008015\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007006\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008035\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005035\n",
      "Time to compute disparity map using StereoBM: 0:00:00.004987\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007000\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006003\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006999\n",
      "Time to compute disparity map using StereoBM: 0:00:00.004999\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006024\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005034\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007015\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005999\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006997\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005005\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007992\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006002\n",
      "Time to compute disparity map using StereoBM: 0:00:00.004994\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006998\n",
      "Time to compute disparity map using StereoBM: 0:00:00.009008\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005032\n",
      "Time to compute disparity map using StereoBM: 0:00:00.004998\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008012\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008000\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008001\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005029\n",
      "Time to compute disparity map using StereoBM: 0:00:00.009059\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005999\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005999\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005002\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005032\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007999\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008033\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007001\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005010\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005015\n",
      "Time to compute disparity map using StereoBM: 0:00:00.004998\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006998\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008027\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005032\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008001\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008007\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005975\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006999\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006000\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007000\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007033\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006002\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005990\n",
      "Time to compute disparity map using StereoBM: 0:00:00.009987\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007999\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005026\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006980\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008000\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006998\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005025\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005005\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008969\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005910\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006022\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007034\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007029\n",
      "Time to compute disparity map using StereoBM: 0:00:00.004999\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006002\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005030\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007997\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007996\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005998\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005034\n",
      "Time to compute disparity map using StereoBM: 0:00:00.011007\n",
      "Time to compute disparity map using StereoBM: 0:00:00.010080\n",
      "Time to compute disparity map using StereoBM: 0:00:00.010015\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006014\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008999\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008984\n",
      "Time to compute disparity map using StereoBM: 0:00:00.011999\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006001\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008001\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006000\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006001\n",
      "Time to compute disparity map using StereoBM: 0:00:00.009995\n",
      "Time to compute disparity map using StereoBM: 0:00:00.004999\n",
      "Time to compute disparity map using StereoBM: 0:00:00.010038\n",
      "Time to compute disparity map using StereoBM: 0:00:00.009000\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008999\n",
      "Time to compute disparity map using StereoBM: 0:00:00.010001\n",
      "Time to compute disparity map using StereoBM: 0:00:00.009986\n",
      "Time to compute disparity map using StereoBM: 0:00:00.009037\n",
      "Time to compute disparity map using StereoBM: 0:00:00.004969\n",
      "Time to compute disparity map using StereoBM: 0:00:00.010025\n",
      "Time to compute disparity map using StereoBM: 0:00:00.010075\n",
      "Time to compute disparity map using StereoBM: 0:00:00.009000\n",
      "Time to compute disparity map using StereoBM: 0:00:00.009035\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007008\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008998\n",
      "Time to compute disparity map using StereoBM: 0:00:00.011004\n",
      "Time to compute disparity map using StereoBM: 0:00:00.009992\n",
      "Time to compute disparity map using StereoBM: 0:00:00.011044\n",
      "Time to compute disparity map using StereoBM: 0:00:00.010033\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005036\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006029\n",
      "Time to compute disparity map using StereoBM: 0:00:00.009970\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005899\n",
      "Time to compute disparity map using StereoBM: 0:00:00.010003\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005994\n",
      "Time to compute disparity map using StereoBM: 0:00:00.010007\n",
      "Time to compute disparity map using StereoBM: 0:00:00.004999\n",
      "Time to compute disparity map using StereoBM: 0:00:00.006975\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005996\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008987\n",
      "Time to compute disparity map using StereoBM: 0:00:00.009998\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005997\n",
      "Time to compute disparity map using StereoBM: 0:00:00.011049\n",
      "Time to compute disparity map using StereoBM: 0:00:00.009039\n",
      "Time to compute disparity map using StereoBM: 0:00:00.010054\n",
      "Time to compute disparity map using StereoBM: 0:00:00.011027\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007001\n",
      "Time to compute disparity map using StereoBM: 0:00:00.009010\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005970\n",
      "Time to compute disparity map using StereoBM: 0:00:00.009969\n",
      "Time to compute disparity map using StereoBM: 0:00:00.011046\n",
      "Time to compute disparity map using StereoBM: 0:00:00.010001\n",
      "Time to compute disparity map using StereoBM: 0:00:00.010000\n",
      "Time to compute disparity map using StereoBM: 0:00:00.009003\n",
      "Time to compute disparity map using StereoBM: 0:00:00.011048\n",
      "Time to compute disparity map using StereoBM: 0:00:00.010034\n",
      "Time to compute disparity map using StereoBM: 0:00:00.010001\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008964\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007957\n",
      "Time to compute disparity map using StereoBM: 0:00:00.009027\n",
      "Time to compute disparity map using StereoBM: 0:00:00.011011\n",
      "Time to compute disparity map using StereoBM: 0:00:00.007996\n",
      "Time to compute disparity map using StereoBM: 0:00:00.005027\n",
      "Time to compute disparity map using StereoBM: 0:00:00.010033\n",
      "Time to compute disparity map using StereoBM: 0:00:00.010005\n",
      "Time to compute disparity map using StereoBM: 0:00:00.008010\n",
      "Time to compute disparity map using StereoBM: 0:00:00.009967\n",
      "Time to compute disparity map using StereoBM: 0:00:00.009992\n",