Ros bag to Apollo record or reverse conversion
#!/usr/bin/env python

# Copyright 2021 daohu527 <daohu527@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging

import rosbag
from cyber_record.record import Record
from tqdm import tqdm

from bag_convert.bag2record.imu import to_imu
from bag_convert.bag2record.localization import to_localization
from bag_convert.bag2record.pointcloud import to_pointcloud
from bag_convert.bag2record.sensor_image import to_image, to_compressed_image


def convert_msg(ros_msg):
    msg_type = ros_msg.__class__.__name__
    if "Imu" in msg_type:
        return to_imu(ros_msg)
    elif "Pose" in msg_type:
        return to_localization(ros_msg)
    elif "PointCloud2" in msg_type:
        return to_pointcloud(ros_msg)
    elif "Image" in msg_type:
        return to_image(ros_msg)
    elif "CompressedImage" in msg_type:
        return to_compressed_image(ros_msg)
    else:
        return None


def convert(bag_file, record_file="result.record"):
    bag = rosbag.Bag(bag_file)
    pbar = tqdm(desc="bag to record progress")
    with Record(record_file, mode='w') as record:
        for topic, msg, t in bag.read_messages():
            logging.debug("{},{},{}".format(topic, type(msg), t))
            msg = convert_msg(msg)
            t = t.secs * (10**9) + t.nsecs
            if msg:
                record.write(topic, msg, t)
            pbar.update()
    bag.close()
    pbar.close()


if __name__ == "__main__":
    convert("../../data/b0-2014-07-11-10-58-16.bag")

bag_convert/bag2record/bag2record.py
daohu527/bag_convert
#!/usr/bin/env python

# Copyright 2021 daohu527 <daohu527@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from modules.common.proto.header_pb2 import Header


def to_timestamp(ros_stamp):
    return ros_stamp.secs + (ros_stamp.nsecs / 1e9)


def to_header(ros_header):
    cyber_header = Header()
    cyber_header.sequence_num = ros_header.seq
    cyber_header.timestamp_sec = to_timestamp(ros_header.stamp)
    cyber_header.frame_id = ros_header.frame_id
    return cyber_header


def add_header(func):
    def inner(ros_msg):
        cyber_msg = func(ros_msg)
        cyber_msg.header.CopyFrom(to_header(ros_msg.header))
        return cyber_msg
    return inner

bag_convert/bag2record/header.py
daohu527/bag_convert
#!/usr/bin/env python

# Copyright 2021 daohu527 <daohu527@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from modules.localization.proto.imu_pb2 import CorrectedImu

from bag_convert.bag2record.header import add_header


@add_header
def to_imu(ros_imu):
    cyber_imu = CorrectedImu()

    pose = cyber_imu.imu
    pose.orientation.qx = ros_imu.orientation.x
    pose.orientation.qy = ros_imu.orientation.y
    pose.orientation.qz = ros_imu.orientation.z
    pose.orientation.qw = ros_imu.orientation.w

    pose.angular_velocity.x = ros_imu.angular_velocity.x
    pose.angular_velocity.y = ros_imu.angular_velocity.y
    pose.angular_velocity.z = ros_imu.angular_velocity.z

    pose.linear_acceleration.x = ros_imu.linear_acceleration.x
    pose.linear_acceleration.y = ros_imu.linear_acceleration.y
    pose.linear_acceleration.z = ros_imu.linear_acceleration.z
    return cyber_imu

bag_convert/bag2record/imu.py
daohu527/bag_convert
#!/usr/bin/env python

# Copyright 2021 daohu527 <daohu527@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from modules.localization.proto.localization_pb2 import LocalizationEstimate

from bag_convert.bag2record.header import add_header


