3. 基于Camera的2D目标检测

任务目标

本案例将基于Autoware的相机驱动获取图像数据，加载预训练的YOLOv4-tiny模型，通过TensorRT加速推理，在ROS 2节点中实现实时2D目标检测，并发布检测框与类别标签的可视化结果。

环境与依赖

• Autoware Universe（基于ROS 2 Humble）
• CUDA 11.8 + cuDNN 8.6 + TensorRT 8.5
• OpenCV 4.x（ROS 2自带）
• 预训练权重：yolov4-tiny.weights与配置文件yolov4-tiny.cfg（从AlexeyAB/darknet获取）
• 相机驱动：usb_cam或Autoware的sensor_kit相机驱动，输出sensor_msgs/Image

核心实现

以下代码实现了一个ROS 2节点，订阅相机图像话题，使用TensorRT加载YOLOv4-tiny引擎进行推理，并发布检测结果与可视化图像。节点采用异步推理模式，避免阻塞图像回调。

// yolo_detector_node.cpp
#include <rclcpp/rclcpp.hpp>
#include <sensor_msgs/msg/image.hpp>
#include <vision_msgs/msg/detection2_d_array.hpp>
#include <cv_bridge/cv_bridge.hpp>
#include <opencv2/opencv.hpp>
#include <NvInfer.h>
#include <fstream>
#include <memory>

using std::placeholders::_1;

// TensorRT 日志回调
class Logger : public nvinfer1::ILogger {
    void log(Severity severity, const char* msg) noexcept override {
        if (severity == Severity::kERROR) std::cerr << "[TRT] " << msg << std::endl;
    }
} g_logger;

class YoloDetector : public rclcpp::Node {
public:
    YoloDetector() : Node("yolo_detector_node") {
        // 参数声明
        this->declare_parameter("engine_path", "yolov4-tiny.engine");
        this->declare_parameter("conf_threshold", 0.5);
        this->declare_parameter("nms_threshold", 0.4);
        this->declare_parameter("input_width", 416);
        this->declare_parameter("input_height", 416);

        std::string engine_path = this->get_parameter("engine_path").as_string();
        conf_th_ = this->get_parameter("conf_threshold").as_double();
        nms_th_ = this->get_parameter("nms_threshold").as_double();
        input_w_ = this->get_parameter("input_width").as_int();
        input_h_ = this->get_parameter("input_height").as_int();

        // 加载TensorRT引擎
        loadEngine(engine_path);

        // 订阅与发布
        image_sub_ = this->create_subscription<sensor_msgs::msg::Image>(
            "/camera/image_raw", 10, std::bind(&YoloDetector::imageCallback, this, _1));
        det_pub_ = this->create_publisher<vision_msgs::msg::Detection2DArray>("/detections", 10);
        vis_pub_ = this->create_publisher<sensor_msgs::msg::Image>("/detection_image", 10);

        RCLCPP_INFO(this->get_logger(), "YOLO detector node started");
    }

private:
    void loadEngine(const std::string& path) {
        std::ifstream file(path, std::ios::binary);
        if (!file.good()) {
            RCLCPP_ERROR(this->get_logger(), "Engine file not found: %s", path.c_str());
            throw std::runtime_error("Engine file missing");
        }
        file.seekg(0, std::ios::end);
        size_t size = file.tellg();
        file.seekg(0, std::ios::beg);
        std::vector<char> blob(size);
        file.read(blob.data(), size);
        file.close();

        runtime_ = std::unique_ptr<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(g_logger));
        engine_ = std::unique_ptr<nvinfer1::ICudaEngine>(
            runtime_->deserializeCudaEngine(blob.data(), size));
        context_ = std::unique_ptr<nvinfer1::IExecutionContext>(engine_->createExecutionContext());

