21 lines
1022 B
C++
21 lines
1022 B
C++
#pragma once
|
|
|
|
#include <vector>
|
|
#include <array>
|
|
#include <opencv2/core/cuda.hpp>
|
|
#include "NvInfer.h" // Include for nvinfer1::Dims and nvinfer1::Dims3
|
|
|
|
template <typename T>
|
|
class IEngine {
|
|
public:
|
|
virtual ~IEngine() = default;
|
|
virtual bool buildLoadNetwork(std::string onnxModelPath, const std::array<float, 3> &subVals = {0.f, 0.f, 0.f},
|
|
const std::array<float, 3> &divVals = {1.f, 1.f, 1.f}, bool normalize = true) = 0;
|
|
virtual bool loadNetwork(std::string trtModelPath, const std::array<float, 3> &subVals = {0.f, 0.f, 0.f},
|
|
const std::array<float, 3> &divVals = {1.f, 1.f, 1.f}, bool normalize = true) = 0;
|
|
virtual bool runInference(const std::vector<std::vector<cv::cuda::GpuMat>> &inputs,
|
|
std::vector<std::vector<std::vector<T>>> &featureVectors) = 0;
|
|
virtual const std::vector<nvinfer1::Dims3> &getInputDims() const = 0;
|
|
virtual const std::vector<nvinfer1::Dims> &getOutputDims() const = 0;
|
|
};
|