code to convert both depth and rgb data to .jpeg in sync

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()
+    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()