mirror of
https://github.com/Karllzy/cotton_color.git
synced 2025-11-09 03:03:53 +00:00
完成OnnxRunner部分以及test_onnx
This commit is contained in:
parent
6e0ad63020
commit
71d2692f1a
@ -1,5 +1,132 @@
|
||||
//
|
||||
// Created by zjc on 24-11-26.
|
||||
//
|
||||
|
||||
#include "OnnxRunner.h"
|
||||
|
||||
// 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 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 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 loadModel(const std::string& modelPath) {
|
||||
cv::dnn::Net net = cv::dnn::readNetFromONNX(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 preprocessImage(const cv::Mat& image, cv::dnn::Net& net, int& padTop, int& padLeft, float& scale) {
|
||||
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> performInference(cv::dnn::Net& net, const cv::Mat& inputImage) {
|
||||
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> applyNMS(std::vector<Detection>& detections) {
|
||||
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;
|
||||
}
|
||||
|
||||
@ -1,16 +1,39 @@
|
||||
//
|
||||
// Created by zjc on 24-11-26.
|
||||
//
|
||||
|
||||
#ifndef ONNXRUNNER_H
|
||||
#define ONNXRUNNER_H
|
||||
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include <opencv2/dnn/dnn.hpp>
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
#include <chrono>
|
||||
|
||||
const float CONFIDENCE_THRESHOLD = 0.2;
|
||||
const float NMS_THRESHOLD = 0.2;
|
||||
const int INPUT_WIDTH = 640;
|
||||
const int INPUT_HEIGHT = 640;
|
||||
|
||||
class OnnxRunner {
|
||||
|
||||
struct Detection {
|
||||
cv::Rect box;
|
||||
float confidence;
|
||||
};
|
||||
|
||||
// Class to measure elapsed time
|
||||
class Timer {
|
||||
public:
|
||||
Timer();
|
||||
void restart();
|
||||
void printElapsedTime(const std::string& message);
|
||||
|
||||
private:
|
||||
std::chrono::high_resolution_clock::time_point start_time;
|
||||
};
|
||||
|
||||
#endif //ONNXRUNNER_H
|
||||
// Function prototypes
|
||||
cv::Mat resizeAndPad(const cv::Mat& image, int targetWidth, int targetHeight, int& padTop, int& padLeft, float& scale, const cv::Scalar& padColor);
|
||||
cv::Mat createDetectionMask(const cv::Mat& originalImage, const std::vector<Detection>& detections, float scale, int padTop, int padLeft);
|
||||
cv::dnn::Net loadModel(const std::string& modelPath);
|
||||
cv::Mat preprocessImage(const cv::Mat& image, cv::dnn::Net& net, int& padTop, int& padLeft, float& scale);
|
||||
std::vector<Detection> performInference(cv::dnn::Net& net, const cv::Mat& inputImage);
|
||||
std::vector<Detection> applyNMS(std::vector<Detection>& detections);
|
||||
|
||||
#endif // ONNXRUNNER_H
|
||||
|
||||
@ -25,4 +25,4 @@ target_link_libraries(test_mask Matrox ${OpenCV_LIBS} ${MIL_LIBS})
|
||||
add_executable(test_onnx
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/test_onnx.cpp)
|
||||
# 链接 Matrox 模块和依赖库
|
||||
target_link_libraries(test_onnx Matrox ${OpenCV_LIBS} ${MIL_LIBS})
|
||||
target_link_libraries(test_onnx CVDL ${OpenCV_LIBS} ${MIL_LIBS})
|
||||
@ -1,66 +1,45 @@
|
||||
#include "vector"
|
||||
#include"iostream"
|
||||
#include"string"
|
||||
#include"Matrox/utils.h"
|
||||
#include"opencv2/opencv.hpp"
|
||||
#include "CVDL/OnnxRunner.h"
|
||||
|
||||
#define IMAGE_PATH MIL_TEXT("C:\\Users\\zjc\\Desktop\\8.bmp")
|
||||
#define SAVE_PATH MIL_TEXT("C:\\Users\\zjc\\Desktop\\suspect.png")
|
||||
|
||||
// int main() {
|
||||
// MIL_ID MilImage = M_NULL;
|
||||
// MIL_ID MilApplication = M_NULL, MilSystem = M_NULL, MilDisplay = M_NULL;
|
||||
//
|
||||
// MappAllocDefault(M_DEFAULT, &MilApplication, &MilSystem, &MilDisplay, M_NULL,
|
||||
// M_NULL);
|
||||
// MbufRestore(IMAGE_PATH, MilSystem, &MilImage);
|
||||
// cv::Mat opencvImage = milToMat(MilImage);
|
||||
// imshow("opencv", opencvImage);
|
||||
// // MosGetch();
|
||||
//
|
||||
// return 0;
|
||||
// }
|
||||
int main() {
|
||||
MIL_ID MilApplication = M_NULL, MilSystem = M_NULL, MilDisplay = M_NULL;
|
||||
MIL_ID MilImage = M_NULL;
|
||||
std::string modelPath = "C:\\Users\\zjc\\Desktop\\dimo_11.14.onnx";
|
||||
std::string imagePath = "C:\\Users\\zjc\\Desktop\\dimo.bmp";
|
||||
Timer timer1;
|
||||
|
||||
// 初始化 MIL 应用程序、系统和显示
|
||||
MappAllocDefault(M_DEFAULT, &MilApplication, &MilSystem, &MilDisplay, M_NULL, M_NULL);
|
||||
if (MilApplication == M_NULL || MilSystem == M_NULL || MilDisplay == M_NULL) {
|
||||
std::cerr << "MIL Initialization failed!" << std::endl;
|
||||
// Load the model
|
||||
cv::dnn::Net net = loadModel(modelPath);
|
||||
timer1.printElapsedTime("Time to load the model");
|
||||
|
||||
// Read the input image
|
||||
cv::Mat image = cv::imread(imagePath);
|
||||
if (image.empty()) {
|
||||
std::cerr << "Could not read the image: " << imagePath << std::endl;
|
||||
return -1;
|
||||
}
|
||||
|
||||
// 加载图像
|
||||
MbufRestore(IMAGE_PATH, MilSystem, &MilImage);
|
||||
if (MilImage == M_NULL) {
|
||||
std::cerr << "Failed to load MIL image!" << std::endl;
|
||||
MappFreeDefault(MilApplication, MilSystem, MilDisplay, M_NULL,M_NULL);
|
||||
return -1;
|
||||
}
|
||||
// Preprocess image
|
||||
int padTop, padLeft;
|
||||
float scale;
|
||||
cv::Mat inputImage = preprocessImage(image, net, padTop, padLeft, scale);
|
||||
timer1.printElapsedTime("Time to preprocess image");
|
||||
|
||||
// 转换并显示
|
||||
cv::Mat opencvImage = milToMat(MilImage);
|
||||
if (!opencvImage.empty()) {
|
||||
cv::imshow("opencv", opencvImage);
|
||||
cv::waitKey(0);
|
||||
}
|
||||
std:: string savepath="C:\\Users\\zjc\\Desktop\\suspect.png";
|
||||
cv::imwrite(savepath,opencvImage);
|
||||
// 释放资源
|
||||
MbufFree(MilImage);
|
||||
MappFreeDefault(MilApplication, MilSystem, MilDisplay, M_NULL,M_NULL); // 使用 MappFreeDefault 代替 MappFree
|
||||
// Perform inference
|
||||
std::vector<Detection> detections = performInference(net, inputImage);
|
||||
|
||||
// Apply Non-Maximum Suppression
|
||||
std::vector<Detection> finalDetections = applyNMS(detections);
|
||||
std::cout << "Number of detections after NMS: " << finalDetections.size() << std::endl;
|
||||
|
||||
// Create and show the detection mask
|
||||
cv::Mat detectionMask = createDetectionMask(image, finalDetections, scale, padTop, padLeft);
|
||||
|
||||
cv::imshow("Detection Mask", detectionMask);
|
||||
|
||||
// Save the result
|
||||
|
||||
std::string savepath = "C:\\Users\\zjc\\Desktop\\suspect_mask.png";
|
||||
cv::imwrite(savepath, detectionMask);
|
||||
timer1.printElapsedTime("Time to run inference");
|
||||
|
||||
cv::waitKey(0);
|
||||
return 0;
|
||||
}
|
||||
|
||||
// int main() {
|
||||
// cv::Mat img = cv::imread("C:\\Users\\zjc\\Desktop\\suspect.png");
|
||||
// if (img.empty()) {
|
||||
// std::cout << "图像加载失败!" << std::endl;
|
||||
// return -1;
|
||||
// }
|
||||
//
|
||||
// // 处理图像
|
||||
// processImage(img);
|
||||
// }
|
||||
Loading…
Reference in New Issue
Block a user