diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..30fcc1f --- /dev/null +++ b/.dockerignore @@ -0,0 +1,109 @@ +# Created by .ignore support plugin (hsz.mobi) +### Python template +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# pyenv +.python-version + +# celery beat schedule file +celerybeat-schedule + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ + +Dockerfile +.travis.yml diff --git a/.gitignore b/.gitignore index f05bb9a..eedf83d 100644 --- a/.gitignore +++ b/.gitignore @@ -95,3 +95,5 @@ ENV/ # Temporary files from text editors *~ + +tests/data diff --git a/.travis.yml b/.travis.yml index e4e2b5e..bb934ad 100644 --- a/.travis.yml +++ b/.travis.yml @@ -1,129 +1,25 @@ -# Generic .travis.yml file for running continuous integration on Travis-CI for -# any ROS package. -# -# Available here: -# - http://felixduvallet.github.io/ros-travis-integration -# - https://github.com/felixduvallet/ros-travis-integration -# -# This installs ROS on a clean Travis-CI virtual machine, creates a ROS -# workspace, resolves all listed dependencies, and sets environment variables -# (setup.bash). Then, it compiles the entire ROS workspace (ensuring there are -# no compilation errors), and runs all the tests. If any of the compilation/test -# phases fail, the build is marked as a failure. -# -# We handle two types of package dependencies specified in the package manifest: -# - system dependencies that can be installed using `rosdep`, including other -# ROS packages and system libraries. These dependencies must be known to -# `rosdistro` and get installed using apt-get. -# - package dependencies that must be checked out from source. These are handled by -# `wstool`, and should be listed in a file named dependencies.rosinstall. -# -# There are two variables you may want to change: -# - ROS_DISTRO (default is indigo). Note that packages must be available for -# ubuntu 14.04 trusty. -# - ROSINSTALL_FILE (default is dependencies.rosinstall inside the repo -# root). This should list all necessary repositories in wstool format (see -# the ros wiki). If the file does not exists then nothing happens. -# -# See the README.md for more information. -# -# Author: Felix Duvallet +language: generic -# NOTE: The build lifecycle on Travis.ci is something like this: -# before_install -# install -# before_script -# script -# after_success or after_failure -# after_script -# OPTIONAL before_deploy -# OPTIONAL deploy -# OPTIONAL after_deploy +services: + - docker -################################################################################ +matrix: + include: + - env: ROS_DISTRO=kinetic + - env: ROS_DISTRO=lunar + - env: ROS_DISTRO=melodic -# Use ubuntu trusty (14.04) with sudo privileges. -dist: - - trusty -sudo: required -language: - - generic cache: - - apt + directories: + # Avoid re-downloading the test data each time + - $TRAVIS_BUILD_DIR/testdata -# Configuration variables. All variables are global now, but this can be used to -# trigger a build matrix for different ROS distributions if desired. -env: - global: - - ROS_DISTRO=indigo - - ROS_CI_DESKTOP="`lsb_release -cs`" # e.g. [precise|trusty|...] - - CI_SOURCE_PATH=$(pwd) - - ROSINSTALL_FILE=$CI_SOURCE_PATH/dependencies.rosinstall - - CATKIN_OPTIONS=$CI_SOURCE_PATH/catkin.options - - ROS_PARALLEL_JOBS='-j8 -l6' - -################################################################################ - -# Install system dependencies, namely a very barebones ROS setup. before_install: - - sudo sh -c "echo \"deb http://packages.ros.org/ros/ubuntu $ROS_CI_DESKTOP main\" > /etc/apt/sources.list.d/ros-latest.list" - - wget http://packages.ros.org/ros.key -O - | sudo apt-key add - - - sudo apt-get update -qq - - sudo apt-get install -y python-catkin-pkg python-rosdep python-wstool ros-$ROS_DISTRO-catkin - - source /opt/ros/$ROS_DISTRO/setup.bash - # Prepare rosdep to install dependencies. - - sudo rosdep init - - rosdep update - -# Create a catkin workspace with the package under integration. -install: - - mkdir -p ~/catkin_ws/src - - cd ~/catkin_ws/src - - catkin_init_workspace - # Create the devel/setup.bash (run catkin_make with an empty workspace) and - # source it to set the path variables. - - cd ~/catkin_ws - - catkin_make - - source devel/setup.bash - # Add the package under integration to the workspace using a symlink. - - cd ~/catkin_ws/src - - ln -s $CI_SOURCE_PATH . - - sudo apt-get install -y python-dev python-numpy python-matplotlib ros-$ROS_DISTRO-tf ros-$ROS_DISTRO-opencv-candidate +- docker build . -t kitti2bag --build-arg ROS_DISTRO=${ROS_DISTRO} -# Install all dependencies, using wstool first and rosdep second. -# wstool looks for a ROSINSTALL_FILE defined in the environment variables. -before_script: - # source dependencies: install using wstool. - - cd ~/catkin_ws/src - - wstool init - - if [[ -f $ROSINSTALL_FILE ]] ; then wstool merge $ROSINSTALL_FILE ; fi - - wstool up - # package depdencies: install using rosdep. - - cd ~/catkin_ws - - rosdep install -y --from-paths src --ignore-src --rosdistro $ROS_DISTRO - -# Compile and test (mark the build as failed if any step fails). If the -# CATKIN_OPTIONS file exists, use it as an argument to catkin_make, for example -# to blacklist certain packages. -# -# NOTE on testing: `catkin_make run_tests` will show the output of the tests -# (gtest, nosetest, etc..) but always returns 0 (success) even if a test -# fails. Running `catkin_test_results` aggregates all the results and returns -# non-zero when a test fails (which notifies Travis the build failed). script: - - source /opt/ros/$ROS_DISTRO/setup.bash - - cd ~/catkin_ws - - catkin_make $( [ -f $CATKIN_OPTIONS ] && cat $CATKIN_OPTIONS ) - # Run the tests, ensuring the path is set correctly. - - source devel/setup.bash - - catkin_make run_tests && catkin_test_results - - cd $CI_SOURCE_PATH - - sudo python setup.py install - - mkdir -p $HOME/kitti - - cd $HOME/kitti - - wget http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0048/2011_09_26_drive_0048_sync.zip - - unzip 2011_09_26_drive_0048_sync.zip - - wget http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_calib.zip - - unzip 2011_09_26_calib.zip - - kitti2bag -t 2011_09_26 -r 0048 raw_synced . - +- docker run -v $TRAVIS_BUILD_DIR/testdata:/kitti2bag/tests/data --entrypoint '/bin/bash' kitti2bag -c ' + source /opt/ros/$ROS_DISTRO/setup.bash && + cd /kitti2bag && + PYTHONIOENCODING=utf-8 python setup.py test --addopts=-svv + ' diff --git a/Dockerfile b/Dockerfile index 04bec59..882e797 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,16 +1,24 @@ -FROM ros:lunar-ros-base +ARG ROS_DISTRO=kinetic +FROM ros:${ROS_DISTRO}-perception -RUN apt-get update \ - && DEBIAN_FRONTEND=noninteractive apt-get -y install \ - ros-lunar-cv-bridge \ - ros-lunar-opencv3 \ - ros-lunar-tf \ - python-pip python-matplotlib \ - && rm -rf /var/lib/apt/lists/* +ARG DEBIAN_FRONTEND=noninteractive +RUN apt-get -y update \ + # && apt-get -y upgrade \ + && apt-get -y install --no-install-recommends \ + ros-${ROS_DISTRO}-cv-bridge \ + ros-${ROS_DISTRO}-tf \ + python-pip \ + && apt-get -y autoremove && apt-get -y clean && rm -rf /var/lib/apt/lists/* +# setup.py will fail on older Ubuntu distros if pip and setuptools are not updated +RUN python -m pip install --upgrade pip setuptools \ + && pip install --upgrade \ + # Upgrade numpy for pykitti's pandas requirement + numpy \ + # Use the development version since a Python 2 incompatibility is not yet fixed in 0.3.1 + git+https://github.com/utiasSTARS/pykitti.git COPY . /kitti2bag -RUN pip install -e /kitti2bag +RUN pip install /kitti2bag WORKDIR /data ENTRYPOINT ["/kitti2bag/docker_entrypoint.sh"] - diff --git a/docker_entrypoint.sh b/docker_entrypoint.sh index d8a3ba6..812db84 100755 --- a/docker_entrypoint.sh +++ b/docker_entrypoint.sh @@ -1,6 +1,10 @@ #!/bin/bash set -e +# unset $@ since setup.bash catches some args like -h +args="$@" +set -- + # setup ros environment source "/opt/ros/$ROS_DISTRO/setup.bash" -exec kitti2bag "$@" +exec kitti2bag $args diff --git a/kitti2bag/__init__.py b/kitti2bag/__init__.py old mode 100644 new mode 100755 index e69de29..cb4c44c --- a/kitti2bag/__init__.py +++ b/kitti2bag/__init__.py @@ -0,0 +1,348 @@ +#!env python +# -*- coding: utf-8 -*- +from __future__ import division, print_function + +import os +import sys +from collections import namedtuple +from datetime import datetime + +import click +import cv2 +import numpy as np +import pykitti +import rosbag +import rospy +import sensor_msgs.point_cloud2 as pcl2 +from cv_bridge import CvBridge +from geometry_msgs.msg import Transform, TransformStamped, TwistStamped +from sensor_msgs.msg import CameraInfo, Imu, NavSatFix, PointField +from std_msgs.msg import Header +from tf.transformations import quaternion_from_euler, quaternion_from_matrix +from tf2_msgs.msg import TFMessage +from tqdm import tqdm + +CameraDetails = namedtuple('CameraDetails', ['nr', 'frame_id', 'topic_id', 'is_rgb']) +cameras = [ + CameraDetails(0, 'camera_gray_left', '/kitti/camera_gray/left', False), + CameraDetails(1, 'camera_gray_right', '/kitti/camera_gray/right', False), + CameraDetails(2, 'camera_color_left', '/kitti/camera_color/left', True), + CameraDetails(3, 'camera_color_right', '/kitti/camera_color/right', True) +] +# All of the cameras share the same camera #0 frame after rectification, +# which is the case for both the raw synced and odometry datasets. +rectified_camera_frame_id = cameras[0].frame_id + + +def to_rostime(dt): + """Convert datetime from python format to ROS format.""" + tsecs = (dt - datetime.utcfromtimestamp(0)).total_seconds() + return rospy.Time.from_sec(tsecs) + + +def read_timestamps(directory): + with open(os.path.join(directory, 'timestamps.txt')) as f: + timestamps = [] + for line in f: + if len(line) == 1: + continue + dt = datetime.strptime(line[:-4], '%Y-%m-%d %H:%M:%S.%f') + timestamps.append(dt) + return timestamps + + +def inv(transform): + """Invert rigid body transformation matrix""" + R = transform[0:3, 0:3] + t = transform[0:3, 3] + t_inv = -1 * R.T.dot(t) + transform_inv = np.eye(4) + transform_inv[0:3, 0:3] = R.T + transform_inv[0:3, 3] = t_inv + return transform_inv + + +def save_imu_data(bag, kitti, imu_frame_id, topic): + print("Exporting IMU") + for timestamp, oxts in zip(kitti.timestamps, kitti.oxts): + q = quaternion_from_euler(oxts.packet.roll, oxts.packet.pitch, oxts.packet.yaw) + imu = Imu() + imu.header.frame_id = imu_frame_id + imu.header.stamp = to_rostime(timestamp) + imu.orientation.x = q[0] + imu.orientation.y = q[1] + imu.orientation.z = q[2] + imu.orientation.w = q[3] + imu.linear_acceleration.x = oxts.packet.af + imu.linear_acceleration.y = oxts.packet.al + imu.linear_acceleration.z = oxts.packet.au + imu.angular_velocity.x = oxts.packet.wf + imu.angular_velocity.y = oxts.packet.wl + imu.angular_velocity.z = oxts.packet.wu + bag.write(topic, imu, t=imu.header.stamp) + + +def save_dynamic_tf(bag, timestamps, tf_matrices, child_frame_id): + print("Exporting time dependent transformations") + for timestamp, tf_matrix in zip(timestamps, tf_matrices): + tf_msg = TFMessage() + tf_stamped = TransformStamped() + tf_stamped.header.stamp = to_rostime(timestamp) + tf_stamped.header.frame_id = 'world' + tf_stamped.child_frame_id = child_frame_id + + t = tf_matrix[0:3, 3] + q = quaternion_from_matrix(tf_matrix) + transform = Transform() + + transform.translation.x = t[0] + transform.translation.y = t[1] + transform.translation.z = t[2] + + transform.rotation.x = q[0] + transform.rotation.y = q[1] + transform.rotation.z = q[2] + transform.rotation.w = q[3] + + tf_stamped.transform = transform + tf_msg.transforms.append(tf_stamped) + + bag.write('/tf', tf_msg, tf_msg.transforms[0].header.stamp) + + +def save_camera_data(bag, kitti, camera, timestamps): + print("Exporting camera {}".format(camera.nr)) + + camera_info = CameraInfo() + camera_info.header.frame_id = rectified_camera_frame_id + camera_info.K = list(getattr(kitti.calib, 'K_cam{}'.format(camera.nr)).flat) + camera_info.P = list(getattr(kitti.calib, 'P_rect_{}0'.format(camera.nr)).flat) + # We do not include the D and R parameters from the calibration data since the images are + # already undistorted and rectified to the camera #0 frame. + camera_info.R = list(np.eye(3).flat) + + cv_bridge = CvBridge() + + image_paths = getattr(kitti, 'cam{}_files'.format(camera.nr)) + for timestamp, image_path in tqdm(list(zip(timestamps, image_paths))): + cv_image = cv2.imread(image_path, cv2.IMREAD_UNCHANGED) + camera_info.height, camera_info.width = cv_image.shape[:2] + encoding = 'bgr8' if camera.is_rgb else 'mono8' + image_message = cv_bridge.cv2_to_imgmsg(cv_image, encoding=encoding) + image_message.header.frame_id = rectified_camera_frame_id + t = to_rostime(timestamp) + image_message.header.stamp = t + camera_info.header.stamp = t + # Follow the naming conventions from + # http://docs.ros.org/melodic/api/sensor_msgs/html/msg/CameraInfo.html + image_topic_ext = '/image_rect_color' if camera.is_rgb else '/image_rect' + bag.write(camera.topic_id + image_topic_ext, image_message, t=t) + bag.write(camera.topic_id + '/camera_info', camera_info, t=t) + + +def save_velo_data(bag, kitti, velo_frame_id, topic): + print("Exporting velodyne data") + velo_path = os.path.join(kitti.data_path, 'velodyne_points') + velo_datetimes = read_timestamps(velo_path) + + for dt, velo_filename in tqdm(list(zip(velo_datetimes, kitti.velo_files))): + if dt is None: + continue + + # read binary data + scan = (np.fromfile(velo_filename, dtype=np.float32)).reshape(-1, 4) + + # create header + header = Header() + header.frame_id = velo_frame_id + header.stamp = to_rostime(dt) + + # fill PCL msg + # uses the PCL PointXYZI struct's field names + # http://docs.pointclouds.org/1.7.0/point__types_8hpp_source.html#l00380 + fields = [PointField('x', 0, PointField.FLOAT32, 1), + PointField('y', 4, PointField.FLOAT32, 1), + PointField('z', 8, PointField.FLOAT32, 1), + PointField('intensity', 12, PointField.FLOAT32, 1)] + pcl_msg = pcl2.create_cloud(header, fields, scan) + + bag.write(topic + '/pointcloud', pcl_msg, t=pcl_msg.header.stamp) + + +def get_static_transform(from_frame_id, to_frame_id, transform): + t = transform[0:3, 3] + q = quaternion_from_matrix(transform) + tf_msg = TransformStamped() + tf_msg.header.frame_id = from_frame_id + tf_msg.child_frame_id = to_frame_id + tf_msg.transform.translation.x = float(t[0]) + tf_msg.transform.translation.y = float(t[1]) + tf_msg.transform.translation.z = float(t[2]) + tf_msg.transform.rotation.x = float(q[0]) + tf_msg.transform.rotation.y = float(q[1]) + tf_msg.transform.rotation.z = float(q[2]) + tf_msg.transform.rotation.w = float(q[3]) + return tf_msg + + +def save_static_transforms(bag, kitti, imu_frame_id, velo_frame_id): + print("Exporting static transformations") + + T_base_link_to_imu = np.eye(4, 4) + T_base_link_to_imu[0:3, 3] = [-2.71 / 2.0 - 0.05, 0.32, 0.93] + + # tf_static + transforms = [ + (imu_frame_id, 'base_link', inv(T_base_link_to_imu)), + (imu_frame_id, velo_frame_id, inv(kitti.calib.T_velo_imu)), + (imu_frame_id, cameras[0].frame_id, inv(kitti.calib.T_cam0_imu)), + (imu_frame_id, cameras[1].frame_id, inv(kitti.calib.T_cam1_imu)), + (imu_frame_id, cameras[2].frame_id, inv(kitti.calib.T_cam2_imu)), + (imu_frame_id, cameras[3].frame_id, inv(kitti.calib.T_cam3_imu)) + ] + + tfm = TFMessage() + for transform in transforms: + t = get_static_transform(from_frame_id=transform[0], to_frame_id=transform[1], transform=transform[2]) + tfm.transforms.append(t) + for timestamp in kitti.timestamps: + time = to_rostime(timestamp) + for transform in tfm.transforms: + transform.header.stamp = time + bag.write('/tf_static', tfm, t=time) + + +def save_gps_fix_data(bag, kitti, gps_frame_id, topic): + for timestamp, oxts in zip(kitti.timestamps, kitti.oxts): + navsatfix_msg = NavSatFix() + navsatfix_msg.header.frame_id = gps_frame_id + navsatfix_msg.header.stamp = to_rostime(timestamp) + navsatfix_msg.latitude = oxts.packet.lat + navsatfix_msg.longitude = oxts.packet.lon + navsatfix_msg.altitude = oxts.packet.alt + navsatfix_msg.status.service = 1 + bag.write(topic, navsatfix_msg, t=navsatfix_msg.header.stamp) + + +def save_gps_vel_data(bag, kitti, gps_frame_id, topic): + for timestamp, oxts in zip(kitti.timestamps, kitti.oxts): + twist_msg = TwistStamped() + twist_msg.header.frame_id = gps_frame_id + twist_msg.header.stamp = to_rostime(timestamp) + twist_msg.twist.linear.x = oxts.packet.vf + twist_msg.twist.linear.y = oxts.packet.vl + twist_msg.twist.linear.z = oxts.packet.vu + twist_msg.twist.angular.x = oxts.packet.wf + twist_msg.twist.angular.y = oxts.packet.wl + twist_msg.twist.angular.z = oxts.packet.wu + bag.write(topic, twist_msg, t=twist_msg.header.stamp) + + +compression_choices = [rosbag.Compression.NONE, rosbag.Compression.BZ2, rosbag.Compression.LZ4] + + +@click.group() +@click.help_option('-h', '--help') +def cli(): + """Convert KITTI dataset to ROS bag file the easy way!""" + pass + + +def common_options(f): + f = click.option("-i", "--input-dir", required=False, default='.', show_default=True, + type=click.Path(exists=True, dir_okay=True, file_okay=False), + help="base directory of the dataset")(f) + f = click.option("-o", "--output-dir", required=False, default='.', show_default=True, + type=click.Path(exists=True, dir_okay=True, file_okay=False), + help="output directory for the created bag file")(f) + f = click.option("--compression", required=False, + type=click.Choice(compression_choices), + default=rosbag.Compression.NONE, show_default=True, + help="which compression to use for the created bag")(f) + f = click.help_option('-h', '--help')(f) + return f + + +@cli.command('raw') +@click.option("-t", "--date", required=True, metavar='DATE', + help="date of the raw dataset (i.e. 2011_09_26)") +@click.option("-r", "--drive", required=True, type=int, + help="drive number of the raw dataset (i.e. 1)") +@common_options +def convert_raw(date, drive, input_dir='.', output_dir='.', compression=rosbag.Compression.NONE): + """Convert a raw synced+rectified KITTI dataset to a bag""" + drive = str(drive).zfill(4) + kitti = pykitti.raw(input_dir, date, drive) + + if len(kitti.timestamps) == 0: + print('Dataset is empty? Exiting.', file=sys.stderr) + sys.exit(1) + + # IMU + imu_frame_id = 'imu_link' + imu_topic = '/kitti/oxts/imu' + gps_fix_topic = '/kitti/oxts/gps/fix' + gps_vel_topic = '/kitti/oxts/gps/vel' + velo_frame_id = 'velo_link' + velo_topic = '/kitti/velo' + + # Export + bag_name = "kitti_{}_drive_{}_synced.bag".format(date, drive) + with rosbag.Bag(os.path.join(output_dir, bag_name), 'w', compression=compression) as bag: + save_static_transforms(bag, kitti, imu_frame_id, velo_frame_id) + imu_tf_matrices = [oxt.T_w_imu for oxt in kitti.oxts] + save_dynamic_tf(bag, kitti.timestamps, imu_tf_matrices, imu_frame_id) + save_imu_data(bag, kitti, imu_frame_id, imu_topic) + save_gps_fix_data(bag, kitti, imu_frame_id, gps_fix_topic) + save_gps_vel_data(bag, kitti, imu_frame_id, gps_vel_topic) + for camera in cameras: + camera_dir = os.path.join(kitti.data_path, 'image_{0:02d}'.format(camera.nr)) + timestamps = read_timestamps(camera_dir) + save_camera_data(bag, kitti, camera, timestamps) + save_velo_data(bag, kitti, velo_frame_id, velo_topic) + + print("## OVERVIEW ##") + print(bag) + + +@cli.command('odom') +@click.option("-s", "--sequence", type=click.IntRange(0, 21), metavar='0...21', required=True, + help="sequence number") +@click.option("-c", "--color", type=click.Choice(['gray', 'color', 'all']), required=True, + help="which camera images to include in the bag") +@common_options +def convert_odom(sequence, color, input_dir='.', output_dir='.', compression=rosbag.Compression.NONE): + """Convert an odometry KITTI dataset to a bag""" + sequence_str = str(sequence).zfill(2) + kitti = pykitti.odometry(input_dir, sequence_str) + + if len(kitti.timestamps) == 0: + print('Dataset is empty? Exiting.', file=sys.stderr) + sys.exit(1) + + # The odom timestamps are relative, add an arbitrary datetime to make them absolute + base_timestamp = datetime(2011, 9, 26, 12, 0, 0) + timestamps = [base_timestamp + timestamp for timestamp in kitti.timestamps] + + if 0 <= sequence < 10: + print("Odometry dataset sequence {} has ground truth information (poses).".format(sequence_str)) + + if color == 'color': + camera_nrs = (2, 3) + elif color == 'gray': + camera_nrs = (0, 1) + else: + camera_nrs = list(range(4)) + + # Export + bag_name = "kitti_data_odometry_{}_sequence_{}.bag".format(color, sequence_str) + with rosbag.Bag(os.path.join(output_dir, bag_name), 'w', compression=compression) as bag: + save_dynamic_tf(bag, timestamps, kitti.poses, rectified_camera_frame_id) + for camera_nr in camera_nrs: + save_camera_data(bag, kitti, cameras[camera_nr], timestamps) + print("## OVERVIEW ##") + print(bag) + + +if __name__ == '__main__': + cli() diff --git a/kitti2bag/__main__.py b/kitti2bag/__main__.py deleted file mode 100644 index 66010d1..0000000 --- a/kitti2bag/__main__.py +++ /dev/null @@ -1,10 +0,0 @@ -from .kitti2bag import run_kitti2bag - - -def main(): - run_kitti2bag() - - -if __name__ == '__main__': - main() - diff --git a/kitti2bag/kitti2bag.py b/kitti2bag/kitti2bag.py deleted file mode 100755 index ea0e81a..0000000 --- a/kitti2bag/kitti2bag.py +++ /dev/null @@ -1,391 +0,0 @@ -#!env python -# -*- coding: utf-8 -*- - -import sys - -try: - import pykitti -except ImportError as e: - print('Could not load module \'pykitti\'. Please run `pip install pykitti`') - sys.exit(1) - -import tf -import os -import cv2 -import rospy -import rosbag -import progressbar -from tf2_msgs.msg import TFMessage -from datetime import datetime -from std_msgs.msg import Header -from sensor_msgs.msg import CameraInfo, Imu, PointField, NavSatFix -import sensor_msgs.point_cloud2 as pcl2 -from geometry_msgs.msg import TransformStamped, TwistStamped, Transform -from cv_bridge import CvBridge -import numpy as np -import argparse - -def save_imu_data(bag, kitti, imu_frame_id, topic): - print("Exporting IMU") - for timestamp, oxts in zip(kitti.timestamps, kitti.oxts): - q = tf.transformations.quaternion_from_euler(oxts.packet.roll, oxts.packet.pitch, oxts.packet.yaw) - imu = Imu() - imu.header.frame_id = imu_frame_id - imu.header.stamp = rospy.Time.from_sec(float(timestamp.strftime("%s.%f"))) - imu.orientation.x = q[0] - imu.orientation.y = q[1] - imu.orientation.z = q[2] - imu.orientation.w = q[3] - imu.linear_acceleration.x = oxts.packet.af - imu.linear_acceleration.y = oxts.packet.al - imu.linear_acceleration.z = oxts.packet.au - imu.angular_velocity.x = oxts.packet.wf - imu.angular_velocity.y = oxts.packet.wl - imu.angular_velocity.z = oxts.packet.wu - bag.write(topic, imu, t=imu.header.stamp) - - -def save_dynamic_tf(bag, kitti, kitti_type, initial_time): - print("Exporting time dependent transformations") - if kitti_type.find("raw") != -1: - for timestamp, oxts in zip(kitti.timestamps, kitti.oxts): - tf_oxts_msg = TFMessage() - tf_oxts_transform = TransformStamped() - tf_oxts_transform.header.stamp = rospy.Time.from_sec(float(timestamp.strftime("%s.%f"))) - tf_oxts_transform.header.frame_id = 'world' - tf_oxts_transform.child_frame_id = 'base_link' - - transform = (oxts.T_w_imu) - t = transform[0:3, 3] - q = tf.transformations.quaternion_from_matrix(transform) - oxts_tf = Transform() - - oxts_tf.translation.x = t[0] - oxts_tf.translation.y = t[1] - oxts_tf.translation.z = t[2] - - oxts_tf.rotation.x = q[0] - oxts_tf.rotation.y = q[1] - oxts_tf.rotation.z = q[2] - oxts_tf.rotation.w = q[3] - - tf_oxts_transform.transform = oxts_tf - tf_oxts_msg.transforms.append(tf_oxts_transform) - - bag.write('/tf', tf_oxts_msg, tf_oxts_msg.transforms[0].header.stamp) - - elif kitti_type.find("odom") != -1: - timestamps = map(lambda x: initial_time + x.total_seconds(), kitti.timestamps) - for timestamp, tf_matrix in zip(timestamps, kitti.T_w_cam0): - tf_msg = TFMessage() - tf_stamped = TransformStamped() - tf_stamped.header.stamp = rospy.Time.from_sec(timestamp) - tf_stamped.header.frame_id = 'world' - tf_stamped.child_frame_id = 'camera_left' - - t = tf_matrix[0:3, 3] - q = tf.transformations.quaternion_from_matrix(tf_matrix) - transform = Transform() - - transform.translation.x = t[0] - transform.translation.y = t[1] - transform.translation.z = t[2] - - transform.rotation.x = q[0] - transform.rotation.y = q[1] - transform.rotation.z = q[2] - transform.rotation.w = q[3] - - tf_stamped.transform = transform - tf_msg.transforms.append(tf_stamped) - - bag.write('/tf', tf_msg, tf_msg.transforms[0].header.stamp) - - -def save_camera_data(bag, kitti_type, kitti, util, bridge, camera, camera_frame_id, topic, initial_time): - print("Exporting camera {}".format(camera)) - if kitti_type.find("raw") != -1: - camera_pad = '{0:02d}'.format(camera) - image_dir = os.path.join(kitti.data_path, 'image_{}'.format(camera_pad)) - image_path = os.path.join(image_dir, 'data') - image_filenames = sorted(os.listdir(image_path)) - with open(os.path.join(image_dir, 'timestamps.txt')) as f: - image_datetimes = map(lambda x: datetime.strptime(x[:-4], '%Y-%m-%d %H:%M:%S.%f'), f.readlines()) - - calib = CameraInfo() - calib.header.frame_id = camera_frame_id - calib.width, calib.height = tuple(util['S_rect_{}'.format(camera_pad)].tolist()) - calib.distortion_model = 'plumb_bob' - calib.K = util['K_{}'.format(camera_pad)] - calib.R = util['R_rect_{}'.format(camera_pad)] - calib.D = util['D_{}'.format(camera_pad)] - calib.P = util['P_rect_{}'.format(camera_pad)] - - elif kitti_type.find("odom") != -1: - camera_pad = '{0:01d}'.format(camera) - image_path = os.path.join(kitti.sequence_path, 'image_{}'.format(camera_pad)) - image_filenames = sorted(os.listdir(image_path)) - image_datetimes = map(lambda x: initial_time + x.total_seconds(), kitti.timestamps) - - calib = CameraInfo() - calib.header.frame_id = camera_frame_id - calib.P = util['P{}'.format(camera_pad)] - - iterable = zip(image_datetimes, image_filenames) - bar = progressbar.ProgressBar() - for dt, filename in bar(iterable): - image_filename = os.path.join(image_path, filename) - cv_image = cv2.imread(image_filename) - calib.height, calib.width = cv_image.shape[:2] - if camera in (0, 1): - cv_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2GRAY) - encoding = "mono8" if camera in (0, 1) else "bgr8" - image_message = bridge.cv2_to_imgmsg(cv_image, encoding=encoding) - image_message.header.frame_id = camera_frame_id - if kitti_type.find("raw") != -1: - image_message.header.stamp = rospy.Time.from_sec(float(datetime.strftime(dt, "%s.%f"))) - topic_ext = "/image_raw" - elif kitti_type.find("odom") != -1: - image_message.header.stamp = rospy.Time.from_sec(dt) - topic_ext = "/image_rect" - calib.header.stamp = image_message.header.stamp - bag.write(topic + topic_ext, image_message, t = image_message.header.stamp) - bag.write(topic + '/camera_info', calib, t = calib.header.stamp) - -def save_velo_data(bag, kitti, velo_frame_id, topic): - print("Exporting velodyne data") - velo_path = os.path.join(kitti.data_path, 'velodyne_points') - velo_data_dir = os.path.join(velo_path, 'data') - velo_filenames = sorted(os.listdir(velo_data_dir)) - with open(os.path.join(velo_path, 'timestamps.txt')) as f: - lines = f.readlines() - velo_datetimes = [] - for line in lines: - if len(line) == 1: - continue - dt = datetime.strptime(line[:-4], '%Y-%m-%d %H:%M:%S.%f') - velo_datetimes.append(dt) - - iterable = zip(velo_datetimes, velo_filenames) - bar = progressbar.ProgressBar() - for dt, filename in bar(iterable): - if dt is None: - continue - - velo_filename = os.path.join(velo_data_dir, filename) - - # read binary data - scan = (np.fromfile(velo_filename, dtype=np.float32)).reshape(-1, 4) - - # create header - header = Header() - header.frame_id = velo_frame_id - header.stamp = rospy.Time.from_sec(float(datetime.strftime(dt, "%s.%f"))) - - # fill pcl msg - fields = [PointField('x', 0, PointField.FLOAT32, 1), - PointField('y', 4, PointField.FLOAT32, 1), - PointField('z', 8, PointField.FLOAT32, 1), - PointField('i', 12, PointField.FLOAT32, 1)] - pcl_msg = pcl2.create_cloud(header, fields, scan) - - bag.write(topic + '/pointcloud', pcl_msg, t=pcl_msg.header.stamp) - - -def get_static_transform(from_frame_id, to_frame_id, transform): - t = transform[0:3, 3] - q = tf.transformations.quaternion_from_matrix(transform) - tf_msg = TransformStamped() - tf_msg.header.frame_id = from_frame_id - tf_msg.child_frame_id = to_frame_id - tf_msg.transform.translation.x = float(t[0]) - tf_msg.transform.translation.y = float(t[1]) - tf_msg.transform.translation.z = float(t[2]) - tf_msg.transform.rotation.x = float(q[0]) - tf_msg.transform.rotation.y = float(q[1]) - tf_msg.transform.rotation.z = float(q[2]) - tf_msg.transform.rotation.w = float(q[3]) - return tf_msg - - -def inv(transform): - "Invert rigid body transformation matrix" - R = transform[0:3, 0:3] - t = transform[0:3, 3] - t_inv = -1 * R.T.dot(t) - transform_inv = np.eye(4) - transform_inv[0:3, 0:3] = R.T - transform_inv[0:3, 3] = t_inv - return transform_inv - - -def save_static_transforms(bag, transforms, timestamps): - print("Exporting static transformations") - tfm = TFMessage() - for transform in transforms: - t = get_static_transform(from_frame_id=transform[0], to_frame_id=transform[1], transform=transform[2]) - tfm.transforms.append(t) - for timestamp in timestamps: - time = rospy.Time.from_sec(float(timestamp.strftime("%s.%f"))) - for i in range(len(tfm.transforms)): - tfm.transforms[i].header.stamp = time - bag.write('/tf_static', tfm, t=time) - - -def save_gps_fix_data(bag, kitti, gps_frame_id, topic): - for timestamp, oxts in zip(kitti.timestamps, kitti.oxts): - navsatfix_msg = NavSatFix() - navsatfix_msg.header.frame_id = gps_frame_id - navsatfix_msg.header.stamp = rospy.Time.from_sec(float(timestamp.strftime("%s.%f"))) - navsatfix_msg.latitude = oxts.packet.lat - navsatfix_msg.longitude = oxts.packet.lon - navsatfix_msg.altitude = oxts.packet.alt - navsatfix_msg.status.service = 1 - bag.write(topic, navsatfix_msg, t=navsatfix_msg.header.stamp) - - -def save_gps_vel_data(bag, kitti, gps_frame_id, topic): - for timestamp, oxts in zip(kitti.timestamps, kitti.oxts): - twist_msg = TwistStamped() - twist_msg.header.frame_id = gps_frame_id - twist_msg.header.stamp = rospy.Time.from_sec(float(timestamp.strftime("%s.%f"))) - twist_msg.twist.linear.x = oxts.packet.vf - twist_msg.twist.linear.y = oxts.packet.vl - twist_msg.twist.linear.z = oxts.packet.vu - twist_msg.twist.angular.x = oxts.packet.wf - twist_msg.twist.angular.y = oxts.packet.wl - twist_msg.twist.angular.z = oxts.packet.wu - bag.write(topic, twist_msg, t=twist_msg.header.stamp) - - -def run_kitti2bag(): - parser = argparse.ArgumentParser(description = "Convert KITTI dataset to ROS bag file the easy way!") - # Accepted argument values - kitti_types = ["raw_synced", "odom_color", "odom_gray"] - odometry_sequences = [] - for s in range(22): - odometry_sequences.append(str(s).zfill(2)) - - parser.add_argument("kitti_type", choices = kitti_types, help = "KITTI dataset type") - parser.add_argument("dir", nargs = "?", default = os.getcwd(), help = "base directory of the dataset, if no directory passed the deafult is current working directory") - parser.add_argument("-t", "--date", help = "date of the raw dataset (i.e. 2011_09_26), option is only for RAW datasets.") - parser.add_argument("-r", "--drive", help = "drive number of the raw dataset (i.e. 0001), option is only for RAW datasets.") - parser.add_argument("-s", "--sequence", choices = odometry_sequences,help = "sequence of the odometry dataset (between 00 - 21), option is only for ODOMETRY datasets.") - args = parser.parse_args() - - bridge = CvBridge() - compression = rosbag.Compression.NONE - # compression = rosbag.Compression.BZ2 - # compression = rosbag.Compression.LZ4 - - # CAMERAS - cameras = [ - (0, 'camera_gray_left', '/kitti/camera_gray_left'), - (1, 'camera_gray_right', '/kitti/camera_gray_right'), - (2, 'camera_color_left', '/kitti/camera_color_left'), - (3, 'camera_color_right', '/kitti/camera_color_right') - ] - - if args.kitti_type.find("raw") != -1: - - if args.date == None: - print("Date option is not given. It is mandatory for raw dataset.") - print("Usage for raw dataset: kitti2bag raw_synced [dir] -t -r ") - sys.exit(1) - elif args.drive == None: - print("Drive option is not given. It is mandatory for raw dataset.") - print("Usage for raw dataset: kitti2bag raw_synced [dir] -t -r ") - sys.exit(1) - - bag = rosbag.Bag("kitti_{}_drive_{}_{}.bag".format(args.date, args.drive, args.kitti_type[4:]), 'w', compression=compression) - kitti = pykitti.raw(args.dir, args.date, args.drive) - if not os.path.exists(kitti.data_path): - print('Path {} does not exists. Exiting.'.format(kitti.data_path)) - sys.exit(1) - - if len(kitti.timestamps) == 0: - print('Dataset is empty? Exiting.') - sys.exit(1) - - try: - # IMU - imu_frame_id = 'imu_link' - imu_topic = '/kitti/oxts/imu' - gps_fix_topic = '/kitti/oxts/gps/fix' - gps_vel_topic = '/kitti/oxts/gps/vel' - velo_frame_id = 'velo_link' - velo_topic = '/kitti/velo' - - T_base_link_to_imu = np.eye(4, 4) - T_base_link_to_imu[0:3, 3] = [-2.71/2.0-0.05, 0.32, 0.93] - - # tf_static - transforms = [ - ('base_link', imu_frame_id, T_base_link_to_imu), - (imu_frame_id, velo_frame_id, inv(kitti.calib.T_velo_imu)), - (imu_frame_id, cameras[0][1], inv(kitti.calib.T_cam0_imu)), - (imu_frame_id, cameras[1][1], inv(kitti.calib.T_cam1_imu)), - (imu_frame_id, cameras[2][1], inv(kitti.calib.T_cam2_imu)), - (imu_frame_id, cameras[3][1], inv(kitti.calib.T_cam3_imu)) - ] - - util = pykitti.utils.read_calib_file(os.path.join(kitti.calib_path, 'calib_cam_to_cam.txt')) - - # Export - save_static_transforms(bag, transforms, kitti.timestamps) - save_dynamic_tf(bag, kitti, args.kitti_type, initial_time=None) - save_imu_data(bag, kitti, imu_frame_id, imu_topic) - save_gps_fix_data(bag, kitti, imu_frame_id, gps_fix_topic) - save_gps_vel_data(bag, kitti, imu_frame_id, gps_vel_topic) - for camera in cameras: - save_camera_data(bag, args.kitti_type, kitti, util, bridge, camera=camera[0], camera_frame_id=camera[1], topic=camera[2], initial_time=None) - save_velo_data(bag, kitti, velo_frame_id, velo_topic) - - finally: - print("## OVERVIEW ##") - print(bag) - bag.close() - - elif args.kitti_type.find("odom") != -1: - - if args.sequence == None: - print("Sequence option is not given. It is mandatory for odometry dataset.") - print("Usage for odometry dataset: kitti2bag {odom_color, odom_gray} [dir] -s ") - sys.exit(1) - - bag = rosbag.Bag("kitti_data_odometry_{}_sequence_{}.bag".format(args.kitti_type[5:],args.sequence), 'w', compression=compression) - - kitti = pykitti.odometry(args.dir, args.sequence) - if not os.path.exists(kitti.sequence_path): - print('Path {} does not exists. Exiting.'.format(kitti.sequence_path)) - sys.exit(1) - - kitti.load_calib() - kitti.load_timestamps() - - if len(kitti.timestamps) == 0: - print('Dataset is empty? Exiting.') - sys.exit(1) - - if args.sequence in odometry_sequences[:11]: - print("Odometry dataset sequence {} has ground truth information (poses).".format(args.sequence)) - kitti.load_poses() - - try: - util = pykitti.utils.read_calib_file(os.path.join(args.dir,'sequences',args.sequence, 'calib.txt')) - current_epoch = (datetime.utcnow() - datetime(1970, 1, 1)).total_seconds() - # Export - if args.kitti_type.find("gray") != -1: - used_cameras = cameras[:2] - elif args.kitti_type.find("color") != -1: - used_cameras = cameras[-2:] - - save_dynamic_tf(bag, kitti, args.kitti_type, initial_time=current_epoch) - for camera in used_cameras: - save_camera_data(bag, args.kitti_type, kitti, util, bridge, camera=camera[0], camera_frame_id=camera[1], topic=camera[2], initial_time=current_epoch) - - finally: - print("## OVERVIEW ##") - print(bag) - bag.close() - diff --git a/setup.cfg b/setup.cfg index b88034e..6afcf47 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,2 +1,5 @@ [metadata] description-file = README.md + +[aliases] +test=pytest diff --git a/setup.py b/setup.py old mode 100644 new mode 100755 index 1250f79..b9ac076 --- a/setup.py +++ b/setup.py @@ -9,10 +9,13 @@ author='Tomas Krejci', author_email='tomas@krej.ci', url='https://github.com/tomas789/kitti2bag/', - download_url = 'https://github.com/tomas789/kitti2bag/archive/1.5.zip', - keywords = ['dataset', 'ros', 'rosbag', 'kitti'], - entry_points = { - 'console_scripts': ['kitti2bag=kitti2bag.__main__:main'], + download_url='https://github.com/tomas789/kitti2bag/archive/1.5.zip', + keywords=['dataset', 'ros', 'rosbag', 'kitti'], + packages=['kitti2bag'], + entry_points={ + 'console_scripts': ['kitti2bag=kitti2bag:cli'], }, - install_requires=['pykitti', 'progressbar2'] + install_requires=['pykitti', 'tqdm', 'click'], + setup_requires=['pytest-runner'], + tests_require=['pytest', 'requests', 'six', 'pyyaml', 'click>=7'] ) diff --git a/tests/test_kitti2bag.py b/tests/test_kitti2bag.py new file mode 100644 index 0000000..118c6d2 --- /dev/null +++ b/tests/test_kitti2bag.py @@ -0,0 +1,250 @@ +import os +import shutil +from os.path import abspath, dirname, exists, join +from zipfile import ZipFile + +import pytest +import requests +import rosbag +import tqdm +import yaml +from click.testing import CliRunner +from six import BytesIO + +# Disable tqdm background monitor which does not exit cleanly with pytest in python 2 +tqdm.monitor_interval = 0 + +from kitti2bag import cli + +TESTS_DIR = dirname(abspath(__file__)) +DATA_DIR = join(TESTS_DIR, 'data') + +DATA_URL_BASE = 'https://s3.eu-central-1.amazonaws.com/avg-kitti' + + +def download_and_extract(filename, output_dir): + url = DATA_URL_BASE + filename + print('Downloading {}...'.format(url)) + response = requests.get(url) + response.raise_for_status() + zipfile = ZipFile(BytesIO(response.content)) + zipfile.extractall(output_dir) + + +@pytest.fixture(scope='session') +def raw_data(): + raw_data_dir = join(DATA_DIR, 'raw') + params = { + 'dir': raw_data_dir, + 'date': '2011_09_26', + 'drive': 48, + 'drive_dir': join(raw_data_dir, '2011_09_26/2011_09_26_drive_0048_sync') + } + if exists(raw_data_dir): + return params + + os.makedirs(raw_data_dir) + download_and_extract('/raw_data/2011_09_26_drive_0048/2011_09_26_drive_0048_sync.zip', raw_data_dir) + download_and_extract('/raw_data/2011_09_26_calib.zip', raw_data_dir) + + return params + + +@pytest.fixture(scope='session') +def odom_data(raw_data): + odom_data_dir = join(DATA_DIR, 'odom') + params = { + 'dir': odom_data_dir, + 'sequence': 4, + 'sequence_dir': join(odom_data_dir, 'sequences/04') + } + if exists(odom_data_dir): + return params + + if not exists(DATA_DIR): + os.makedirs(DATA_DIR) + download_and_extract('/data_odometry_calib.zip', DATA_DIR) + shutil.move(join(DATA_DIR, 'dataset'), odom_data_dir) + + # create dummy image data for odom + for i in range(4): + shutil.copytree( + join(raw_data['drive_dir'], 'image_0{}/data').format(i), + join(params['sequence_dir'], 'image_{}').format(i) + ) + + return params + + +def clean_bag_info(info): + del info['path'] + del info['version'] + del info['size'] + for t in info['types']: + del t['md5'] + return info + + +def test_raw_synced(raw_data, tmpdir): + runner = CliRunner() + result = runner.invoke( + cli, [ + 'raw', + '--date', raw_data['date'], + '--drive', raw_data['drive'], + '--input-dir', raw_data['dir'], + '--output-dir', str(tmpdir), + ], + catch_exceptions=False + ) + print(result.stdout_bytes.decode('utf8')) + assert result.exit_code == 0 + + expected_bagfile = tmpdir.join('kitti_2011_09_26_drive_0048_synced.bag') + assert expected_bagfile.exists() + with rosbag.Bag(str(expected_bagfile), 'r') as bag: + info = yaml.safe_load(bag._get_yaml_info()) + + # remove irrelevant fields + info = clean_bag_info(info) + + expected = { + 'compression': 'none', + 'duration': 2.151645, + 'end': 1317046453.072943, + 'indexed': True, + 'messages': 308, + 'start': 1317046450.921298, + 'topics': [ + {'frequency': 9.6728, + 'messages': 22, + 'topic': '/kitti/camera_color/left/camera_info', + 'type': 'sensor_msgs/CameraInfo'}, + {'frequency': 9.6728, + 'messages': 22, + 'topic': '/kitti/camera_color/left/image_rect_color', + 'type': 'sensor_msgs/Image'}, + {'frequency': 9.6725, + 'messages': 22, + 'topic': '/kitti/camera_color/right/camera_info', + 'type': 'sensor_msgs/CameraInfo'}, + {'frequency': 9.6725, + 'messages': 22, + 'topic': '/kitti/camera_color/right/image_rect_color', + 'type': 'sensor_msgs/Image'}, + {'frequency': 9.6725, + 'messages': 22, + 'topic': '/kitti/camera_gray/left/camera_info', + 'type': 'sensor_msgs/CameraInfo'}, + {'frequency': 9.6725, + 'messages': 22, + 'topic': '/kitti/camera_gray/left/image_rect', + 'type': 'sensor_msgs/Image'}, + {'frequency': 9.6709, + 'messages': 22, + 'topic': '/kitti/camera_gray/right/camera_info', + 'type': 'sensor_msgs/CameraInfo'}, + {'frequency': 9.6709, + 'messages': 22, + 'topic': '/kitti/camera_gray/right/image_rect', + 'type': 'sensor_msgs/Image'}, + {'frequency': 9.9951, + 'messages': 22, + 'topic': '/kitti/oxts/gps/fix', + 'type': 'sensor_msgs/NavSatFix'}, + {'frequency': 9.9951, + 'messages': 22, + 'topic': '/kitti/oxts/gps/vel', + 'type': 'geometry_msgs/TwistStamped'}, + {'frequency': 9.9951, + 'messages': 22, + 'topic': '/kitti/oxts/imu', + 'type': 'sensor_msgs/Imu'}, + {'frequency': 9.6715, + 'messages': 22, + 'topic': '/kitti/velo/pointcloud', + 'type': 'sensor_msgs/PointCloud2'}, + {'frequency': 9.9951, + 'messages': 22, + 'topic': '/tf', + 'type': 'tf2_msgs/TFMessage'}, + {'frequency': 9.9951, + 'messages': 22, + 'topic': '/tf_static', + 'type': 'tf2_msgs/TFMessage'} + ], + 'types': [ + {'type': 'geometry_msgs/TwistStamped'}, + {'type': 'sensor_msgs/CameraInfo'}, + {'type': 'sensor_msgs/Image'}, + {'type': 'sensor_msgs/Imu'}, + {'type': 'sensor_msgs/NavSatFix'}, + {'type': 'sensor_msgs/PointCloud2'}, + {'type': 'tf2_msgs/TFMessage'} + ] + } + + assert info == expected + + tmpdir.remove() + + +@pytest.mark.parametrize('color', ['gray', 'color']) +def test_odom(odom_data, tmpdir, color): + runner = CliRunner() + result = runner.invoke( + cli, [ + 'odom', + '--sequence', odom_data['sequence'], + '--color', color, + '--input-dir', odom_data['dir'], + '--output-dir', str(tmpdir), + ], + catch_exceptions=False + ) + print(result.stdout_bytes.decode('utf8')) + assert result.exit_code == 0 + + expected_bagfile = tmpdir.join('kitti_data_odometry_{}_sequence_04.bag'.format(color)) + assert expected_bagfile.exists() + with rosbag.Bag(str(expected_bagfile), 'r') as bag: + info = yaml.safe_load(bag._get_yaml_info()) + + info = clean_bag_info(info) + + expected = { + 'compression': 'none', + 'duration': 2.187759, + 'end': 1317038402.187759, + 'indexed': True, + 'messages': 88, + 'start': 1317038400.0, + 'topics': [ + {'frequency': 9.6034, + 'messages': 22, + 'topic': '/kitti/camera_{}/left/camera_info'.format(color), + 'type': 'sensor_msgs/CameraInfo'}, + {'frequency': 9.6034, + 'messages': 22, + 'topic': '/kitti/camera_{}/left/image_rect{}'.format( + color, '' if color == 'gray' else '_color'), + 'type': 'sensor_msgs/Image'}, + {'frequency': 9.6034, + 'messages': 22, + 'topic': '/kitti/camera_{}/right/camera_info'.format(color), + 'type': 'sensor_msgs/CameraInfo'}, + {'frequency': 9.6034, + 'messages': 22, + 'topic': '/kitti/camera_{}/right/image_rect{}'.format( + color, '' if color == 'gray' else '_color'), + 'type': 'sensor_msgs/Image'} + ], + 'types': [ + {'type': 'sensor_msgs/CameraInfo'}, + {'type': 'sensor_msgs/Image'} + ] + } + + assert info == expected + + tmpdir.remove()