diff --git a/ros2_ws/src/convert_2_img/convert_2_img/convert_to_img.py b/ros2_ws/src/convert_2_img/convert_2_img/convert_to_img.py index a5a78205..36b8272c 100644 --- a/ros2_ws/src/convert_2_img/convert_2_img/convert_to_img.py +++ b/ros2_ws/src/convert_2_img/convert_2_img/convert_to_img.py @@ -1,25 +1,60 @@ -import rclpy -from sensor_msgs.msg import Image +#!/usr/bin/env python +import os import cv2 -import cv_bridge +from cv_bridge import CvBridge +from sensor_msgs.msg import Image +import rclpy +from rclpy.node import Node +import threading -count = 0 -def image_callback(data): - bridge = cv_bridge.CvBridge() - try: - cv_image = bridge.imgmsg_to_cv2(data, "bgr8") - except cv_bridge.CvBridgeError as e: - print(e) - - global count - cv2.imwrite('stereo_image' + str(count)+ '.jpeg', cv_image) - count = count + 1 +class ImageSaver(Node): + def __init__(self): + super().__init__('image_saver') + self.depth_sub = self.create_subscription(Image, '/zed2i/zed_node/depth/depth_registered', self.depth_callback, 10) + self.image_sub = self.create_subscription(Image, '/zed2i/zed_node/left_raw/image_raw_color', self.image_callback, 10) + self.cv_bridge = CvBridge() + self.depth_image = None + self.color_image = None + self.count_stereo = 0 + self.count_depth = 0 -def main(args=None): - rclpy.init(args=args) - node = rclpy.create_node('stereo_image_subscriber') - sub = node.create_subscription(Image, '/camera/color/image_raw', image_callback, False) - rclpy.spin(node) + def depth_callback(self, data): + print("depth cb") + self.depth_image = self.cv_bridge.imgmsg_to_cv2(data) + self.save_images('d') + + + def image_callback(self, data): + print("image cb") + self.color_image = self.cv_bridge.imgmsg_to_cv2(data) + self.save_images('s') + + + def save_images(self, img_type): + print('save images') + + if img_type == 's': + filename_stereo = "/stereo/stereo_{}.jpg".format(self.count_stereo) + # import pdb; pdb.set_trace() + path_stereo = os.getcwd() + filename_stereo + cv2.imwrite(path_stereo, self.color_image) + self.get_logger().info("Saved image to {}".format(path_stereo)) + self.count_stereo += 1 + + if img_type == 'd': + filename_depth = "/depth/depth_{}.jpg".format(self.count_depth) + path_depth = os.getcwd() + filename_depth + cv2.imwrite(path_depth, self.depth_image) + self.get_logger().info("Saved image to {}".format(path_depth)) + self.count_depth += 1 if __name__ == '__main__': - main() \ No newline at end of file + rclpy.init() + image_saver = ImageSaver() + rate = image_saver.create_rate(1) # 10 Hz + rclpy.spin(image_saver) + # while rclpy.ok(): + # image_saver.save_images() + # # rate.sleep() + # image_saver.destroy_node() + # rclpy.shutdown()