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

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


static const float confThreshold = 0.4f;
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;

    mClassNames = {
        "Armoured vehicle",
        "Truck",
        "Vehicle",
        "Artillery",
        "Shadow artillery",
        "Trenches",
        "Military man",
        "Tyre tracks",
        "Additional protection tank",
        "Smoke"
    };
}


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);
        std::string label = mClassNames[detection.classId] + ": " + std::to_string(detection.conf);

        int baseLine;
        cv::Size labelSize = cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);

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

        cv::putText(
            result,
            label,
            cv::Point(
                detection.box.x,
                detection.box.y),
            cv::FONT_HERSHEY_SIMPLEX,
            0.5,
            cv::Scalar(0, 0, 0),
            1);
    }

    return result;
}


void AiEngineInferencevOnnxRuntime::performInferenceSlot(cv::Mat frame)
{
    mActive = true;

    cv::Mat scaledImage                  = resizeAndPad(frame);
    std::vector<Yolov8Result> detections = mPredictor.predict(scaledImage);
    AiEngineInferenceResult result;

    for (uint i = 0; i < detections.size(); i++) {
        const Yolov8Result &detection = detections[i];

        // Add detected objects to the results
        AiEngineObject object;
        object.classId          = 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) {
        qDebug() << __PRETTY_FUNCTION__ << "detections:" << detections.size();
        result.frame = drawLabels(scaledImage, detections);
        emit resultsReady(result);
    }

    mActive = false;
}