ros2_driver.py

import cv2
import numpy as np
import open3d as o3d
import time
import matplotlib.pyplot as plt
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import PointCloud2, PointField, Image, CameraInfo
from std_msgs.msg import Header
import sensor_msgs_py.point_cloud2 as pc2
from cv_bridge import CvBridge

# ==== CONFIGURATION ====
LEFT_ID = 2
RIGHT_ID = 0
width, height = 640, 480

# ROS2 Setup
class PointCloudPublisher(Node):
    def __init__(self):
        super().__init__('stereo_pointcloud_publisher')
        self.publisher_ = self.create_publisher(PointCloud2, '/stereo/points', 10)
        self.left_pub = self.create_publisher(Image, '/stereo/left/image_rect_color', 10)
        self.right_pub = self.create_publisher(Image, '/stereo/right/image_rect_color', 10)
        self.left_info_pub = self.create_publisher(CameraInfo, '/stereo/left/camera_info', 10)
        self.right_info_pub = self.create_publisher(CameraInfo, '/stereo/right/camera_info', 10)
        self.bridge = CvBridge()
        print("[*] ROS2 Node started: publishing on /stereo/points and image topics")

    def publish_pointcloud(self, points, colors):
        header = Header()
        header.stamp = self.get_clock().now().to_msg()
        header.frame_id = "camera_link"

        # Create packed RGB float32 field
        cloud_points = []
        for pt, clr in zip(points, colors):
            r, g, b = (clr * 255).astype(np.uint8)
            rgb = int(r) << 16 | int(g) << 8 | int(b)
            rgb_float = np.frombuffer(np.array([rgb], dtype=np.uint32).tobytes(), dtype=np.float32)[0]
            cloud_points.append((pt[0], pt[1], pt[2], rgb_float))

        fields = [
            PointField(name='x', offset=0, datatype=PointField.FLOAT32, count=1),
            PointField(name='y', offset=4, datatype=PointField.FLOAT32, count=1),
            PointField(name='z', offset=8, datatype=PointField.FLOAT32, count=1),
            PointField(name='rgb', offset=12, datatype=PointField.FLOAT32, count=1),
        ]

        cloud_msg = pc2.create_cloud(header, fields, cloud_points)
        self.publisher_.publish(cloud_msg)

    def publish_images(self, left, right):
        stamp = self.get_clock().now().to_msg()
        left_msg = self.bridge.cv2_to_imgmsg(left, encoding="bgr8")
        right_msg = self.bridge.cv2_to_imgmsg(right, encoding="bgr8")
        left_msg.header.stamp = stamp
        right_msg.header.stamp = stamp
        left_msg.header.frame_id = "camera_left_optical_frame"
        right_msg.header.frame_id = "camera_right_optical_frame"
        self.left_pub.publish(left_msg)
        self.right_pub.publish(right_msg)

        # Dummy camera info
        cam_info = CameraInfo()
        cam_info.width = width
        cam_info.height = height
        cam_info.header.stamp = stamp
        cam_info.header.frame_id = "camera_left_optical_frame"
        self.left_info_pub.publish(cam_info)

        cam_info.header.frame_id = "camera_right_optical_frame"
        self.right_info_pub.publish(cam_info)

# Load calibration
data = np.load("calibration_data.npz")
left_map1, left_map2 = data["left_map1"], data["left_map2"]
right_map1, right_map2 = data["right_map1"], data["right_map2"]
Q = data["Q"]

# Stereo matching config
stereo = cv2.StereoSGBM_create(
    minDisparity=0,
    numDisparities=16 * 7,
    blockSize=3,
    P1=8 * 3 * 3**2,
    P2=32 * 3 * 3**2,
    disp12MaxDiff=1,
    uniquenessRatio=20,
    speckleWindowSize=0,
    speckleRange=64
)

