【opencv】dnn示例-vit_tracker.cpp 使用OpenCV库和ViTTrack模型实现的视频追踪程序

这段代码是一个使用OpenCV库和ViTTrack模型实现的视频追踪程序。程序通过摄像头或视频文件获取图像序列，并对选定的目标对象进行实时追踪。

代码主要分为以下几个部分：

导入必要的库：程序开始时先导入了iostream，cmath以及相关OpenCV的库。

参数解析：通过定义一个参数字典keys来解析命令行输入的参数，这些参数包括输入的视频源、模型的路径以及计算后端等。

追踪器的创建与设置：使用给定的模型文件以及命令行参数中的后端和目标设备选项，创建一个ViT追踪器对象。

视频捕获：程序尝试打开摄像头或读取视频文件。如果指定了摄像头或视频文件但不能成功打开，程序将输出错误信息。

目标对象的选择：获取视频的第一帧，并让用户选择一个区域作为追踪的初始目标。

追踪循环：程序进入一个循环，不断从视频源获取新的帧，并使用追踪器更新目标位置。每一次追踪完成后，都会输出框架的计数、预测分数、目标位置方框以及每帧追踪耗时。

追踪结果展示：将追踪的结果（目标位置的矩形框及其分数）覆盖在原图像上，并在窗口中展示出来。

退出处理：程序会在按下ESC键后退出追踪循环，并关闭程序。

整体而言，代码实现了一个基于ViT追踪模型的视频目标追踪系统，包含了从视频源获取图像、初始化追踪目标、连续追踪目标以及展示追踪结果等功能。程序使用OpenCV的DNN模块加载ONNX格式的深度学习模型，并通过选择合适的后端（例如CUDA）和目标设备来加速计算过程。

// VitTracker
// model: https://github.com/opencv/opencv_zoo/tree/main/models/object_tracking_vittrack


#include <iostream> // 引入输入输出流库
#include <cmath> // 引入数学运算库


#include <opencv2/dnn.hpp> // 引入OpenCV深度神经网络模块
#include <opencv2/imgproc.hpp> // 引入OpenCV图像处理模块
#include <opencv2/highgui.hpp> // 引入OpenCV高级用户界面模块
#include <opencv2/video.hpp> // 引入OpenCV视频处理模块


using namespace cv; // 使用cv命名空间简化代码
using namespace cv::dnn; // 使用cv::dnn命名空间简化代码


const char *keys = // 定义命令行参数
        "{ help     h  |   | Print help message }" // 打印帮助信息
        "{ input    i  |   | Full path to input video folder, the specific camera index. (empty for camera 0) }" // 输入视频路径或摄像头索引（默认为0）
        "{ net    | object_tracking_vittrack_2023sep.onnx | Path to onnx model of vitTracker.onnx}" // vitTracker.onnx模型的路径
        "{ backend     | 5 | Choose one of computation backends: " // 选择计算后端
                            "0: automatically (by default), " // 自动选择
                            "1: Halide language (http://halide-lang.org/), " // Halide语言
                            "2: Intel's Deep Learning Inference Engine (https://software.intel.com/openvino-toolkit), " // 英特尔深度学习推理引擎
                            "3: OpenCV implementation, " // OpenCV实现
                            "4: VKCOM, " // VKCOM
                            "5: CUDA }," // CUDA
        "{ target      | 6 | Choose one of target computation devices: " // 选择目标计算设备
                            "0: CPU target (by default), " // CPU
                            "1: OpenCL, " // OpenCL
                            "2: OpenCL fp16 (half-float precision), " // OpenCL fp16 (半精度浮点)
                            "3: VPU, " // VPU
                            "4: Vulkan, " // Vulkan
                            "6: CUDA, " // CUDA
                            "7: CUDA fp16 (half-float preprocess) }" // CUDA fp16 (半精度预处理)
;