        // 分配输入输出缓冲区
        cudaMalloc(&input_buf_, input_w_ * input_h_ * 3 * sizeof(float));
        cudaMalloc(&output_buf_, 2535 * 85 * sizeof(float)); // 3个尺度 x 845个anchor x 85
        cudaStreamCreate(&stream_);
    }

    void imageCallback(const sensor_msgs::msg::Image::SharedPtr msg) {
        auto cv_img = cv_bridge::toCvCopy(msg, "bgr8")->image;
        cv::Mat resized;
        cv::resize(cv_img, resized, cv::Size(input_w_, input_h_));

        // HWC -> CHW 并归一化到[0,1]
        cv::Mat float_img;
        resized.convertTo(float_img, CV_32FC3, 1.0 / 255.0);
        std::vector<cv::Mat> channels(3);
        cv::split(float_img, channels);
        std::vector<float> input_data(input_w_ * input_h_ * 3);
        for (int c = 0; c < 3; c++) {
            memcpy(input_data.data() + c * input_w_ * input_h_,
                   channels[c].data, input_w_ * input_h_ * sizeof(float));
        }

        // 异步推理
        cudaMemcpyAsync(input_buf_, input_data.data(),
                        input_w_ * input_h_ * 3 * sizeof(float),
                        cudaMemcpyHostToDevice, stream_);
        context_->enqueueV2(&input_buf_, stream_, nullptr);
        cudaMemcpyAsync(output_buf_, output_buf_,
                        2535 * 85 * sizeof(float),
                        cudaMemcpyDeviceToHost, stream_);
        cudaStreamSynchronize(stream_);

        // 后处理：解析输出并执行NMS
        std::vector<vision_msgs::msg::Detection2D> detections;
        std::vector<cv::Rect> boxes;
        std::vector<float> scores;
        std::vector<int> class_ids;

        float* output = static_cast<float*>(output_buf_);
        for (int i = 0; i < 2535; i++) {
            float* row = output + i * 85;
            float obj_conf = row[4];
            if (obj_conf < conf_th_) continue;

            // 找出最高类别分数
            float max_cls = 0;
            int max_id = 0;
            for (int j = 5; j < 85; j++) {
                if (row[j] > max_cls) {
                    max_cls = row[j];
                    max_id = j - 5;
                }
            }
            float final_conf = obj_conf * max_cls;
            if (final_conf < conf_th_) continue;

            // 解码边界框（yolo格式：中心x,y + 宽高）
            float cx = row[0] * input_w_;
            float cy = row[1] * input_h_;
            float w = row[2] * input_w_;
            float h = row[3] * input_h_;
            int x = static_cast<int>(cx - w / 2);
            int y = static_cast<int>(cy - h / 2);
            boxes.emplace_back(x, y, static_cast<int>(w), static_cast<int>(h));
            scores.push_back(final_conf);
            class_ids.push_back(max_id);
        }

        // NMS
        std::vector<int> keep;
        cv::dnn::NMSBoxes(boxes, scores, conf_th_, nms_th_, keep);

        for (int idx : keep) {
            vision_msgs::msg::Detection2D det;
            det.bbox.center.position.x = boxes[idx].x + boxes[idx].width / 2.0;
            det.bbox.center.position.y = boxes[idx].y + boxes[idx].height / 2.0;
            det.bbox.size_x = boxes[idx].width;
            det.bbox.size_y = boxes[idx].height;
            det.results.resize(1);
            det.results[0].id = class_ids[idx];
            det.results[0].score = scores[idx];
            detections.push_back(det);
        }

        // 发布检测结果
        auto det_msg = vision_msgs::msg::Detection2DArray();
        det_msg.header = msg->header;
        det_msg.detections = detections;
        det_pub_->publish(det_msg);

        // 绘制并发布可视化图像
        for (int idx : keep) {
            cv::rectangle(cv_img, boxes[idx], cv::Scalar(0, 255, 0), 2);
            std::string label = "ID:" + std::to_string(class_ids[idx]) +
                                " " + std::to_string(scores[idx]).substr(0, 4);
            cv::putText(cv_img, label, cv::Point(boxes[idx].x, boxes[idx].y - 5),
                        cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 255, 0), 1);
        }
        vis_pub_->publish(*cv_bridge::CvImage(msg->header, "bgr8", cv_img).toImageMsg());
    }

    // 成员变量
    rclcpp::Subscription<sensor_msgs::msg::Image>::SharedPtr image_sub_;
    rclcpp::Publisher<vision_msgs::msg::Detection2DArray>::SharedPtr det_pub_;
    rclcpp::Publisher<sensor_msgs::msg::Image>::SharedPtr vis_pub_;
    std::unique_ptr<nvinfer1::IRuntime> runtime_;
    std::unique_ptr<nvinfer1::ICudaEngine> engine_;
    std::unique_ptr<nvinfer1::IExecutionContext> context_;
    float *input_buf_, *output_buf_;
    cudaStream_t stream_;
    double conf_th_, nms_th_;
    int input_w_, input_h_;
};

int main(int argc, char** argv) {
    rclcpp::init(argc, argv);
    try {
        rclcpp::spin(std::make_shared<YoloDetector>());
    } catch (const std::exception& e) {
        RCLCPP_ERROR(rclcpp::get_logger("yolo_detector"), "Exception: %s", e.what());
    }
    rclcpp::shutdown();
    return 0;
}

构建与运行

# 1. 生成TensorRT引擎（首次运行前执行）
cd ~/autoware_ws
trtexec --onnx=yolov4-tiny.onnx --saveEngine=yolov4-tiny.engine --fp16

# 2. 编译节点
colcon build --packages-select yolo_detector

# 3. 启动相机驱动（以usb_cam为例）
ros2 launch usb_cam camera.launch.py

# 4. 运行检测节点
ros2 run yolo_detector yolo_detector_node --ros-args \
  -p engine_path:=/path/to/yolov4-tiny.engine

# 5. 可视化结果（新终端）
rqt_image_view /detection_image

验证方法

• 话题检查：运行ros2 topic echo /detections，当画面中有目标物体时，应看到包含bbox坐标和score的检测消息输出。
• 可视化确认：在rqt_image_view中观察/detection_image话题，图像上应绘制绿色边界框和类别ID标签，框位置与画面中物体位置一致。
• 性能指标：在节点日志中观察推理耗时（可自行添加计时），在GPU环境下应达到实时处理（>15 FPS）。

常见问题

• 引擎文件缺失：需先通过trtexec将ONNX模型转换为TensorRT引擎文件，转换时需要与运行环境相同的GPU型号。
• 输出尺寸不匹配：YOLOv4-tiny在416x416输入下输出2535个候选框（3个尺度），若修改输入尺寸需同步调整输出缓冲区大小。
• 类别ID映射：当前输出为COCO 80类索引（0-79），如需显示类别名称，需加载coco.names文件进行映射。