@add_header
def to_localization(ros_pose):
    cyber_localization = LocalizationEstimate()
    # cyber_localization.header.CopyFrom(to_header(ros_pose.header))

    pose = cyber_localization.pose
    pose.position.x = ros_pose.position.x
    pose.position.y = ros_pose.position.y
    pose.position.z = ros_pose.position.z

    pose.orientation.qx = ros_pose.orientation.x
    pose.orientation.qy = ros_pose.orientation.y
    pose.orientation.qz = ros_pose.orientation.z
    pose.orientation.qw = ros_pose.orientation.w
    return cyber_localization

bag_convert/bag2record/localization.py
daohu527/bag_convert
#!/usr/bin/env python

# Copyright 2021 daohu527 <daohu527@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from modules.drivers.proto.pointcloud_pb2 import PointCloud
from sensor_msgs import point_cloud2

from bag_convert.bag2record.header import add_header


@add_header
def to_pointcloud(ros_pointcloud2):
    cyber_pointcloud = PointCloud()

    cyber_pointcloud.height = ros_pointcloud2.height
    cyber_pointcloud.width = ros_pointcloud2.width

    # todo(zero): Need to parse "ros_pointcloud2.fields" for assignment
    for p in point_cloud2.read_points(ros_pointcloud2):
        point_xyzit = cyber_pointcloud.point.add()
        point_xyzit.x = p[0]
        point_xyzit.y = p[1]
        point_xyzit.z = p[2]
        if len(p) > 3:
            point_xyzit.intensity = p[3]
        if len(p) > 4:
            point_xyzit.timestamp = p[4]

    cyber_pointcloud.is_dense = ros_pointcloud2.is_dense
    return cyber_pointcloud

bag_convert/bag2record/pointcloud.py
daohu527/bag_convert
#!/usr/bin/env python

# Copyright 2021 daohu527 <daohu527@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from modules.drivers.proto.sensor_image_pb2 import Image, CompressedImage

from bag_convert.bag2record.header import add_header


@add_header
def to_image(ros_image):
    cyber_image = Image()
    # cyber_image.header = to_header(ros_image.header)

    cyber_image.height = ros_image.height
    cyber_image.width = ros_image.width
    cyber_image.encoding = ros_image.encoding
    cyber_image.step = ros_image.step

    cyber_image.data = ros_image.data
    return cyber_image


@add_header
def to_compressed_image(ros_compressed_image):
    cyber_compressed_image = CompressedImage()
    # cyber_compressed_image.header = to_header(ros_compressed_image.header)
    cyber_compressed_image.format = ros_compressed_image.format
    cyber_compressed_image.data = ros_compressed_image.data
    return cyber_compressed_image

bag_convert/bag2record/sensor_image.py
daohu527/bag_convert

# example

'''
def convert_func(msg):
  pass

register.add("custom_topic", convert_func)
'''

bag_convert/configs/custom_convert.py
daohu527/bag_convert
# -*- coding: utf-8 -*-
#!/usr/bin/env python

# Copyright 2021 daohu527 <daohu527@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import logging
import sys
from pathlib import Path

import bag_convert.bag2record.bag2record as bag2record
import bag_convert.record2bag.record2bag as record2bag


def main(args=sys.argv):
    parser = argparse.ArgumentParser(
        description="bag_convert is a tool to convert rosbag to record or reverse.",
        prog="main.py")

    parser.add_argument(
        "-m", "--mode", action="store", type=str, required=False,
        help="conversion method")
    parser.add_argument(
        "-b", "--bag_file", action="store", type=str, required=False,
        default="result.bag", help="Ros bag file path")
    parser.add_argument(
        "-r", "--record_file", action="store", type=str, required=False,
        default="result.record", help="Cyber record file path")

    args = parser.parse_args(args[1:])

    # 1. Check if parameters are valid
    bag_file = Path(args.bag_file)
    record_file = Path(args.record_file)

    # 2. conversion package
    if args.mode == 'b2r':
        if not bag_file or not bag_file.is_file():
            logging.error("Bag file not exist! '{}'".format(args.bag_file))
            return
        bag2record.convert(bag_file, record_file)
    elif args.mode == 'r2b':
        if not record_file or not record_file.is_file():
            logging.error(
                "Record file not exist! '{}'".format(args.record_file))
            return
        record2bag.convert(record_file, bag_file)
    else:
        logging.error("Mode not exist! '{}'".format(args.mode))

