write_alg_bag_to_data.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import rosbag\n",
    "from cfg.dataset.cfg_sequence import dataset_sequence_calib_used_dict\n",
    "from data_loader.ros_msg.pointcloud import PointCloud\n",
    "from data_loader.ros_msg.image import Image\n",
    "from data_loader.ros_msg.odometry import Odometry\n",
    "from data_loader.calib.intrinsic_extrinsic_loader import IntrinsicExtrinsicLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_alg_bag_to_data(sequence_name, algorithm, rosbag_path, dataset_path, rostopic_msg_frameid_dict):\n",
    "  ##### Open the rosbag\n",
    "  input_bag = rosbag.Bag(rosbag_path)\n",
    "  print('Finish reading bag, start loading messages, and writing message to data folder')\n",
    "  print('Will save messages to algorithm_result'.format(os.path.join(dataset_path, sequence_name, algorithm)))\n",
    "\n",
    "  ##### Initialize the lidar object of Sensor class\n",
    "  # Ouster_undistorted\n",
    "  print('Loading Ouster_undistorted messages...')\n",
    "  pointcloud = PointCloud(sensor_type='vlp')\n",
    "  output_data_path = os.path.join(dataset_path, sequence_name, 'algorithm_result'.format(algorithm), 'ouster00_undistort', 'points')\n",
    "  num_msg = pointcloud.load_messages_write_to_file(bag=input_bag, output_path=output_data_path, topic=rostopic_msg_frameid_dict['ouster_points_undistorted'][0])\n",
    "  print('     Saving {} Ouster_undistorted points messages to {}!'.format(num_msg, output_data_path))\n",
    "\n",
    "  # Odometry\n",
    "  print('Loading Odometry messages...')\n",
    "  odometry = Odometry(traj_type='TUM', msg_type='nav_msgs/Odometry')\n",
    "  output_data_path = os.path.join(dataset_path, sequence_name, 'algorithm_result'.format(algorithm), 'odometry')\n",
    "  num_msg = odometry.load_messages_write_to_file(bag=input_bag, output_path=output_data_path, topic=rostopic_msg_frameid_dict['odometry'][0])\n",
    "  print('     Saving {} Odometry messages to {}!'.format(num_msg, output_data_path))\n",
    "  print('Finishing loading Odometry messages!')\n",
    "\n",
    "  # Camera Odometry\n",
    "  if algorithm == 'r3live':\n",
    "    print('Loading Camera_Odometry messages...')\n",
    "    odometry = Odometry(traj_type='TUM', msg_type='nav_msgs/Odometry')\n",
    "    output_data_path = os.path.join(dataset_path, sequence_name, 'algorithm_result'.format(algorithm), 'camera_odometry')\n",
    "    num_msg = odometry.load_messages_write_to_file(bag=input_bag, output_path=output_data_path, topic=rostopic_msg_frameid_dict['camera_odometry'][0])\n",
    "    print('     Saving {} Camera_Odometry messages to {}!'.format(num_msg, output_data_path))\n",
    "    print('Finishing loading Camera_Odometry messages!')\n",
    "\n",
    "  ##### Close the rosbag\n",
    "  input_bag.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Start processing sequence: handheld_room00 with algorithm_result: fastlio2\n",
      "Finish loading parameters\n",
      "Finish reading bag, start loading messages, and writing message to data folder\n",
      "Will save messages to algorithm_result\n",
      "Loading Ouster_undistorted messages...\n",
      "     Saving 1361 Ouster_undistorted points messages to /Rocket_ssd/dataset/FusionPortable_dataset_develop/sensor_data/data_refined/handheld_room00/algorithm_result/ouster00_undistort/points!\n",
      "Loading Odometry messages...\n",
      "     Saving 1361 Odometry messages to /Rocket_ssd/dataset/FusionPortable_dataset_develop/sensor_data/data_refined/handheld_room00/algorithm_result/odometry!\n",
      "Finishing loading Odometry messages!\n",
      "Start processing sequence: legged_room00 with algorithm_result: fastlio2\n",
      "Finish loading parameters\n",
      "Finish reading bag, start loading messages, and writing message to data folder\n",
      "Will save messages to algorithm_result\n",
      "Loading Ouster_undistorted messages...\n",
      "     Saving 1736 Ouster_undistorted points messages to /Rocket_ssd/dataset/FusionPortable_dataset_develop/sensor_data/data_refined/legged_room00/algorithm_result/ouster00_undistort/points!\n",
      "Loading Odometry messages...\n",
      "     Saving 1736 Odometry messages to /Rocket_ssd/dataset/FusionPortable_dataset_develop/sensor_data/data_refined/legged_room00/algorithm_result/odometry!\n",
      "Finishing loading Odometry messages!\n",
      "Start processing sequence: legged_grass00 with algorithm_result: fastlio2\n",
      "Finish loading parameters\n",
      "Finish reading bag, start loading messages, and writing message to data folder\n",
      "Will save messages to algorithm_result\n",
      "Loading Ouster_undistorted messages...\n",
      "     Saving 3016 Ouster_undistorted points messages to /Rocket_ssd/dataset/FusionPortable_dataset_develop/sensor_data/data_refined/legged_grass00/algorithm_result/ouster00_undistort/points!\n",
      "Loading Odometry messages...\n",
      "     Saving 3016 Odometry messages to /Rocket_ssd/dataset/FusionPortable_dataset_develop/sensor_data/data_refined/legged_grass00/algorithm_result/odometry!\n",
      "Finishing loading Odometry messages!\n",
      "Start processing sequence: ugv_campus00 with algorithm_result: fastlio2\n",
      "Finish loading parameters\n",
      "Finish reading bag, start loading messages, and writing message to data folder\n",
      "Will save messages to algorithm_result\n",
      "Loading Ouster_undistorted messages...\n",
      "     Saving 3336 Ouster_undistorted points messages to /Rocket_ssd/dataset/FusionPortable_dataset_develop/sensor_data/data_refined/ugv_campus00/algorithm_result/ouster00_undistort/points!\n",
      "Loading Odometry messages...\n",
      "     Saving 3336 Odometry messages to /Rocket_ssd/dataset/FusionPortable_dataset_develop/sensor_data/data_refined/ugv_campus00/algorithm_result/odometry!\n",
      "Finishing loading Odometry messages!\n",
      "Start processing sequence: ugv_parking00 with algorithm_result: fastlio2\n",
      "Finish loading parameters\n",
      "Finish reading bag, start loading messages, and writing message to data folder\n",
      "Will save messages to algorithm_result\n",
      "Loading Ouster_undistorted messages...\n",
      "     Saving 1786 Ouster_undistorted points messages to /Rocket_ssd/dataset/FusionPortable_dataset_develop/sensor_data/data_refined/ugv_parking00/algorithm_result/ouster00_undistort/points!\n",
      "Loading Odometry messages...\n",
      "     Saving 1786 Odometry messages to /Rocket_ssd/dataset/FusionPortable_dataset_develop/sensor_data/data_refined/ugv_parking00/algorithm_result/odometry!\n",
      "Finishing loading Odometry messages!\n",
      "Start processing sequence: vehicle_campus00 with algorithm_result: fastlio2\n",
      "Finish loading parameters\n",
      "Finish reading bag, start loading messages, and writing message to data folder\n",
      "Will save messages to algorithm_result\n",
      "Loading Ouster_undistorted messages...\n",
      "     Saving 6106 Ouster_undistorted points messages to /Rocket_ssd/dataset/FusionPortable_dataset_develop/sensor_data/data_refined/vehicle_campus00/algorithm_result/ouster00_undistort/points!\n",
      "Loading Odometry messages...\n",
      "     Saving 6106 Odometry messages to /Rocket_ssd/dataset/FusionPortable_dataset_develop/sensor_data/data_refined/vehicle_campus00/algorithm_result/odometry!\n",
      "Finishing loading Odometry messages!\n",
      "Start processing sequence: vehicle_highway00 with algorithm_result: fastlio2\n",
      "Finish loading parameters\n",
      "Finish reading bag, start loading messages, and writing message to data folder\n",
      "Will save messages to algorithm_result\n",
      "Loading Ouster_undistorted messages...\n",
      "     Saving 6946 Ouster_undistorted points messages to /Rocket_ssd/dataset/FusionPortable_dataset_develop/sensor_data/data_refined/vehicle_highway00/algorithm_result/ouster00_undistort/points!\n",
      "Loading Odometry messages...\n",
      "     Saving 6946 Odometry messages to /Rocket_ssd/dataset/FusionPortable_dataset_develop/sensor_data/data_refined/vehicle_highway00/algorithm_result/odometry!\n",
      "Finishing loading Odometry messages!\n"
     ]
    }
   ],
   "source": [
    "algorithms = ['fastlio2']\n",
    "for sequence_name in dataset_sequence_calib_used_dict.keys():\n",
    "  for algorithm in algorithms:\n",
    "    print('Start processing sequence: {} with algorithm_result: {}'.format(sequence_name, algorithm))\n",
    "\n",
    "    ##### Set up the rosbag path\n",
    "    input_dataset_path = '/Rocket_ssd/dataset/FusionPortable_dataset_develop/evaluation/algorithm_result'\n",
    "    rosbag_path = os.path.join(input_dataset_path, sequence_name + '_{}.bag'.format(algorithm))\n",
    "\n",
    "    ##### Set up the output data path\n",
    "    dataset_path = '/Rocket_ssd/dataset/FusionPortable_dataset_develop/sensor_data/data_refined'\n",
    "\n",
    "    ##### Set up the message topic list for different algorithms\n",
    "    if algorithm =='r3live':\n",
    "      from cfg.algorithm.cfg_r3live import algorithm_rostopic_msg_frameid_dict\n",
    "    elif algorithm =='fastlio2':\n",
    "      from cfg.algorithm.cfg_fastlio2 import algorithm_rostopic_msg_frameid_dict\n",
    "    print('Finish loading parameters')\n",
    "\n",
    "    ##### Main function\n",
    "    write_alg_bag_to_data(sequence_name, algorithm, rosbag_path, dataset_path, algorithm_rostopic_msg_frameid_dict)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base_py3",
   "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
}