diff --git a/misc/rtsp_ai_player/aiengine.cpp b/misc/rtsp_ai_player/aiengine.cpp index ff17f36..d4357c4 100644 --- a/misc/rtsp_ai_player/aiengine.cpp +++ b/misc/rtsp_ai_player/aiengine.cpp @@ -2,6 +2,7 @@ #include #include "aiengine.h" #include "aiengineinference.h" +#include "aiengineimagesaver.h" #if defined(OPI5_BUILD) #include "src-opi5/aiengineinferenceopi5.h" @@ -82,6 +83,13 @@ void AiEngine::inferenceResultsReceivedSlot(AiEngineInferenceResult result) } cv::imshow("Received Frame", result.frame); + +#ifdef SAVE_IMAGES + static int imageCounter = 0; + AiEngineImageSaver *saver = new AiEngineImageSaver(result.frame, ++imageCounter); + saver->start(); + connect(saver, &AiEngineImageSaver::finished, saver, &QObject::deleteLater); +#endif } diff --git a/misc/rtsp_ai_player/aiengineimagesaver.cpp b/misc/rtsp_ai_player/aiengineimagesaver.cpp new file mode 100644 index 0000000..8ff29cb --- /dev/null +++ b/misc/rtsp_ai_player/aiengineimagesaver.cpp @@ -0,0 +1,31 @@ +#include +#include +#include +#include "aiengineimagesaver.h" + +AiEngineImageSaver::AiEngineImageSaver(const cv::Mat &image, int imageNumber, QObject *parent) + : QThread(parent), image(image.clone()), imageNumber(imageNumber) +{ +} + +AiEngineImageSaver::~AiEngineImageSaver() +{ + wait(); +} + +void AiEngineImageSaver::run() +{ + if (image.empty() || image.channels() == 0) { + qWarning() << "AiEngineImageSaver. Empty image or no channels, nothing to save."; + return; + } + + // Calculate the size of the upper half + int halfHeight = image.rows / 2; + cv::Mat upperHalf = image(cv::Rect(0, 0, image.cols, halfHeight)); + + // Use bpm to reduce encoding time. + QString filePath = QString("/tmp/image-%1.bmp").arg(imageNumber, 5, 10, QChar('0')); + + cv::imwrite(filePath.toStdString(), upperHalf); +} diff --git a/misc/rtsp_ai_player/aiengineimagesaver.h b/misc/rtsp_ai_player/aiengineimagesaver.h new file mode 100644 index 0000000..93d9c65 --- /dev/null +++ b/misc/rtsp_ai_player/aiengineimagesaver.h @@ -0,0 +1,21 @@ +#pragma once + +#include +#include +#include + +class AiEngineImageSaver : public QThread +{ + Q_OBJECT + +public: + AiEngineImageSaver(const cv::Mat &image, int imageNumber, QObject *parent = nullptr); + ~AiEngineImageSaver(); + +protected: + void run() override; + +private: + cv::Mat image; + int imageNumber; +}; diff --git a/misc/rtsp_ai_player/aiengineinference.cpp b/misc/rtsp_ai_player/aiengineinference.cpp index 9a8f968..38f059a 100644 --- a/misc/rtsp_ai_player/aiengineinference.cpp +++ b/misc/rtsp_ai_player/aiengineinference.cpp @@ -5,7 +5,105 @@ AiEngineInference::AiEngineInference(QString modelPath, QObject *parent) : QObject{parent}, mModelPath(modelPath), mActive(false) -{} +{ +#ifdef YOLO_ONNX + mClassNames = { + "person", + "bicycle", + "car", + "motorcycle", + "airplane", + "bus", + "train", + "truck", + "boat", + "traffic light", + "fire hydrant", + "stop sign", + "parking meter", + "bench", + "bird", + "cat", + "dog", + "horse", + "sheep", + "cow", + "elephant", + "bear", + "zebra", + "giraffe", + "backpack", + "umbrella", + "handbag", + "tie", + "suitcase", + "frisbee", + "skis", + "snowboard", + "sports ball", + "kite", + "baseball bat", + "baseball glove", + "skateboard", + "surfboard", + "tennis racket", + "bottle", + "wine glass", + "cup", + "fork", + "knife", + "spoon", + "bowl", + "banana", + "apple", + "sandwich", + "orange", + "broccoli", + "carrot", + "hot dog", + "pizza", + "donut", + "cake", + "chair", + "couch", + "potted plant", + "bed", + "dining table", + "toilet", + "tv", + "laptop", + "mouse", + "remote", + "keyboard", + "cell phone", + "microwave", + "oven", + "toaster", + "sink", + "refrigerator", + "book", + "clock", + "vase", + "scissors", + "teddy bear", + "hair drier", + "toothbrush" + }; +#else + mClassNames = { + "Armoured vehicle", + "Truck", + "Vehicle", + "Artillery", + "Shadow artillery", + "Trenches", + "Military man", + "Tyre tracks", + "Additional protection tank", + "Smoke" + }; +#endif +} bool AiEngineInference::isActive(void) diff --git a/misc/rtsp_ai_player/aiengineinference.h b/misc/rtsp_ai_player/aiengineinference.h index d27f7c6..192a7b1 100644 --- a/misc/rtsp_ai_player/aiengineinference.h +++ b/misc/rtsp_ai_player/aiengineinference.h @@ -40,6 +40,7 @@ protected: QString mModelPath; bool mActive; int mNumber; + QVector mClassNames; public slots: virtual void performInferenceSlot(cv::Mat frame) = 0; diff --git a/misc/rtsp_ai_player/rtsp_ai_player.pro b/misc/rtsp_ai_player/rtsp_ai_player.pro index 05b14f1..1c2366b 100644 --- a/misc/rtsp_ai_player/rtsp_ai_player.pro +++ b/misc/rtsp_ai_player/rtsp_ai_player.pro @@ -16,9 +16,15 @@ else { } yolo_onnx { + message("Using YOLOv8 ONNX models for indoor testing.") QMAKE_CXXFLAGS += -DYOLO_ONNX } +save_images { + message("Saving inference images to /tmp.") + QMAKE_CXXFLAGS += -DSAVE_IMAGES +} + opi5 { message("OPI5 build") CONFIG += link_pkgconfig diff --git a/misc/rtsp_ai_player/src-onnx-runtime/aiengineinferenceonnxruntime.cpp b/misc/rtsp_ai_player/src-onnx-runtime/aiengineinferenceonnxruntime.cpp index 15d658b..8f7d343 100644 --- a/misc/rtsp_ai_player/src-onnx-runtime/aiengineinferenceonnxruntime.cpp +++ b/misc/rtsp_ai_player/src-onnx-runtime/aiengineinferenceonnxruntime.cpp @@ -14,104 +14,7 @@ AiEngineInferencevOnnxRuntime::AiEngineInferencevOnnxRuntime(QString modelPath, mPredictor(modelPath.toStdString(), confThreshold, iouThreshold, maskThreshold) { qDebug() << "TUOMAS AiEngineInferencevOnnxRuntime() mModelPath=" << mModelPath; - -#ifdef YOLO_ONNX - mClassNames = { - "person", - "bicycle", - "car", - "motorcycle", - "airplane", - "bus", - "train", - "truck", - "boat", - "traffic light", - "fire hydrant", - "stop sign", - "parking meter", - "bench", - "bird", - "cat", - "dog", - "horse", - "sheep", - "cow", - "elephant", - "bear", - "zebra", - "giraffe", - "backpack", - "umbrella", - "handbag", - "tie", - "suitcase", - "frisbee", - "skis", - "snowboard", - "sports ball", - "kite", - "baseball bat", - "baseball glove", - "skateboard", - "surfboard", - "tennis racket", - "bottle", - "wine glass", - "cup", - "fork", - "knife", - "spoon", - "bowl", - "banana", - "apple", - "sandwich", - "orange", - "broccoli", - "carrot", - "hot dog", - "pizza", - "donut", - "cake", - "chair", - "couch", - "potted plant", - "bed", - "dining table", - "toilet", - "tv", - "laptop", - "mouse", - "remote", - "keyboard", - "cell phone", - "microwave", - "oven", - "toaster", - "sink", - "refrigerator", - "book", - "clock", - "vase", - "scissors", - "teddy bear", - "hair drier", - "toothbrush" - }; -#else - mClassNames = { - "Armoured vehicle", - "Truck", - "Vehicle", - "Artillery", - "Shadow artillery", - "Trenches", - "Military man", - "Tyre tracks", - "Additional protection tank", - "Smoke" - }; -#endif + qDebug() << "AiEngineInferencevOnnxRuntime() mClassNames.size() =" << mClassNames.size(); } @@ -183,6 +86,11 @@ void AiEngineInferencevOnnxRuntime::performInferenceSlot(cv::Mat frame) for (uint i = 0; i < detections.size(); i++) { const Yolov8Result &detection = detections[i]; + if (detection.classId >= mClassNames.size()) { + qDebug() << "performInferenceSlot() invalid classId =" << detection.classId; + continue; + } + // Add detected objects to the results AiEngineObject object; object.classId = detection.classId; diff --git a/misc/rtsp_ai_player/src-onnx-runtime/aiengineinferenceonnxruntime.h b/misc/rtsp_ai_player/src-onnx-runtime/aiengineinferenceonnxruntime.h index b010da9..71b953f 100644 --- a/misc/rtsp_ai_player/src-onnx-runtime/aiengineinferenceonnxruntime.h +++ b/misc/rtsp_ai_player/src-onnx-runtime/aiengineinferenceonnxruntime.h @@ -17,5 +17,4 @@ public slots: private: cv::Mat drawLabels(const cv::Mat &image, const std::vector &detections); YOLOPredictor mPredictor; - QVector mClassNames; }; diff --git a/misc/rtsp_ai_player/src-opi5/aiengineinferenceopi5.cpp b/misc/rtsp_ai_player/src-opi5/aiengineinferenceopi5.cpp index 15077bf..94f3b7f 100644 --- a/misc/rtsp_ai_player/src-opi5/aiengineinferenceopi5.cpp +++ b/misc/rtsp_ai_player/src-opi5/aiengineinferenceopi5.cpp @@ -90,6 +90,12 @@ void AiEngineInferenceOpi5::drawObjects(cv::Mat& image, const object_detect_resu for (int i = 0; i < result_list.count; i++) { const object_detect_result& result = result_list.results[i]; + if (result.cls_id >= mClassNames.size()) { + continue; + } + + fprintf(stderr, "TUOMAS [%d] prop = %f\n", i, result.prop); + int left = result.box.left; int top = result.box.top; int right = result.box.right; @@ -99,9 +105,10 @@ void AiEngineInferenceOpi5::drawObjects(cv::Mat& image, const object_detect_resu // Text char c_text[256]; - sprintf(c_text, "%s %.1f%%", coco_cls_to_name(result.cls_id), result.prop * 100); + //sprintf(c_text, "%s %d%%", coco_cls_to_name(result.cls_id), (int)(round(result.prop * 100))); + sprintf(c_text, "%s %d%%", mClassNames[result.cls_id].toStdString().c_str(), (int)(round(result.prop * 100))); cv::Point textOrg(left, top - 5); - cv::putText(image, std::string(c_text), textOrg, cv::FONT_HERSHEY_SIMPLEX, result.prop, cv::Scalar(0, 0, 255), 1); + cv::putText(image, std::string(c_text), textOrg, cv::FONT_HERSHEY_COMPLEX, result.prop, cv::Scalar(0, 0, 255), 1, cv::LINE_AA); } } @@ -117,6 +124,7 @@ void AiEngineInferenceOpi5::performInferenceSlot(cv::Mat frame) object_detect_result_list od_results; int ret = inference_yolov8_model(&mRrknnAppCtx0, &imgBuffer, &od_results, mNumber); + freeImageBuffer(imgBuffer); if (ret != 0) { qDebug() << "AiEngineInferenceOpi5::performInferenceSlot() failure! ret: " << ret; mActive = false; @@ -138,8 +146,6 @@ void AiEngineInferenceOpi5::performInferenceSlot(cv::Mat frame) } drawObjects(scaledFrame, od_results); - freeImageBuffer(imgBuffer); - result.frame = scaledFrame.clone(); emit resultsReady(result);