static
int run(int argc, char** argv) // 定义主要运行函数
{
    // Parse command line arguments.
    CommandLineParser parser(argc, argv, keys); // 命令行参数解析


    if (parser.has("help")) // 如果存在帮助参数
    {
        parser.printMessage(); // 打印帮助信息
        return 0;
    }


    std::string inputName = parser.get<String>("input"); // 获取输入源
    std::string net = parser.get<String>("net"); // 获取模型路径
    int backend = parser.get<int>("backend"); // 获取后端参数
    int target = parser.get<int>("target"); // 获取目标计算设备参数


    Ptr<TrackerVit> tracker; // 声明一个ViT跟踪器的智能指针
    try // 尝试加载跟踪器和模型
    {
        TrackerVit::Params params; // 创建跟踪器参数结构
        params.net = samples::findFile(net); // 找到或保持模型文件路径
        params.backend = backend; // 设置后端参数
        params.target = target; // 设置目标计算设备参数
        tracker = TrackerVit::create(params); // 创建跟踪器实例
    }
    catch (const cv::Exception& ee) // 如果捕获到异常
    {
        std::cerr << "Exception: " << ee.what() << std::endl; // 打印异常信息
        std::cout << "Can't load the network by using the following files:" << std::endl; // 说明无法加载网络
        std::cout << "net : " << net << std::endl; // 打印网络路径
        return 2;
    }


    const std::string winName = "vitTracker"; // 窗口名称
    namedWindow(winName, WINDOW_AUTOSIZE); // 创建具有自动大小的窗口


    // Open a video file or an image file or a camera stream.
    VideoCapture cap; // 创建视频捕获对象


    if (inputName.empty() || (isdigit(inputName[0]) && inputName.size() == 1)) // 如果输入为空或仅有单个数字（表示摄像头索引）
    {
        int c = inputName.empty() ? 0 : inputName[0] - '0'; // 默认为0或转换数字为摄像头索引
        std::cout << "Trying to open camera #" << c << " ..." << std::endl;
        if (!cap.open(c)) // 尝试打开摄像头
        {
            std::cout << "Capture from camera #" << c << " didn't work. Specify -i=<video> parameter to read from video file" << std::endl; // 打开失败提示
            return 2;
        }
    }
    else if (inputName.size()) // 如果输入非空
    {
        inputName = samples::findFileOrKeep(inputName); // 查找或保持文件路径
        if (!cap.open(inputName)) // 尝试打开视频文件
        {
            std::cout << "Could not open: " << inputName << std::endl; // 打开失败提示
            return 2;
        }
    }


    // Read the first image.
    Mat image; // 创建一个Mat对象用于存放图像
    cap >> image; // 捕获一帧图像
    if (image.empty()) // 如果图像为空
    {
        std::cerr << "Can't capture frame!" << std::endl; // 抓取图像失败提示
        return 2;
    }


    Mat image_select = image.clone(); // 克隆捕获的图像以便用户选择跟踪区域
    putText(image_select, "Select initial bounding box you want to track.", Point(0, 15), FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0, 255, 0)); // 在图像上添加文字指导用户选择跟踪区域
    putText(image_select, "And Press the ENTER key.", Point(0, 35), FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0, 255, 0)); // 提示用户按下回车键确认选择


    Rect selectRect = selectROI(winName, image_select); // 让用户选择一个区域并返回矩形坐标
    std::cout << "ROI=" << selectRect << std::endl; // 输出用户选择的区域


    tracker->init(image, selectRect); // 初始化跟踪器并使用用户选择的区域


    TickMeter tickMeter; // 创建一个时钟计时器对象


    for (int count = 0; ; ++count) // 循环处理视频帧
    {
        cap >> image; // 捕获一帧图像
        if (image.empty()) // 如果图像为空
        {
            std::cerr << "Can't capture frame " << count << ". End of video stream?" << std::endl; // 抓取图像失败提示
            break;
        }


        Rect rect; // 创建一个矩形对象


        tickMeter.start(); // 开始计时
        bool ok = tracker->update(image, rect); // 更新跟踪器状态并返回跟踪框的位置
        tickMeter.stop(); // 停止计时


        float score = tracker->getTrackingScore(); // 获取跟踪评分


        std::cout << "frame " << count << // 输出当前帧数
            ": predicted score=" << score << // 输出预测评分
            "  rect=" << rect << // 输出跟踪框位置
            "  time=" << tickMeter.getTimeMilli() << "ms" << // 输出计时时间
            std::endl;


        Mat render_image = image.clone(); // 克隆当前帧以便绘制跟踪结果


        if (ok) // 如果跟踪成功
        {
            rectangle(render_image, rect, Scalar(0, 255, 0), 2); // 在渲染图像上绘制跟踪框


            std::string timeLabel = format("Inference time: %.2f ms", tickMeter.getTimeMilli()); // 格式化输出计时标签
            std::string scoreLabel = format("Score: %f", score); // 格式化输出评分标签
            putText(render_image, timeLabel, Point(0, 15), FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0, 255, 0)); // 将计时标签绘制在图像上
            putText(render_image, scoreLabel, Point(0, 35), FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0, 255, 0)); // 将评分标签绘制在图像上
        }


        imshow(winName, render_image); // 显示渲染后图像


        tickMeter.reset(); // 重置计时器


        int c = waitKey(1); // 等待按键
        if (c == 27 /*ESC*/) // 如果按下ESC键
            break; // 退出循环
    }


    std::cout << "Exit" << std::endl; // 输出退出提示
    return 0; // 返回成功
}


int main(int argc, char **argv) // 主程序入口函数
{
    try // 尝试运行主功能函数
    {
        return run(argc, argv); // 调用run函数并返回结果
    }
    catch (const std::exception& e) // 捕获任何抛出的异常
    {
        std::cerr << "FATAL: C++ exception: " << e.what() << std::endl; // 输出异常信息
        return 1; // 返回失败
    }
}

该代码是使用OpenCV库实现的一个简单的视频跟踪应用程序。它主要通过解析命令行输入不同的参数来调整视频输入源、选择神经网络模型、计算后端以及目标计算设备。代码首先尝试加载一个基于ViT（Vision Transformer）的目标跟踪器模型。接下来，它通过用户选择的图像区域初始化跟踪器并开始对视频流中的目标进行跟踪。每个循环中，都会更新跟踪器并计算其评分以及跟踪目标的位置，然后将跟踪结果绘制在每帧上并实时显示给用户。如果用户按下ESC键，则会退出循环结束程序。