# === CAMERA SETUP ===
camL = cv2.VideoCapture(LEFT_ID)
camR = cv2.VideoCapture(RIGHT_ID)
camL.set(3, width)
camL.set(4, height)
camR.set(3, width)
camR.set(4, height)

# === POINT CLOUD VIS ===
vis = o3d.visualization.Visualizer()
vis.create_window("Live Point Cloud", width=960, height=540)
pcd = o3d.geometry.PointCloud()
vis.add_geometry(pcd)

render_option = vis.get_render_option()
render_option.background_color = np.asarray([0, 0, 0])
render_option.point_size = 3.0
render_option.show_coordinate_frame = True
render_option.point_show_normal = False

coord_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.1, origin=[0, 0, 0])
vis.add_geometry(coord_frame)
first_frame = True

print("[*] Running live stereo → point cloud. Press Ctrl+C or close window to exit.")

rclpy.init()
node = PointCloudPublisher()

try:
    while rclpy.ok():
        start = time.time()
        retL, frameL = camL.read()
        retR, frameR = camR.read()
        if not retL or not retR:
            print("[!] Camera frame failed")
            break

        rectL = cv2.remap(frameL, left_map1, left_map2, cv2.INTER_LINEAR)
        rectR = cv2.remap(frameR, right_map1, right_map2, cv2.INTER_LINEAR)

        grayL = cv2.cvtColor(rectL, cv2.COLOR_BGR2GRAY)
        grayR = cv2.cvtColor(rectR, cv2.COLOR_BGR2GRAY)

        disp = stereo.compute(grayL, grayR).astype(np.float32) / 16.0

        disp_vis = cv2.normalize(disp, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
        cv2.imshow("Disparity", disp_vis)
        cv2.imshow("Left Rectified", rectL)
        cv2.imshow("Right Rectified", rectR)
        cv2.waitKey(1)

        points_3D = cv2.reprojectImageTo3D(disp, Q)
        max_disp = stereo.getNumDisparities()
        mask = (disp > 0.1) & (disp < max_disp) & np.isfinite(points_3D[:, :, 2])
        points = points_3D[mask]
        rectL_rgb = cv2.cvtColor(rectL, cv2.COLOR_BGR2RGB)
        colors = rectL_rgb[mask]

        z = points[:, 2]
        valid_z = np.logical_and(z > 0.1, z < 100.0)
        points = points[valid_z]
        colors = colors[valid_z]

        if points.shape[0] == 0:
            print("[!] No valid 3D points")
            continue

        points[:, 1] *= -1
        points[:, 2] *= -1

        sample_idx = np.random.choice(len(points), size=min(30000, len(points)), replace=False)
        sampled_pts = points[sample_idx]
        sampled_clr = colors[sample_idx] / 255.0

        pcd.points = o3d.utility.Vector3dVector(sampled_pts)
        pcd.colors = o3d.utility.Vector3dVector(sampled_clr)

        vis.update_geometry(pcd)
        if first_frame:
            bbox = pcd.get_axis_aligned_bounding_box()
            vis.reset_view_point(True)
            vis.get_view_control().set_lookat(bbox.get_center())
            vis.get_view_control().set_zoom(0.7)
            first_frame = False

        vis.poll_events()
        vis.update_renderer()

        elapsed = time.time() - start
        fps = 1.0 / elapsed
        mean_depth = np.mean(points[:, 2])
        std_depth = np.std(points[:, 2])
        print(f"[FPS: {fps:.1f}] Points: {len(points):,}  Depth Mean: {mean_depth:.2f}m  Std: {std_depth:.2f}m")

        node.publish_pointcloud(sampled_pts, sampled_clr)
        node.publish_images(rectL, rectR)
        rclpy.spin_once(node, timeout_sec=0)

except KeyboardInterrupt:
    print("\n[✓] Exiting...")

camL.release()
camR.release()
vis.destroy_window()
cv2.destroyAllWindows()
node.destroy_node()
rclpy.shutdown()