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

class ObjectDetectionNode(Node):
    def __init__(self):
        super().__init__('object_detection_node')
        
        # 创建图像订阅者
        self.image_sub = self.create_subscription(
            Image, 'camera/image_raw', self.image_callback, 10)
        
        # 创建检测结果发布者
        self.detection_pub = self.create_publisher(Image, 'detection_image', 10)
        
        # 初始化OpenCV桥接器
        self.bridge = CvBridge()
        
        # 加载TensorFlow模型
        self.model = tf.saved_model.load('path/to/saved_model')
        
        # 获取类别名称
        self.class_names = ['person', 'bicycle', 'car', ...]  # 根据实际模型调整
        
        self.get_logger().info('Object detection node initialized')
    
    def image_callback(self, msg):
        try:
            # 将ROS图像消息转换为OpenCV格式
            cv_image = self.bridge.imgmsg_to_cv2(msg, 'bgr8')
            
            # 预处理图像
            input_tensor = tf.convert_to_tensor(cv_image)
            input_tensor = input_tensor[tf.newaxis, ...]
            
            # 运行目标检测
            detections = self.model(input_tensor)
            
            # 处理检测结果
            num_detections = int(detections.pop('num_detections'))
            detections = {key: value[0, :num_detections].numpy()
                          for key, value in detections.items()}
            detections['num_detections'] = num_detections
            
            # 过滤置信度低的检测结果
            detections['detection_classes'] = detections['detection_classes'].astype(np.int64)
            scores = detections['detection_scores']
            classes = detections['detection_classes']
            boxes = detections['detection_boxes']
            
            # 在图像上绘制检测结果
            for i in range(len(scores)):
                if scores[i] > 0.5:  # 置信度阈值
                    class_id = classes[i]
                    class_name = self.class_names[class_id]
                    box = boxes[i]
                    
                    # 转换边界框坐标
                    h, w, _ = cv_image.shape
                    ymin, xmin, ymax, xmax = box
                    xmin = int(xmin * w)
                    xmax = int(xmax * w)
                    ymin = int(ymin * h)
                    ymax = int(ymax * h)
                    
                    # 绘制边界框和标签
                    cv2.rectangle(cv_image, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)
                    cv2.putText(cv_image, f'{class_name}: {scores[i]:.2f}',
                               (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
            
            # 发布结果图像
            output_msg = self.bridge.cv2_to_imgmsg(cv_image, 'bgr8')
            output_msg.header = msg.header
            self.detection_pub.publish(output_msg)
            
        except Exception as e:
            self.get_logger().error(f'Error processing image: {str(e)}')

def main(args=None):
    rclpy.init(args=args)
    node = ObjectDetectionNode()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()