217 lines
8.2 KiB
C++
217 lines
8.2 KiB
C++
#include "onnxrunner.h"
|
|
#include<Qdebug>
|
|
|
|
// Timer class implementation
|
|
Timer::Timer() : start_time(std::chrono::high_resolution_clock::now()) {}
|
|
|
|
void Timer::restart() {
|
|
start_time = std::chrono::high_resolution_clock::now();
|
|
}
|
|
|
|
void Timer::printElapsedTime(const std::string& message) {
|
|
auto end_time = std::chrono::high_resolution_clock::now();
|
|
std::chrono::duration<double> elapsed = end_time - start_time;
|
|
std::cout << message << ": " << elapsed.count() << " seconds" << std::endl;
|
|
start_time = end_time;
|
|
}
|
|
|
|
// Resize and pad input image
|
|
cv::Mat ONNXRunner::resizeAndPad(const cv::Mat& image, int targetWidth, int targetHeight, int& padTop, int& padLeft, float& scale, const cv::Scalar& padColor) {
|
|
int originalWidth = image.cols;
|
|
int originalHeight = image.rows;
|
|
|
|
scale = std::min((float)targetWidth / originalWidth, (float)targetHeight / originalHeight);
|
|
int newWidth = static_cast<int>(originalWidth * scale);
|
|
int newHeight = static_cast<int>(originalHeight * scale);
|
|
|
|
cv::Mat resizedImage;
|
|
cv::resize(image, resizedImage, cv::Size(newWidth, newHeight));
|
|
|
|
padTop = (targetHeight - newHeight) / 2;
|
|
int padBottom = targetHeight - newHeight - padTop;
|
|
padLeft = (targetWidth - newWidth) / 2;
|
|
int padRight = targetWidth - newWidth - padLeft;
|
|
|
|
cv::Mat paddedImage;
|
|
cv::copyMakeBorder(resizedImage, paddedImage, padTop, padBottom, padLeft, padRight, cv::BORDER_CONSTANT, padColor);
|
|
|
|
return paddedImage;
|
|
}
|
|
|
|
// Create detection mask
|
|
cv::Mat ONNXRunner::createDetectionMask(const cv::Mat& originalImage, const std::vector<Detection>& detections, float scale, int padTop, int padLeft) {
|
|
cv::Mat mask = cv::Mat::zeros(originalImage.size(), CV_8UC1); // Single channel mask
|
|
|
|
for (const auto& detection : detections) {
|
|
int x = static_cast<int>((detection.box.x - padLeft) / scale);
|
|
int y = static_cast<int>((detection.box.y - padTop) / scale);
|
|
int w = static_cast<int>(detection.box.width / scale);
|
|
int h = static_cast<int>(detection.box.height / scale);
|
|
|
|
x = std::max(0, std::min(x, originalImage.cols - 1));
|
|
y = std::max(0, std::min(y, originalImage.rows - 1));
|
|
w = std::min(w, originalImage.cols - x);
|
|
h = std::min(h, originalImage.rows - y);
|
|
|
|
cv::rectangle(mask, cv::Rect(x, y, w, h), cv::Scalar(255), cv::FILLED); // White color for detections
|
|
}
|
|
|
|
return mask;
|
|
}
|
|
|
|
// Load the ONNX model
|
|
cv::dnn::Net ONNXRunner::loadModel(const std::string& modelPath) {
|
|
// Load the ONNX model
|
|
cv::dnn::Net net = cv::dnn::readNetFromONNX(modelPath);
|
|
if (net.empty()) {
|
|
qCritical() << "加载ONNX模型失败:" << QString::fromStdString(modelPath);
|
|
} else {
|
|
qDebug() << "成功加载ONNX模型:" << QString::fromStdString(modelPath);
|
|
}
|
|
// net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA); // Use CUDA backend
|
|
// net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA); // Run on GPU
|
|
return net;
|
|
}
|
|
|
|
// Preprocess image for model input
|
|
cv::Mat ONNXRunner::preprocessImage(const cv::Mat& image, cv::dnn::Net& net, int& padTop, int& padLeft, float& scale) const {
|
|
cv::Scalar padColor(128, 128, 128); // Gray padding
|
|
cv::Mat inputImage = resizeAndPad(image, INPUT_WIDTH, INPUT_HEIGHT, padTop, padLeft, scale, padColor);
|
|
cv::Mat blob = cv::dnn::blobFromImage(inputImage, 1 / 255.0, cv::Size(INPUT_WIDTH, INPUT_HEIGHT), cv::Scalar(0, 0, 0), true, false);
|
|
net.setInput(blob);
|
|
return inputImage;
|
|
}
|
|
|
|
// Perform inference on the input image
|
|
std::vector<Detection> ONNXRunner::performInference(cv::dnn::Net& net, const cv::Mat& inputImage) const {
|
|
std::vector<Detection> detections;
|
|
cv::Mat output = net.forward();
|
|
float* data = (float*)output.data;
|
|
|
|
for (int i = 0; i < 25200; ++i) {
|
|
float confidence = data[i * 6 + 4];
|
|
if (confidence >= CONFIDENCE_THRESHOLD) {
|
|
float cx = data[i * 6];
|
|
float cy = data[i * 6 + 1];
|
|
float w = data[i * 6 + 2];
|
|
float h = data[i * 6 + 3];
|
|
|
|
cx = cx * inputImage.cols / INPUT_WIDTH;
|
|
cy = cy * inputImage.rows / INPUT_HEIGHT;
|
|
w = w * inputImage.cols / INPUT_WIDTH;
|
|
h = h * inputImage.rows / INPUT_HEIGHT;
|
|
|
|
int left = static_cast<int>(cx - w / 2);
|
|
int top = static_cast<int>(cy - h / 2);
|
|
int width = static_cast<int>(w);
|
|
int height = static_cast<int>(h);
|
|
|
|
left = std::max(0, std::min(left, inputImage.cols - 1));
|
|
top = std::max(0, std::min(top, inputImage.rows - 1));
|
|
width = std::min(width, inputImage.cols - left);
|
|
height = std::min(height, inputImage.rows - top);
|
|
|
|
detections.push_back({cv::Rect(left, top, width, height), confidence});
|
|
}
|
|
}
|
|
|
|
return detections;
|
|
}
|
|
|
|
// Apply Non-Maximum Suppression
|
|
std::vector<Detection> ONNXRunner::applyNMS(std::vector<Detection>& detections) const {
|
|
std::vector<int> indices;
|
|
std::vector<cv::Rect> boxes;
|
|
std::vector<float> scores;
|
|
|
|
for (const auto& detection : detections) {
|
|
boxes.push_back(detection.box);
|
|
scores.push_back(detection.confidence);
|
|
}
|
|
|
|
cv::dnn::NMSBoxes(boxes, scores, CONFIDENCE_THRESHOLD, NMS_THRESHOLD, indices);
|
|
std::vector<Detection> finalDetections;
|
|
for (int idx : indices) {
|
|
finalDetections.push_back(detections[idx]);
|
|
}
|
|
|
|
return finalDetections;
|
|
}
|
|
|
|
|
|
void ONNXRunner::load(const std::string& modelPath) {
|
|
QString qModelPath = QString::fromStdString(modelPath);
|
|
|
|
// 检查模型文件是否存在
|
|
if (!QFile::exists(qModelPath)) {
|
|
qCritical() << "ONNX模型文件不存在:" << qModelPath;
|
|
// 根据需求，可以选择抛出异常、设置状态标志或采取其他错误处理措施
|
|
return;
|
|
}
|
|
|
|
// 加载模型
|
|
this->net = ONNXRunner::loadModel(modelPath);
|
|
|
|
// 检查模型是否成功加载
|
|
if (this->net.empty()) {
|
|
qCritical() << "加载ONNX模型失败:" << qModelPath;
|
|
// 根据需求，可以选择抛出异常、设置状态标志或采取其他错误处理措施
|
|
return;
|
|
}
|
|
|
|
// 创建一个空的输入矩阵作为预热数据（假定模型输入是 RGB 图像）
|
|
cv::Mat dummyInput = cv::Mat::zeros(INPUT_HEIGHT, INPUT_WIDTH, CV_8UC3); // 640x640 的全零矩阵
|
|
|
|
// 调用 predict 方法进行预热，避免第一次推理时加载模型导致延迟
|
|
try {
|
|
for(int i = 0; i < 2; i++)
|
|
this->predict(dummyInput);
|
|
qDebug() << "ONNX模型预热完成";
|
|
}
|
|
catch (const std::exception& e) {
|
|
qCritical() << "ONNXRunner::predict 异常:" << e.what();
|
|
// 处理异常，例如记录日志或设置状态标志
|
|
}
|
|
};
|
|
|
|
std::vector<Detection> ONNXRunner::predict(const cv::Mat &image) {
|
|
cv::dnn::Net net = this->net;
|
|
cv::Mat inputImage = preprocessImage(image, net, this->pad_top, this->pad_left, this->scale);
|
|
|
|
// Perform inference
|
|
std::vector<Detection> detections = performInference(net, inputImage);
|
|
// Apply Non-Maximum Suppression
|
|
std::vector<Detection> finalDetections = applyNMS(detections);
|
|
return finalDetections;
|
|
}
|
|
|
|
cv::Mat ONNXRunner::postProcess(const std::vector<Detection>& detections, const cv::Mat& image) const {
|
|
// Create and show the detection mask
|
|
cv::Mat detectionMask = ONNXRunner::createDetectionMask(image, detections, this->scale, this->pad_top, this->pad_left);
|
|
return detectionMask;
|
|
}
|
|
|
|
void ONNXRunner::warm_up()
|
|
{
|
|
// 检查模型是否成功加载
|
|
if (this->net.empty()) {
|
|
qCritical() << "加载ONNX模型失败" ;
|
|
// 根据需求，可以选择抛出异常、设置状态标志或采取其他错误处理措施
|
|
return;
|
|
}
|
|
|
|
// 创建一个空的输入矩阵作为预热数据（假定模型输入是 RGB 图像）
|
|
cv::Mat dummyInput = cv::Mat::zeros(INPUT_HEIGHT, INPUT_WIDTH, CV_8UC3); // 640x640 的全零矩阵
|
|
|
|
// 调用 predict 方法进行预热，避免第一次推理时加载模型导致延迟
|
|
try {
|
|
for(int i = 0; i < 10; i++)
|
|
this->predict(dummyInput);
|
|
qDebug() << "ONNX模型预热完成";
|
|
}
|
|
catch (const std::exception& e) {
|
|
qCritical() << "ONNXRunner::predict 异常:" << e.what();
|
|
// 处理异常，例如记录日志或设置状态标志
|
|
}
|
|
}
|