python 旋转平移rosbag中点云数据

#!/usr/bin/env python3

import rospy
import rosbag
import sensor_msgs.point_cloud2 as pc2
from sensor_msgs.msg import PointCloud2, PointField
import numpy as np
import tf.transformations as tf
import sys

def transform_pointcloud(input_bagfile, output_bagfile, rotation_angle, translation):
    rotation_matrix = tf.euler_matrix(0, 0, rotation_angle)  # 创建旋转矩阵
    translation_vector = np.array(translation)

    with rosbag.Bag(output_bagfile, 'w') as outbag:
        for topic, msg, t in rosbag.Bag(input_bagfile).read_messages():
            if topic == '/livox/lidar':
                points = pc2.read_points(msg, field_names=("x", "y", "z","intensity"), skip_nans=True)
                points_transformed = []
                for pt in points:
                    pt_xyz = np.array([pt[0], pt[1], pt[2]])
                    pt_rotated = np.dot(rotation_matrix[:3, :3], pt_xyz)  # 应用旋转矩阵
                    pt_transformed = translation_vector + pt_rotated # 应用平移向量
                    pt_transformed = np.append(pt_transformed, pt[3]) #添加强度信息
                    points_transformed.append(pt_transformed)

                # 创建一个新的PointCloud2消息并写入输出bag文件
                msg_transformed = PointCloud2()
                msg_transformed.header = msg.header
                msg_transformed.height = 1
                msg_transformed.width = len(points_transformed)
                msg_transformed.fields = [
                    PointField(name="x", offset=0, datatype=PointField.FLOAT32, count=1),
                    PointField(name="y", offset=4, datatype=PointField.FLOAT32, count=1),
                    PointField(name="z", offset=8, datatype=PointField.FLOAT32, count=1),
                    PointField(name="intensity", offset=12, datatype=PointField.FLOAT32, count=1),
                ]
                msg_transformed.is_bigendian = False
                msg_transformed.point_step = 16
                msg_transformed.row_step = msg_transformed.point_step * len(points_transformed)
                msg_transformed.is_dense = True
                msg_transformed.data = np.asarray(points_transformed, np.float32).tobytes()

                outbag.write(topic, msg_transformed, t)
            else:
                outbag.write(topic, msg, t)

if __name__ == '__main__':
    rospy.init_node('transform_pointcloud_node')
    if len(sys.argv) < 2:
        print("You should input one bag filepath at least.")
        exit()
    input_bag = sys.argv[1]
    name = sys.argv[2]
    output_bag = input_bag[:input_bag.rfind(".")] + "_" + name + ".bag"
    rotation_angle = rospy.get_param('~rotation_angle', np.pi / 2.0)  # 旋转角度，默认为90度
    translation = rospy.get_param('~translation', [0.0, 0.0, 0.0])  # 平移向量，默认为[1.0, 0.0, 0.0]

    try:
        transform_pointcloud(input_bag, output_bag, rotation_angle, translation)
        rospy.loginfo("Point cloud transformed and saved successfully!")
    except Exception as e:
        rospy.logerr("Failed to transform point cloud: %s", str(e))