OSC_ROS2/test/joint_states_recording.py
Alexander Schaefer 33f21d3096 remove gitignore
2025-04-22 17:10:24 +02:00

79 lines
2.7 KiB
Python

import rclpy
from rclpy.node import Node
from sensor_msgs.msg import JointState
import csv
import os
import roboticstoolbox as rtb
class JointStateLogger(Node):
def __init__(self):
super().__init__('joint_state_logger_split')
# Subscribe to /joint_states topic
self.subscription = self.create_subscription(
JointState,
'/joint_states',
self.joint_states_callback,
1000
)
self.robot = rtb.ERobot.URDF('/BA/ur10e.urdf')
# Directory to store logs
self.log_dir = '/BA/joint_states_logs'
os.makedirs(self.log_dir, exist_ok=True)
self.log_dir_cart = '/BA/joint_states_cart_logs'
os.makedirs(self.log_dir_cart, exist_ok=True)
def joint_states_callback(self, msg: JointState):
# Use ROS time as timestamp
stamp = msg.header.stamp
timestamp = f"{stamp.sec}.{str(stamp.nanosec).zfill(9)}"
for i, name in enumerate(msg.name):
# File per joint
file_path = os.path.join(self.log_dir, f"{name}.csv")
# If file doesn't exist, create with header
if not os.path.exists(file_path):
with open(file_path, mode='w', newline='') as f:
writer = csv.writer(f)
writer.writerow(["timestamp", "position", "velocity", "effort"])
# Get values if they exist
position = msg.position[i] if i < len(msg.position) else ''
velocity = msg.velocity[i] if i < len(msg.velocity) else ''
effort = msg.effort[i] if i < len(msg.effort) else ''
# Append to joint-specific CSV
with open(file_path, mode='a', newline='') as f:
writer = csv.writer(f)
writer.writerow([timestamp, position, velocity, effort])
file_path_cart = os.path.join(self.log_dir_cart, "cartesian.csv")
if not os.path.exists(file_path_cart):
with open(file_path_cart, mode='w', newline='') as f:
writer = csv.writer(f)
writer.writerow(["timestamp", "x", "y", "z", "roll", "pitch", "yaw"])
with open(file_path_cart, mode='a', newline='') as f:
writer = csv.writer(f)
[x,y,z] = self.robot.fkine(list([msg.position[-1]])+list(msg.position[:4])).t
[roll, pitch, yaw] = self.robot.fkine(list([msg.position[-1]])+list(msg.position[:4])).rpy()
writer.writerow([timestamp, x, y, z, roll, pitch, yaw])
def main():
rclpy.init()
node = JointStateLogger()
try:
rclpy.spin(node)
except KeyboardInterrupt:
print("")
finally:
node.destroy_node()
rclpy.shutdown()
if __name__ == '__main__':
main()