autopilot/misc/rtsp_ai_player/src-onnx-runtime/aiengineinferenceonnxruntime.cpp

#include <QDebug>
#include <QThread>
#include <vector>
#include "aiengineinferenceonnxruntime.h"


static const float confThreshold = 0.2f;
static const float iouThreshold  = 0.4f;
static const float maskThreshold = 0.5f;


AiEngineInferencevOnnxRuntime::AiEngineInferencevOnnxRuntime(QString modelPath, QObject *parent) :
    AiEngineInference{modelPath, parent},
    mPredictor(modelPath.toStdString(), confThreshold, iouThreshold, maskThreshold)
{
    qDebug() << "TUOMAS AiEngineInferencevOnnxRuntime() mModelPath=" << mModelPath;
    qDebug() << "AiEngineInferencevOnnxRuntime() mClassNames.size() =" << mClassNames.size();
}


cv::Mat AiEngineInferencevOnnxRuntime::drawLabels(const cv::Mat &image, const std::vector<Yolov8Result> &detections)
{
    cv::Mat result = image.clone();

    for (const auto &detection : detections)
    {
        cv::rectangle(result, detection.box, cv::Scalar(0, 255, 0), 2);

        int confidence = roundf(detection.conf * 100);
        std::string label = mClassNames[detection.classId].toStdString() + ": " + std::to_string(confidence) + "%";

        int baseLine;
        cv::Size labelSize = cv::getTextSize(label, cv::FONT_HERSHEY_COMPLEX, 0.5, 1, &baseLine);
        cv::Point labelOrigin(detection.box.x, detection.box.y - labelSize.height - baseLine);

        cv::rectangle(
            result,
            labelOrigin,
            cv::Point(detection.box.x + labelSize.width, detection.box.y),
            cv::Scalar(255, 255, 255),
            cv::FILLED);

        cv::putText(
            result,
            label,
            cv::Point(detection.box.x, detection.box.y - baseLine + 2),
            cv::FONT_HERSHEY_COMPLEX,
            0.5,
            cv::Scalar(0, 0, 0),
            1,
            cv::LINE_AA);

    }

    return result;
}


void AiEngineInferencevOnnxRuntime::performInferenceSlot(cv::Mat frame)
{
    qDebug() << __PRETTY_FUNCTION__;

    try {
        mActive = true;
        cv::Mat scaledImage                  = resizeAndPad(frame);
        std::vector<Yolov8Result> detections = mPredictor.predict(scaledImage);

#ifdef YOLO_ONNX
        // Only keep following detected objects.
        // car    = 2
        // train  = 6
        // cup    = 41
        // banana = 46
        auto it = std::remove_if(detections.begin(), detections.end(),
                                 [](const Yolov8Result& result) {
                                     return result.classId != 2 &&
                                            result.classId != 6 &&
                                            result.classId != 41 &&
                                            result.classId != 46;
                                 });
        detections.erase(it, detections.end());
#endif

        AiEngineInferenceResult result;

        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;
            object.classStr         = mClassNames[detection.classId];
            object.propability      = detection.conf;
            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);
        }

        if (result.objects.empty() == false) {
            result.frame = drawLabels(scaledImage, detections);
            emit resultsReady(result);
        }

        mActive = false;
    }
    catch (const cv::Exception& e) {
        std::cout << "OpenCV exception caught: " << e.what() << std::endl;
    }
    catch (const std::exception& e) {
        std::cout << "Standard exception caught: " << e.what() << std::endl;
    }
    catch (...) {
        std::cout << "Unknown exception caught" << std::endl;
    }
}


void AiEngineInferencevOnnxRuntime::initialize(int number)
{
    (void)number;
}