Merge pull request #18 from Tang-JingWei/ROS2

Add ros2 publish uw image function.

Author: haoyuma2002814

docs/CHANGELOG.md

@@ -5,3 +5,10 @@
 ### Added
 
 - Initial version of OceanSim Extension
+
+## [0.1.0] - 2025-08-05
+
+### Added
+
+- Add ros2 control function
+- Add ros2 publish uw image

docs/README.md

@@ -11,10 +11,25 @@ We suggest that make sure the extension isaacsim.ros2.bridge is being setup to "
 
 
 ## Usage:
+### ros2 control:
 We provided an exmaple util located at `isaacsim/oceansim/utils/ros2_control.py` for user to consult and develop on.
 
 This util extends the control mode to ros control in the **sensor_example** extension. 
 
+### ros2 publish uw image:
+We add ros2 publish uw image function in the UW_Camera class, located at `isaacsim/oceansim/sensors/UW_Camera.py`.
+
+For testing, we provide a ros2 subscriber example located at `isaacsim/oceansim/utils/ros2_image_subscriber.py`.
+
+Test steps:
+1. check the Underwater Camera checkbox in the **sensor_example** extension.
+2. run the simulation.
+3. open a terminal and run the ros2_image_subscriber.py.
+```
+cd /path/to/oceansim/utils
+python3 ros2_image_subscriber.py
+```
+
 ## Acknowledgement:
 Great appreciation to [Tang-JingWei](https://github.com/Tang-JingWei) for contributng the ros bridge example for OceanSim.
 

isaacsim/oceansim/sensors/UW_Camera.py

@@ -13,6 +13,18 @@
 # Custom import
 from isaacsim.oceansim.utils.UWrenderer_utils import UW_render
 
+'''
+Attention:
+
+Before OceanSim extension being activated, the extension isaacsim.ros2.bridge should be activated, otherwise rclpy will
+fail to be loaded.
+
+so, we suggest that make sure the extension isaacsim.ros2.bridge is being setup to "AUTOLOADED" in Window->Extension.
+'''
+import rclpy
+from sensor_msgs.msg import CompressedImage
+import time
+import cv2
 
 class UW_Camera(Camera):
 
@@ -64,7 +76,8 @@ def initialize(self,
                    viewport: bool = True,
                    writing_dir: str = None,
                    UW_yaml_path: str = None,
-                   physics_sim_view=None):
+                   physics_sim_view=None,
+                   enable_ros2_pub=True, uw_img_topic="/oceansim/robot/uw_img", ros2_pub_frequency=5, ros2_pub_jpeg_quality=50):
 
         """Configure underwater rendering properties and initialize pipelines.
     
@@ -77,7 +90,11 @@ def initialize(self,
             viewport (bool, optional): Enable viewport visualization. Defaults to True.
             writing_dir (str, optional): Directory to save rendered images. Defaults to None.
             UW_yaml_path (str, optional): Path to YAML file with water properties. Defaults to None.
-            physics_sim_view (_type_, optional): _description_. Defaults to None.            
+            physics_sim_view (_type_, optional): _description_. Defaults to None.          
+            enable_ros2_pub (bool, optional): Enable ROS2 communication. Defaults to True.
+            uw_img_topic (str, optional): ROS2 topic name for UW image. Defaults to "/oceansim/robot/uw_img".
+            ros2_pub_frequency (int, optional): ROS2 publish frequency. Defaults to 5.
+            ros2_pub_jpeg_quality (int, optional): ROS2 publish jpeg quality. Defaults to 50.
     
         """
         self._id = 0