bag_convert/main.py
daohu527/bag_convert
#!/usr/bin/env python

# Copyright 2021 daohu527 <daohu527@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from sensor_msgs.msg import CompressedImage

from bag_convert.record2bag.header import add_header


@add_header
def to_compressed_image(cyber_compressed_image):
    ros_compressed_image = CompressedImage()
    ros_compressed_image.format = cyber_compressed_image.format
    ros_compressed_image.data = cyber_compressed_image.data
    return ros_compressed_image

bag_convert/record2bag/compressed_image.py
daohu527/bag_convert
#!/usr/bin/env python

# Copyright 2021 daohu527 <daohu527@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import rospy
from std_msgs.msg import Header


def to_header(ros_header, cyber_header):
    ros_header.seq = cyber_header.sequence_num
    ros_header.stamp = rospy.Time.from_sec(cyber_header.timestamp_sec)
    ros_header.frame_id = cyber_header.frame_id


def add_header(func):
    def inner(cyber_msg):
        ros_msg = func(cyber_msg)
        to_header(ros_msg.header, cyber_msg.header)
        return ros_msg
    return inner

bag_convert/record2bag/header.py
daohu527/bag_convert
#!/usr/bin/env python

# Copyright 2021 daohu527 <daohu527@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from sensor_msgs.msg import Image

from bag_convert.record2bag.header import add_header


@add_header
def to_image(cyber_image):
    ros_image = Image()

    ros_image.height = cyber_image.height
    ros_image.width = cyber_image.width
    ros_image.encoding = cyber_image.encoding
    ros_image.step = cyber_image.step

    ros_image.data = cyber_image.data
    return ros_image

bag_convert/record2bag/image.py
daohu527/bag_convert
#!/usr/bin/env python

# Copyright 2021 daohu527 <daohu527@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from sensor_msgs.msg import Imu

from bag_convert.record2bag.header import add_header


@add_header
def to_imu(cyber_imu):
    ros_imu = Imu()

    pose = cyber_imu.imu
    ros_imu.orientation.x = pose.orientation.qx
    ros_imu.orientation.y = pose.orientation.qy
    ros_imu.orientation.z = pose.orientation.qz
    ros_imu.orientation.w = pose.orientation.qw

    ros_imu.angular_velocity.x = pose.angular_velocity.x
    ros_imu.angular_velocity.y = pose.angular_velocity.y
    ros_imu.angular_velocity.z = pose.angular_velocity.z

    ros_imu.linear_acceleration.x = pose.linear_acceleration.x
    ros_imu.linear_acceleration.y = pose.linear_acceleration.y
    ros_imu.linear_acceleration.z = pose.linear_acceleration.z
    return ros_imu

bag_convert/record2bag/imu.py
daohu527/bag_convert
#!/usr/bin/env python

# Copyright 2021 daohu527 <daohu527@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from std_msgs.msg import Header
from sensor_msgs.msg import PointField
from sensor_msgs import point_cloud2

from bag_convert.record2bag.header import to_header


def to_pointcloud(cyber_pointcloud):
    ros_header = Header()
    to_header(ros_header, cyber_pointcloud.header)

    # todo(zero): Need to parse "ros_pointcloud2.fields" for assignment
    points = []
    for point_xyzit in cyber_pointcloud.point:
        points.append([point_xyzit.x, point_xyzit.y, point_xyzit.z,
                       point_xyzit.intensity, point_xyzit.timestamp])

    fields = [PointField('x', 0, PointField.FLOAT32, 1),
              PointField('y', 4, PointField.FLOAT32, 1),
              PointField('z', 8, PointField.FLOAT32, 1),
              PointField('intensity', 12, PointField.UINT32, 1),
              PointField('timestamp', 16, PointField.FLOAT64, 1)]

    return point_cloud2.create_cloud(ros_header, fields, points)

