Fully working opi_rtsp on PC with YOLOv8 ONNX models

This commit is contained in:
Tuomas Järvinen
2024-07-10 18:37:33 +02:00
parent 896307d296
commit 3d39d8fd99
5 changed files with 272 additions and 220 deletions
+60 -40
View File
@@ -3,34 +3,39 @@
#include "aiengineinferenceonnx.h"
const int INFERENCE_SQUARE_WIDTH = 640;
const int INFERENCE_SQUARE_HEIGHT = 640;
AiEngineInferenceOnnx::AiEngineInferenceOnnx(QString modelPath, QObject *parent)
: AiEngineInference{modelPath, parent}
: AiEngineInference{modelPath, parent},
mInference(modelPath.toStdString(), cv::Size(640, 640), "classes.txt")
{
//qDebug() << "TUOMAS test mModelPath=" << mModelPath;
mEngine = new InferenceEngine(modelPath.toStdString());
//mEngine = new InferenceEngine(modelPath.toStdString());
//mInference = new Inference(modelPath.toStdString(), cv::Size(INFERENCE_SQUARE_WIDTH, INFERENCE_SQUARE_HEIGHT), "classes.txt");
}
cv::Mat resizeAndPad(const cv::Mat& src) {
// Desired size
const int targetWidth = 640;
const int targetHeight = 640;
cv::Mat resizeAndPad(const cv::Mat& src)
{
// Calculate the aspect ratio
float aspectRatio = static_cast<float>(src.cols) / src.rows;
// Determine new size while maintaining aspect ratio
int newWidth = src.cols;
int newHeight = src.rows;
if (src.cols > targetWidth || src.rows > targetHeight) {
if (aspectRatio > 1) {
if (src.cols > INFERENCE_SQUARE_WIDTH || src.rows > INFERENCE_SQUARE_HEIGHT) {
if (aspectRatio > 1)
{
// Width is greater than height
newWidth = targetWidth;
newHeight = static_cast<int>(targetWidth / aspectRatio);
} else {
newWidth = INFERENCE_SQUARE_WIDTH;
newHeight = static_cast<int>(INFERENCE_SQUARE_WIDTH / aspectRatio);
}
else {
// Height is greater than or equal to width
newHeight = targetHeight;
newWidth = static_cast<int>(targetHeight * aspectRatio);
newHeight = INFERENCE_SQUARE_HEIGHT;
newWidth = static_cast<int>(INFERENCE_SQUARE_HEIGHT * aspectRatio);
}
}
@@ -39,7 +44,7 @@ cv::Mat resizeAndPad(const cv::Mat& src) {
cv::resize(src, resized, cv::Size(newWidth, newHeight));
// Create a new 640x640 image with a black background
cv::Mat output(targetHeight, targetWidth, src.type(), cv::Scalar(0, 0, 0));
cv::Mat output(INFERENCE_SQUARE_HEIGHT, INFERENCE_SQUARE_WIDTH, src.type(), cv::Scalar(0, 0, 0));
// Copy the resized image to the top-left corner of the new image
resized.copyTo(output(cv::Rect(0, 0, resized.cols, resized.rows)));
@@ -50,35 +55,50 @@ cv::Mat resizeAndPad(const cv::Mat& src) {
void AiEngineInferenceOnnx::performInferenceSlot(cv::Mat frame)
{
//qDebug() << "performInferenceSlot() in thread: " << QThread::currentThreadId();
try {
//qDebug() << "performInferenceSlot() in thread: " << QThread::currentThreadId();
mActive = true;
mActive = true;
cv::Mat scaledImage = resizeAndPad(frame);
cv::Mat scaledImage = resizeAndPad(frame);
std::vector<Detection> detections = mInference.runInference(scaledImage);
AiEngineInferenceResult result;
int orig_width = scaledImage.cols;
int orig_height = scaledImage.rows;
std::vector<float> input_tensor_values = mEngine->preprocessImage(scaledImage);
std::vector<float> results = mEngine->runInference(input_tensor_values);
float confidence_threshold = 0.4;
std::vector<Detection> detections = mEngine->filterDetections(results, confidence_threshold, mEngine->input_shape[2], mEngine->input_shape[3], orig_width, orig_height);
for (uint i = 0; i < detections.size(); ++i) {
const Detection &detection = detections[i];
AiEngineInferenceResult result;
for (uint32_t i = 0; i < detections.size(); i++) {
const Detection &detection = detections[i];
// Add detected objects to the results
AiEngineObject object;
object.classId = detection.class_id;
object.propability = detection.confidence;
object.rectangle.top = detection.box.y;
object.rectangle.left = detection.box.x;
object.rectangle.bottom = detection.box.y + detection.box.height;
object.rectangle.right = detection.box.x + detection.box.width;
result.objects.append(object);
AiEngineObject object;
object.classId = detection.class_id;
object.propability = detection.confidence;
object.rectangle.top = detection.bbox.y;
object.rectangle.left = detection.bbox.x;
object.rectangle.bottom = detection.bbox.y + detection.bbox.height;
object.rectangle.right = detection.bbox.x + detection.bbox.width;
result.objects.append(object);
/*
// Draw box and text
cv::Rect box = detection.box;
cv::Scalar color = detection.color;
cv::rectangle(frame, box, color, 2);
std::string classString = detection.className + ' ' + std::to_string(detection.confidence).substr(0, 4);
//std::cout << "classString:" << classString << std::endl;
cv::Size textSize = cv::getTextSize(classString, cv::FONT_HERSHEY_DUPLEX, 1, 2, 0);
cv::Rect textBox(box.x, box.y - 40, textSize.width + 10, textSize.height + 20);
cv::rectangle(scaledImage, textBox, color, cv::FILLED);
cv::putText(scaledImage, classString, cv::Point(box.x + 5, box.y - 10), cv::FONT_HERSHEY_DUPLEX, 1, cv::Scalar(0, 0, 0), 2, 0);
*/
}
if (result.objects.empty() == false) {
result.frame = mInference.drawLabels(scaledImage, detections);
emit resultsReady(result);
}
mActive = false;
}
catch (const cv::Exception& e) {
std::cerr << "performInferenceSlot() Error: " << e.what() << std::endl;
}
result.frame = mEngine->draw_labels(scaledImage.clone(), detections);
emit resultsReady(result);
mActive = false;
}