write_depthmap_to_kitti360.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from data_loader.calib.intrinsic_extrinsic_loader import IntrinsicExtrinsicLoader\n",
    "from tools.utils import *\n",
    "\n",
    "def generate_depthmap_to_kitti360(dataset_path, calib_path, platform, sequence_name, algorithm):\n",
    "  ##### Set up the output data path\n",
    "  kitti360_path = os.path.join(dataset_path, sequence_name, 'kitti360')\n",
    "\n",
    "  ##### Set up the message topic list for different platforms\n",
    "  # Platform\n",
    "  if platform == 'handheld':\n",
    "    from cfg.dataset.cfg_handheld import dataset_sensor_frameid_dict\n",
    "    from cfg.dataset.cfg_handheld import dataset_rostopic_msg_frameid_dict\n",
    "  elif platform == 'ugv':\n",
    "    from cfg.dataset.cfg_ugv import dataset_sensor_frameid_dict\n",
    "    from cfg.dataset.cfg_ugv import dataset_rostopic_msg_frameid_dict\n",
    "  elif platform == 'legged':\n",
    "    from cfg.dataset.cfg_legged import dataset_sensor_frameid_dict\n",
    "    from cfg.dataset.cfg_legged import dataset_rostopic_msg_frameid_dict\n",
    "    if sequence_name == 'legged_grass00':\n",
    "      dataset_sensor_frameid_dict, dataset_rostopic_msg_frameid_dict = \\\n",
    "        filter_sensor('event', dataset_sensor_frameid_dict, dataset_rostopic_msg_frameid_dict)\n",
    "  elif platform =='vehicle':\n",
    "    from cfg.dataset.cfg_vehicle import dataset_sensor_frameid_dict\n",
    "    from cfg.dataset.cfg_vehicle import dataset_rostopic_msg_frameid_dict\n",
    "\n",
    "  ##### Set up the sensor configuration\n",
    "  int_ext_loader = IntrinsicExtrinsicLoader(is_print=False)\n",
    "  int_ext_loader.load_calibration(calib_path=calib_path, sensor_frameid_dict=dataset_sensor_frameid_dict)\n",
    "  print('Finish loading parameters')\n",
    "\n",
    "  ##### Create the output data path\n",
    "  if platform == 'vehicle':\n",
    "    flag = os.system('mkdir -p '+ os.path.join(kitti360_path, 'vehicle_frame_cam00/color_proj_image/data'))\n",
    "    flag = os.system('mkdir -p '+ os.path.join(kitti360_path, 'vehicle_frame_cam01/color_proj_image/data'))\n",
    "    flag = os.system('mkdir -p '+ os.path.join(kitti360_path, 'vehicle_frame_cam00/depth_image/data'))\n",
    "    flag = os.system('mkdir -p '+ os.path.join(kitti360_path, 'vehicle_frame_cam01/depth_image/data'))\n",
    "  else:\n",
    "    flag = os.system('mkdir -p '+ os.path.join(kitti360_path, 'frame_cam00/color_proj_image/data'))\n",
    "    flag = os.system('mkdir -p '+ os.path.join(kitti360_path, 'frame_cam01/color_proj_image/data'))\n",
    "    flag = os.system('mkdir -p '+ os.path.join(kitti360_path, 'frame_cam00/depth_image/data'))\n",
    "    flag = os.system('mkdir -p '+ os.path.join(kitti360_path, 'frame_cam01/depth_image/data'))\n",
    "  print('Finish creating Folder')\n",
    "\n",
    "  ##### Generate depth map\n",
    "  # Description: Generate depth map\n",
    "  import cv2\n",
    "  import open3d\n",
    "  from PIL import Image\n",
    "  from IPython.display import display\n",
    "  import matplotlib.pyplot as plt\n",
    "  cmap = plt.get_cmap('jet')\n",
    "  cmap_colors = cmap(np.linspace(0, 1, 31)) * 255\n",
    "\n",
    "  for frame_id in range(100000):\n",
    "    ##### Load undistroted point cloud and undistroted image\n",
    "    pcd_path = os.path.join(kitti360_path, 'ouster00_undistort/points/data', '{:06d}.pcd'.format(frame_id))\n",
    "    if not os.path.exists(pcd_path):\n",
    "      break\n",
    "\n",
    "    xyz_points = np.asarray(open3d.io.read_point_cloud(pcd_path).points)\n",
    "    if platform == 'vehicle':\n",
    "      img_path = os.path.join(kitti360_path, 'vehicle_frame_cam00', 'image/data', '{:06d}.png'.format(frame_id)) # images are already undistorted, but not rectified\n",
    "      img = cv2.imread(img_path)\n",
    "      camera = int_ext_loader.sensor_collection['vehicle_frame_left_camera']\n",
    "    else:\n",
    "      img_path = os.path.join(kitti360_path, 'frame_cam00', 'image/data', '{:06d}.png'.format(frame_id)) # images are already undistorted, but not rectified\n",
    "      img = cv2.imread(img_path)\n",
    "      camera = int_ext_loader.sensor_collection['frame_left_camera']\n",
    "\n",
    "    ##### Load extrinsics from the tf_graph\n",
    "    T_cam_lidar = int_ext_loader.tf_graph.get_relative_transform(camera.frame_id, 'body_imu')\n",
    "    xyz_points_cam = np.matmul(T_cam_lidar[:3, :3], xyz_points.T).T + T_cam_lidar[:3, 3].T\n",
    "\n",
    "    ##### Project point cloud onto the camera frame\n",
    "    for p_C in xyz_points_cam:\n",
    "      flag, u_C = camera.project(p_C)\n",
    "      if flag:\n",
    "        i = int(min(p_C[2], 30.0))\n",
    "        color = (int(cmap_colors[i, 0]), int(cmap_colors[i, 1]), int(cmap_colors[i, 2]))\n",
    "        cv2.circle(img, (round(u_C[0]), round(u_C[1])), radius=1, color=color, thickness=-1)\n",
    "    img_pillow = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))\n",
    "    new_img_path = img_path.replace('image', 'color_proj_image')\n",
    "    img_pillow.save(new_img_path)\n",
    "\n",
    "    ##### Project point cloud onto the depth camera frame\n",
    "    depth_img = np.zeros((camera.height, camera.width), dtype=np.uint32)\n",
    "    for p_C in xyz_points_cam:\n",
    "      flag, u_C = camera.project(p_C)\n",
    "      if flag:\n",
    "        depth_img[round(u_C[1])][round(u_C[0])] = np.uint32(p_C[2] * 1000)\n",
    "    img_pillow = Image.fromarray(depth_img)\n",
    "    new_img_path = img_path.replace('image', 'depth_image')\n",
    "    img_pillow.save(new_img_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cfg.dataset.cfg_sequence import dataset_sequence_calib_used_dict\n",
    "\n",
    "algorithms = ['fastlio2']\n",
    "for sequence_name in dataset_sequence_calib_used_dict.keys():\n",
    "  for algorithm in algorithms:\n",
    "    platform = dataset_sequence_calib_used_dict[sequence_name][0]\n",
    "    calib_folder = dataset_sequence_calib_used_dict[sequence_name][1]\n",
    "    used = dataset_sequence_calib_used_dict[sequence_name][2]\n",
    "    if used:\n",
    "      print('Start processing platform: {} sequence: {} with algorithm_result: {}'.format(platform, sequence_name, algorithm))\n",
    "      dataset_path = '/Rocket_ssd/dataset/FusionPortable_dataset_develop/sensor_data/data_refined'\n",
    "      calib_path = os.path.join(dataset_path, '../../calibration_files', calib_folder, 'calib')\n",
    "      generate_depthmap_to_kitti360(dataset_path, calib_path, platform, sequence_name, algorithm)\n",
    "      print('Finish generating depth images')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "fp_dataset",
   "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}