main.cpp

#include <yaml-cpp/yaml.h>

#include "utils/include/draw.hpp"

const std::string modelConfig_path = "/home/lin/workspace/tensorrt_detect/configs/detectConfig.yaml"; // 模型配置文件路径
const YAML::Node modelConfig = YAML::LoadFile(modelConfig_path);                                      // 模型配置

const std::string modelPath = modelConfig["modelPath"].as<std::string>(); // 模型路径
const int inputSize = modelConfig["inputSize"].as<int>();                 // 图像输入大小
const float scoreThreshold = modelConfig["scoreThreshold"].as<float>();   // 置信度阈值
const float nmsThreshold = modelConfig["nmsThreshold"].as<float>();       // 非极大值抑制阈值
const std::vector<std::string> classNames = []()
{
    std::vector<std::string> tmp;
    for (const auto &item : modelConfig["classNames"])
        tmp.emplace_back(item.second.as<std::string>());

    return tmp;
}(); // 模型类别列表
const bool showFlag = modelConfig["showFlag"].as<bool>(); // 绘制, 展示检测结果标志

char waitKey_Flag;

int main(int argc, char const *argv[])
{

    Model detectModel = Model(modelPath, inputSize, scoreThreshold, nmsThreshold);
    bool detectFlag = false;
    while (true)
    {
        cv::Mat frame = cv::imread("/home/lin/workspace/tensorrt_detect/assets/bus.jpg");

        double start_time = static_cast<double>(cv::getTickCount());
        
        detectFlag = detectModel.Detect(frame);

        std::cout << (static_cast<double>(cv::getTickCount()) - start_time) / cv::getTickFrequency() << std::endl;

        // 绘制, 展示检测结果
        if (showFlag)
        {
            drawDetect(frame, detectModel.detectResults, classNames);
            detectFlag = false;
            cv::imshow("img", frame);
            waitKey_Flag = cv::waitKey(1);
            if (waitKey_Flag == 113) // q
                break;
        }
    }

    return 0;
}