save_pcd_odom.py

save_pcd_odom.py

# !/usr/bin/env python3
# coding=utf-8

from __future__ import print_function, division, absolute_import
import rospy
from sensor_msgs.msg import PointCloud2
from nav_msgs.msg import Odometry
import message_filters

import numpy as np
import ros_numpy

import pypcd
import transforms3d

import os
import sys

frame_id = 0


def get_RT_matrix(odom_msg):
    quaternion = np.asarray(
        [odom_msg.pose.pose.orientation.w, odom_msg.pose.pose.orientation.x, odom_msg.pose.pose.orientation.y,
         odom_msg.pose.pose.orientation.z])
    translation = np.asarray(
        [odom_msg.pose.pose.position.x, odom_msg.pose.pose.position.y, odom_msg.pose.pose.position.z])

    rotation = transforms3d.quaternions.quat2mat(quaternion)

    # print(T_qua2rota)
    RT_matrix = np.eye(4)
    RT_matrix[:3, :3] = rotation
    RT_matrix[:3, 3] = translation.T

    # RT_matrix = np.matmul(R_axes, RT_matrix)

    return RT_matrix


def callback(odom_msg, pc_msg):
    global frame_id

    stamp = odom_msg.header.stamp
    print("{:d}.{:d}".format(stamp.secs, stamp.nsecs))

    pc_array = ros_numpy.numpify(pc_msg)
    if len(pc_array.shape) == 2:
        pc = np.zeros((pc_array.shape[0] * pc_array.shape[1], 4))
    else:
        pc = np.zeros((pc_array.shape[0], 4))
    # 解析点格式
    pc[:, 0] = pc_array['x'].reshape(-1)
    pc[:, 1] = pc_array['y'].reshape(-1)
    pc[:, 2] = pc_array['z'].reshape(-1)
    pc[:, 3] = pc_array['intensity'].reshape(-1)

    save_path = "/tmp/odometry"
    if not os.path.exists(save_path):
        os.mkdir(save_path)
    pcd = pypcd.PointCloud.from_msg(pc_msg)

    odom = get_RT_matrix(odom_msg)
    odom_file = os.path.join(save_path, "{:d}.{:d}.odom".format(stamp.secs, stamp.nsecs))
    np.savetxt(odom_file, odom, fmt="%.5f")

    pc_file = os.path.join(save_path, "{:d}.{:d}.pcd".format(stamp.secs, stamp.nsecs))
    pcd.save_pcd(pc_file, compression="binary_compressed")
    print("saving {:03}....".format(frame_id))
    frame_id += 1


if __name__ == '__main__':
    rospy.init_node("save_pcd_odom", anonymous=True)
    rospy.loginfo("Start....")

    odom_sub = message_filters.Subscriber("/lio_sam/mapping/odometry", Odometry)
    points_sub = message_filters.Subscriber("/RS32/velodyne_points", PointCloud2)

    ts = message_filters.ApproximateTimeSynchronizer([odom_sub, points_sub], 20, 0.1, allow_headerless=True)

    ts.registerCallback(callback)

    rospy.spin()