mi-task/test/demo_test/ai_camera_demo.py

119 lines
4.1 KiB
Python
Raw Normal View History

2025-08-18 11:06:42 +08:00
#!/usr/bin/env python3
2025-08-18 20:34:03 +08:00
#ros2 run camera_test camera_server
2025-08-18 11:06:42 +08:00
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
import numpy as np
from protocol.srv import CameraService
import cv2
2025-08-18 20:34:03 +08:00
import os
import time
from datetime import datetime
# 配置显示环境(关键修改)
os.environ['OPENCV_VIDEOIO_PRIORITY_GTK'] = '1'
os.environ['DISPLAY'] = ':0' # 使用主显示器
2025-08-18 11:06:42 +08:00
class AICameraDemo(Node):
def __init__(self):
super().__init__('ai_camera_demo')
2025-08-18 20:34:03 +08:00
# 创建服务客户端
self.camera_client = self.create_client(CameraService, '/camera_service')
2025-08-18 11:06:42 +08:00
while not self.camera_client.wait_for_service(timeout_sec=1.0):
2025-08-18 20:34:03 +08:00
self.get_logger().info('等待AI相机服务...')
2025-08-18 11:06:42 +08:00
2025-08-18 20:34:03 +08:00
# 图像订阅
2025-08-18 11:06:42 +08:00
self.image_subscription = self.create_subscription(
2025-08-18 20:34:03 +08:00
Image, '/image', self.image_callback, 10
2025-08-18 11:06:42 +08:00
)
2025-08-18 20:34:03 +08:00
# 截图控制参数
self.last_capture_time = time.time()
2025-08-18 11:06:42 +08:00
self.start_time = time.time()
2025-08-18 20:34:03 +08:00
self.capture_interval = 3.0 # 截图间隔(秒)
self.total_duration = 30.0 # 总运行时间(秒)
self.save_dir = "captures" # 截图保存目录
os.makedirs(self.save_dir, exist_ok=True)
# 初始化OpenCV窗口
self.window_name = 'AI Camera Feed'
try:
cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
self.gui_enabled = True
except:
self.get_logger().error("无法创建OpenCV窗口将仅保存截图")
self.gui_enabled = False
2025-08-18 11:06:42 +08:00
def start_camera(self):
2025-08-18 20:34:03 +08:00
"""启动相机服务"""
2025-08-18 11:06:42 +08:00
req = CameraService.Request()
req.command = 9 # START_IMAGE_PUBLISH
2025-08-18 20:34:03 +08:00
req.width, req.height, req.fps = 640, 480, 30
2025-08-18 11:06:42 +08:00
future = self.camera_client.call_async(req)
rclpy.spin_until_future_complete(self, future)
2025-08-18 20:34:03 +08:00
return future.result().result == 0
2025-08-18 11:06:42 +08:00
def stop_camera(self):
2025-08-18 20:34:03 +08:00
"""停止相机服务"""
2025-08-18 11:06:42 +08:00
req = CameraService.Request()
req.command = 10 # STOP_IMAGE_PUBLISH
future = self.camera_client.call_async(req)
rclpy.spin_until_future_complete(self, future)
2025-08-18 20:34:03 +08:00
return future.result().result == 0
2025-08-18 11:06:42 +08:00
def image_callback(self, msg):
2025-08-18 20:34:03 +08:00
"""处理图像并定时截图"""
2025-08-18 11:06:42 +08:00
current_time = time.time()
elapsed = current_time - self.start_time
2025-08-18 20:34:03 +08:00
# 超过总时长则停止
if elapsed >= self.total_duration:
self.get_logger().info("已完成30秒截图任务")
raise KeyboardInterrupt
try:
# 转换图像格式
2025-08-18 11:06:42 +08:00
if msg.encoding in ['rgb8', 'bgr8']:
2025-08-18 20:34:03 +08:00
img = np.frombuffer(msg.data, dtype=np.uint8).reshape(
2025-08-18 11:06:42 +08:00
msg.height, msg.width, 3
)
if msg.encoding == 'bgr8':
2025-08-18 20:34:03 +08:00
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
2025-08-18 11:06:42 +08:00
2025-08-18 20:34:03 +08:00
# 定时截图逻辑
if current_time - self.last_capture_time >= self.capture_interval:
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
filename = f"{self.save_dir}/capture_{timestamp}.jpg"
cv2.imwrite(filename, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
self.get_logger().info(f"已保存截图: {filename}")
self.last_capture_time = current_time
2025-08-18 11:06:42 +08:00
2025-08-18 20:34:03 +08:00
# 实时显示
if self.gui_enabled:
cv2.imshow(self.window_name, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
cv2.waitKey(1)
except Exception as e:
self.get_logger().error(f"图像处理错误: {str(e)}")
2025-08-18 11:06:42 +08:00
def main(args=None):
rclpy.init(args=args)
demo = AICameraDemo()
try:
if demo.start_camera():
2025-08-18 20:34:03 +08:00
demo.get_logger().info("开始截图任务 (每3秒截图持续30秒)")
2025-08-18 11:06:42 +08:00
rclpy.spin(demo)
except KeyboardInterrupt:
2025-08-18 20:34:03 +08:00
demo.get_logger().info("手动终止")
2025-08-18 11:06:42 +08:00
finally:
demo.stop_camera()
2025-08-18 20:34:03 +08:00
cv2.destroyAllWindows()
2025-08-18 11:06:42 +08:00
demo.destroy_node()
rclpy.shutdown()
if __name__ == '__main__':
main()