sladro
e13cb3659c
- 创建基本项目结构和目录 - 添加CMake构建系统 - 实现基础的配置解析功能 - 添加YOLO推理框架支持 - 集成RTSP和视频流处理功能 - 添加性能监控和日志系统
87 lines
2.3 KiB
C++
87 lines
2.3 KiB
C++
#pragma once
|
|
|
|
#include <memory>
|
|
#include <string>
|
|
#include <vector>
|
|
#include <NvInfer.h>
|
|
#include <cuda_runtime.h>
|
|
#include "types.hpp"
|
|
#include "preprocess.hpp"
|
|
|
|
namespace pipeline {
|
|
|
|
// 前向声明
|
|
namespace detail {
|
|
class Logger;
|
|
class CudaStream;
|
|
class CudaBuffer;
|
|
}
|
|
|
|
class TrtInference {
|
|
public:
|
|
explicit TrtInference(const InferenceConfig& config);
|
|
~TrtInference();
|
|
|
|
// 禁用拷贝
|
|
TrtInference(const TrtInference&) = delete;
|
|
TrtInference& operator=(const TrtInference&) = delete;
|
|
|
|
// 加载模型
|
|
bool loadEngine();
|
|
|
|
// ONNX转换为TensorRT引擎
|
|
bool convertOnnxToEngine(const std::string& onnx_path, const std::string& engine_path);
|
|
|
|
// 批量推理
|
|
bool infer(const std::vector<cv::Mat>& images, std::vector<DetectionResult>& results);
|
|
|
|
// 获取配置
|
|
const InferenceConfig& getConfig() const { return config_; }
|
|
|
|
// 获取状态
|
|
bool isLoaded() const { return context_ != nullptr; }
|
|
|
|
private:
|
|
// 预处理
|
|
bool preprocess(const std::vector<cv::Mat>& images);
|
|
|
|
// 后处理
|
|
bool postprocess(std::vector<DetectionResult>& results);
|
|
|
|
// NMS处理
|
|
void doNMS(std::vector<BBox>& boxes, float nms_thresh);
|
|
|
|
// 创建执行上下文
|
|
bool createContext();
|
|
|
|
// 分配CUDA内存
|
|
bool allocateBuffers();
|
|
|
|
// 释放资源
|
|
void destroy();
|
|
|
|
private:
|
|
float calculateIOU(const BBox& box1, const BBox& box2);
|
|
|
|
private:
|
|
InferenceConfig config_; // 配置参数
|
|
std::unique_ptr<detail::Logger> logger_; // TensorRT日志器
|
|
std::unique_ptr<detail::CudaStream> stream_; // CUDA流
|
|
|
|
// TensorRT相关
|
|
nvinfer1::IRuntime* runtime_{nullptr}; // TensorRT运行时
|
|
nvinfer1::ICudaEngine* engine_{nullptr}; // TensorRT引擎
|
|
nvinfer1::IExecutionContext* context_{nullptr}; // 执行上下文
|
|
|
|
// CUDA内存
|
|
std::vector<std::unique_ptr<detail::CudaBuffer>> input_buffers_; // 输入缓冲区
|
|
std::vector<std::unique_ptr<detail::CudaBuffer>> output_buffers_; // 输出缓冲区
|
|
std::vector<void*> bindings_; // 绑定指针
|
|
std::vector<std::string> output_names_; // 输出层名称
|
|
|
|
// 预处理器
|
|
std::unique_ptr<Preprocessor> preprocessor_; // 预处理器实例
|
|
};
|
|
|
|
} // namespace pipeline
|