Python ros2 screen capture node

region_capture.py

import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import cv2
import mss
import numpy as np

# マウス選択用の状態を保持する辞書
selection_state = {
    'start_point': None,
    'end_point': None,
    'is_drawing': False,
    'roi': None
}

def select_roi_callback(event, x, y, flags, param):
    """OpenCVのマウスイベントコールバック関数"""
    global selection_state
    
    if event == cv2.EVENT_LBUTTONDOWN:
        selection_state['start_point'] = (x, y)
        selection_state['is_drawing'] = True

    elif event == cv2.EVENT_MOUSEMOVE:
        if selection_state['is_drawing']:
            img_copy = param.copy()
            cv2.rectangle(img_copy, selection_state['start_point'], (x, y), (0, 255, 0), 2)
            cv2.imshow("Select Region", img_copy)

    elif event == cv2.EVENT_LBUTTONUP:
        selection_state['end_point'] = (x, y)
        selection_state['is_drawing'] = False
        
        # 矩形の左上と右下を確定
        x1, y1 = selection_state['start_point']
        x2, y2 = selection_state['end_point']
        selection_state['roi'] = (min(x1, x2), min(y1, y2), abs(x1 - x2), abs(y1 - y2))

class ScreenRegionPublisher(Node):
    def __init__(self):
        super().__init__('screen_region_publisher')

        # パラメータの宣言
        self.declare_parameter('publish_fps', 30.0)
        self.declare_parameter('output_width', 640)
        self.declare_parameter('output_height', 480)

        self.fps = self.get_parameter('publish_fps').get_parameter_value().double_value
        self.width = self.get_parameter('output_width').get_parameter_value().integer_value
        self.height = self.get_parameter('output_height').get_parameter_value().integer_value

        self.bridge = CvBridge()
        self.sct = mss.mss()
        self.capture_region = None
        self.publisher_ = None
        self.timer = None

        # 起動時に領域選択UIを呼び出す
        self.capture_region = self.select_capture_region()

        if self.capture_region:
            self.get_logger().info(f"Region selected: {self.capture_region}")
            self.publisher_ = self.create_publisher(Image, 'image_raw', 10)
            timer_period = 1.0 / self.fps
            self.timer = self.create_timer(timer_period, self.publish_region)
            self.get_logger().info(f"Publishing at {self.fps} FPS with output size {self.width}x{self.height}.")
        else:
            self.get_logger().warn("No region selected. Shutting down.")
            # self.destroy_node() を呼ぶとエラーになることがあるので、rclpy.shutdown()で対応
            rclpy.shutdown()


    def select_capture_region(self):
        """UIを表示してユーザーに領域を選択させる関数"""
        global selection_state
        
        # モニター全体のスクリーンショットを撮る
        monitor = self.sct.monitors[0] # 全モニターを含む領域
        sct_img = self.sct.grab(monitor)
        full_screenshot = np.array(sct_img)

        window_name = "Select Region - Drag a box and press ENTER"
        cv2.namedWindow(window_name, cv2.WINDOW_NORMAL) # リサイズ可能なウィンドウ
        cv2.setMouseCallback(window_name, select_roi_callback, full_screenshot)
        cv2.imshow(window_name, full_screenshot)

        self.get_logger().info("Please draw a rectangle on the screen and press 'ENTER' to confirm or 'ESC' to cancel.")

        while True:
            key = cv2.waitKey(1) & 0xFF
            # Enterキーが押されたらループを抜ける
            if key == 13: # 13 is the Enter Key
                break
            # Escキーが押されたらキャンセル
            if key == 27: # 27 is the ESC Key
                selection_state['roi'] = None
                break
        
        cv2.destroyAllWindows()

        if selection_state['roi']:
            x, y, w, h = selection_state['roi']
            # mssはモニター相対座標ではなく、全画面での絶対座標を使う
            return {'left': monitor['left'] + x, 'top': monitor['top'] + y, 'width': w, 'height': h}
        return None

    def publish_region(self):
        """選択された領域をキャプチャし、リサイズして配信する関数"""
        sct_img = self.sct.grab(self.capture_region)
        frame = np.array(sct_img)
        frame_bgr = cv2.cvtColor(frame, cv2.COLOR_BGRA2BGR)
        
        # 指定された解像度にリサイズ
        resized_frame = cv2.resize(frame_bgr, (self.width, self.height))

        # ROSメッセージに変換して配信
        ros_image_msg = self.bridge.cv2_to_imgmsg(resized_frame, "bgr8")
        ros_image_msg.header.stamp = self.get_clock().now().to_msg()
        self.publisher_.publish(ros_image_msg)

def main(args=None):
    rclpy.init(args=args)
    node = ScreenRegionPublisher()
    if rclpy.ok():
        try:
            rclpy.spin(node)
        except KeyboardInterrupt:
            pass
        finally:
            if node and rclpy.ok():
                node.destroy_node()
    # shutdownは複数回呼ばれても問題ない
    if rclpy.ok():
        rclpy.shutdown()

if __name__ == '__main__':
    main()