Untitled

#!/usr/bin/env python

# picamera stereo ROS node for Raspberry Pi 5 with dual cameras

from picamera2 import Picamera2, Preview
from picamera2.encoders import JpegEncoder
import time
import rospy
from sensor_msgs.msg import CameraInfo, Image
import yaml
import io
import signal
import sys


def parse_calibration_yaml(calib_file):
    with open(calib_file, 'r') as f:
        params = yaml.safe_load(f)

    cam_info = CameraInfo()
    cam_info.height = params['image_height']
    cam_info.width = params['image_width']
    cam_info.distortion_model = params['distortion_model']
    cam_info.K = params['camera_matrix']['data']
    cam_info.D = params['distortion_coefficients']['data']
    cam_info.R = params['rectification_matrix']['data']
    cam_info.P = params['projection_matrix']['data']

    return cam_info


# cam resolution
res_x = 320
res_y = 240
target_FPS = 15

# initialize the cameras
print("Init cameras...")
camera1 = Picamera2(0)  # First camera
camera2 = Picamera2(1)  # Second camera

camera1.configure(camera1.create_still_configuration(main={"size": (res_x, res_y)}, buffer_count=2))
camera2.configure(camera2.create_still_configuration(main={"size": (res_x, res_y)}, buffer_count=2))

camera1.start()
camera2.start()

# setup the publishers
print("init publishers")
left_img_pub = rospy.Publisher('stereo/left/image_raw', Image, queue_size=1)
right_img_pub = rospy.Publisher('stereo/right/image_raw', Image, queue_size=1)

left_cam_pub = rospy.Publisher('stereo/left/camera_info', CameraInfo, queue_size=1)
right_cam_pub = rospy.Publisher('stereo/right/camera_info', CameraInfo, queue_size=1)

rospy.init_node('stereo_pub')

# init messages
left_img_msg = Image()
left_img_msg.height = res_y
left_img_msg.width = res_x
left_img_msg.step = res_x * 3
left_img_msg.encoding = 'rgb8'
left_img_msg.header.frame_id = 'stereo_camera'

right_img_msg = Image()
right_img_msg.height = res_y
right_img_msg.width = res_x
right_img_msg.step = res_x * 3
right_img_msg.encoding = 'rgb8'
right_img_msg.header.frame_id = 'stereo_camera'

# parse the left and right camera calibration yaml files
left_cam_info = parse_calibration_yaml('/home/ubuntu/left.yaml')
right_cam_info = parse_calibration_yaml('/home/ubuntu/right.yaml')

def signal_handler(signal, frame):
    print('CTRL-C caught')
    print('closing cameras')
    camera1.stop()
    camera2.stop()
    time.sleep(1)
    print('cameras closed')
    sys.exit(0)

signal.signal(signal.SIGINT, signal_handler)

print("Setup done, entering main loop")
framecount = 0
frametimer = time.time()
toggle = True

while not rospy.is_shutdown():
    framecount += 1

    left_frame = camera1.capture_array()
    right_frame = camera2.capture_array()

    stamp = rospy.Time.now()
    left_img_msg.header.stamp = stamp
    right_img_msg.header.stamp = stamp
    left_cam_info.header.stamp = stamp
    right_cam_info.header.stamp = stamp

    left_cam_pub.publish(left_cam_info)
    right_cam_pub.publish(right_cam_info)

    left_img_msg.data = left_frame.tobytes()
    right_img_msg.data = right_frame.tobytes()

    left_img_pub.publish(left_img_msg)
    right_img_pub.publish(right_img_msg)

    if time.time() > frametimer + 1.0:
        if toggle:
            indicator = '  o'
        else:
            indicator = '  -'
        toggle = not toggle
        print('approx publish rate:', framecount, 'target FPS:', target_FPS, indicator)
        frametimer = time.time()
        framecount = 0