@@ -116,9 +133,90 @@ def initialize(self,
         if writing_dir is not None:
             self._writing = True
             self._writing_backend = rep.BackendDispatch({"paths": {"out_dir": writing_dir}})
+
+        # ROS2 configuration
+        self._enable_ros2_pub = enable_ros2_pub
+        self._uw_img_topic = uw_img_topic
+        self._last_publish_time = 0.0
+        self._ros2_pub_frequency = ros2_pub_frequency     # publish frequency, hz
+        self._ros2_pub_jpeg_quality = ros2_pub_jpeg_quality
+        self._setup_ros2_publisher()
 
         print(f'[{self._name}] Initialized successfully. Data writing: {self._writing}')
 
+    def _setup_ros2_publisher(self):
+        '''
+        setup the publisher for uw image
+        '''
+        try:
+            if not self._enable_ros2_pub:
+                return
+            
+            # Initialize ROS2 context if not already done
+            if not rclpy.ok():
+                rclpy.init()
+                print(f'[{self._name}] ROS2 context initialized')
+
+            # Create uw image publisher node
+            node_name = f'oceansim_rob_uw_img_pub_{self._name.lower()}'.replace(' ', '_')
+            self._ros2_uw_img_node = rclpy.create_node(node_name)
+            self._uw_img_pub = self._ros2_uw_img_node.create_publisher(
+                CompressedImage, 
+                self._uw_img_topic, 
+                10
+            )
+        
+        except Exception as e:
+            print(f'[{self._name}] ROS2 uw image publisher setup failed: {e}')
+
+    def _ros2_publish_uw_img(self, uw_img):
+        """
+        publish the uw image
+        """
+        try:
+            if self._uw_img_pub is None:
+                return
+
+            # fps control
+            current_time = time.time()
+            if current_time - self._last_publish_time < (1.0 / self._frequency):
+                return
+
+            # Convert the image
+            uw_image_cpu = uw_img.numpy()
+            if uw_image_cpu.dtype != np.uint8:
+                uw_image_cpu = uw_image_cpu.astype(np.uint8)    # UW_render return 'rgba'
+            uw_image_bgr = cv2.cvtColor(uw_image_cpu, cv2.COLOR_RGBA2BGR)
+            encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), self._ros2_pub_jpeg_quality]      # JPEG quality, default 90, can be 0-100
+            result, compressed_img = cv2.imencode('.jpg', uw_image_bgr, encode_param)
+            if not result:
+                print(f'[{self._name}] Failed to compress image to JPEG')
+                return
+
+            # Create a ROS2 Image message
+            msg = CompressedImage()
+            msg.header.stamp = self._ros2_uw_img_node.get_clock().now().to_msg()
+            msg.header.frame_id = 'uw_image'
+            msg.format = 'jpeg'
+            msg.data = compressed_img.tobytes()
+            
+            # Publish the message
+            self._uw_img_pub.publish(msg)
+
+            rclpy.spin_once(self._ros2_uw_img_node, timeout_sec=0.0)
+
+            self._last_publish_time = current_time
+
+            # debug
+            # self._ros2_uw_img_node.get_logger().info(
+            #     f'Published image: encoding={msg.encoding}, '
+            #     f'width={msg.width}, height={msg.height}, step={msg.step}, '
+            #     f'data_size={len(msg.data)}'
+            # )
+
+        except Exception as e:
+            print(f'[{self._name}] ROS2 uw image publish failed: {e}')
+
     def render(self):
         """Process and display a single frame with underwater effects.
     
@@ -150,6 +248,9 @@ def render(self):
             if self._writing:
                 self._writing_backend.schedule(write_image, path=f'UW_image_{self._id}.png', data=uw_image)
                 print(f'[{self._name}] [{self._id}] Rendered image saved to {self._writing_backend.output_dir}')
+            if self._enable_ros2_pub:
+                self._ros2_publish_uw_img(uw_image)
+                pass
 
             self._id += 1
 

isaacsim/oceansim/utils/asset_path.json

@@ -1,3 +1,3 @@
 {
-  "asset_path": "/home/jingwei/jingwei_ssd/OceanSim_assets"
+  "asset_path": "/home/jingwei/OceanSim_assets"
 }

isaacsim/oceansim/utils/ros2_image_subscriber.py

@@ -0,0 +1,65 @@
+import rclpy
+from rclpy.node import Node
+from sensor_msgs.msg import CompressedImage
+import cv2
+import numpy as np
+from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy
+
+class ImageSubscriber(Node):
+    def __init__(self):
+        super().__init__('image_subscriber')
+        
+        qos = QoSProfile(
+            reliability=ReliabilityPolicy.RELIABLE,
+            history=HistoryPolicy.KEEP_LAST,
+            depth=1
+        )
+
+        self.subscription = self.create_subscription(
+            CompressedImage,
+            '/oceansim/robot/uw_img',
+            self.image_callback,
+            qos)
+            
+        self.get_logger().info('Compressed Image Subscriber Node has been started')
+
+    def image_callback(self, msg):
+        try:
+            # debug
+            self.get_logger().info(f'Received compressed image: format={msg.format}, '
+                                 f'data_size={len(msg.data)}')
+
+            # convert
+            np_arr = np.frombuffer(msg.data, np.uint8)
+            # decode jpeg img
+            cv_image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)  # return BGR format
+
+            if cv_image is None:
+                self.get_logger().error('Failed to decode compressed image')
+                return
+
+            # debug
+            self.get_logger().info(f'Image shape: {cv_image.shape}, dtype: {cv_image.dtype}, '
+                                 f'min: {cv_image.min()}, max: {cv_image.max()}')
+
+            # imshow
+            cv2.imshow('Underwater Image', cv_image)
+            cv2.waitKey(1)
+
+        except Exception as e:
+            self.get_logger().error(f'Error processing compressed image: {str(e)}')
+
+def main(args=None):
+    rclpy.init(args=args)
+    image_subscriber = ImageSubscriber()
+    try:
+        rclpy.spin(image_subscriber)
+    except KeyboardInterrupt:
+        pass
+    finally:
+        cv2.destroyAllWindows()
+        image_subscriber.destroy_node()
+        rclpy.shutdown()
+
+if __name__ == '__main__':
+    main()