bag_convert/record2bag/pointcloud2.py
daohu527/bag_convert
#!/usr/bin/env python

# Copyright 2021 daohu527 <daohu527@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from geometry_msgs.msg import Pose

from bag_convert.record2bag.header import add_header


def to_pose(cyber_localization):
    ros_pose = Pose()

    pose = cyber_localization.pose
    ros_pose.position.x = pose.position.x
    ros_pose.position.y = pose.position.y
    ros_pose.position.z = pose.position.z

    ros_pose.orientation.x = pose.orientation.qx
    ros_pose.orientation.y = pose.orientation.qy
    ros_pose.orientation.z = pose.orientation.qz
    ros_pose.orientation.w = pose.orientation.qw
    return ros_pose

bag_convert/record2bag/pose.py
daohu527/bag_convert
#!/usr/bin/env python

# Copyright 2021 daohu527 <daohu527@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging

import rosbag
import rospy
from cyber_record.record import Record
from tqdm import tqdm

from bag_convert.record2bag.compressed_image import to_compressed_image
from bag_convert.record2bag.image import to_image
from bag_convert.record2bag.imu import to_imu
from bag_convert.record2bag.pointcloud2 import to_pointcloud
from bag_convert.record2bag.pose import to_pose


def convert_msg(cyber_msg):
    msg_type = cyber_msg.__class__.__name__
    if "CorrectedImu" in msg_type:
        return to_imu(cyber_msg)
    elif "LocalizationEstimate" in msg_type:
        return to_pose(cyber_msg)
    elif "PointCloud" in msg_type:
        return to_pointcloud(cyber_msg)
    elif "Image" in msg_type:
        return to_image(cyber_msg)
    elif "CompressedImage" in msg_type:
        return to_compressed_image(cyber_msg)
    else:
        return None


def convert(record_file, bag_file="result.bag"):
    record = Record(record_file)
    pbar = tqdm(desc="bag to record progress")
    with rosbag.Bag(bag_file, 'w') as bag:
        for topic, msg, t in record.read_messages():
            logging.debug("{},{},{}".format(topic, type(msg), t))
            msg = convert_msg(msg)
            t = rospy.Time.from_sec(t / (10**9))
            if msg:
                bag.write(topic, msg, t)
            pbar.update()
    record.close()
    pbar.close()


if __name__ == "__main__":
    convert("../../data/demo_sensor_data_for_vision.record")

bag_convert/record2bag/record2bag.py
daohu527/bag_convert
import setuptools

with open("README.md", "r", encoding="utf-8") as fh:
    long_description = fh.read()

setuptools.setup(
    name="bag_convert",
    version="0.0.2",
    author="daohu527",
    author_email="daohu527@gmail.com",
    description="Ros bag to Apollo record or reverse conversion",
    long_description=long_description,
    long_description_content_type="text/markdown",
    url="https://github.com/daohu527/bag_convert",
    project_urls={
        "Bug Tracker": "https://github.com/daohu527/bag_convert/issues",
    },
    classifiers=[
        "Programming Language :: Python :: 3",
        "License :: OSI Approved :: Apache Software License",
        "Operating System :: OS Independent",
    ],
    package_dir={"": "."},
    packages=setuptools.find_packages(where="."),
    install_requires=[
        'cyber_record',
        'record_msg<=0.1.1',
        'py3rosmsgs',
        'rospkg',
        'pycryptodomex',
        'tqdm',
    ],
    entry_points={
        'console_scripts': [
            'bag_convert = bag_convert.main:main',
        ],
    },
    python_requires=">=3.6",
)

setup.py
daohu527/bag_convert