autopilot/tmp/opi_rtsp/src-onnx/aiengineinferenceonnx.cpp

#include <QDebug>
#include <QThread>
#include "aiengineinferenceonnx.h"


AiEngineInferenceOnnx::AiEngineInferenceOnnx(QString modelPath, QObject *parent)
    : AiEngineInference{modelPath, parent}
{
    //qDebug() << "TUOMAS test mModelPath=" << mModelPath;
    mEngine = new InferenceEngine(modelPath.toStdString());
}


cv::Mat resizeAndPad(const cv::Mat& src) {
    // Desired size
    const int targetWidth = 640;
    const int targetHeight = 640;

    // 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) {
            // Width is greater than height
            newWidth = targetWidth;
            newHeight = static_cast<int>(targetWidth / aspectRatio);
        } else {
            // Height is greater than or equal to width
            newHeight = targetHeight;
            newWidth = static_cast<int>(targetHeight * aspectRatio);
        }
    }

    // Resize the original image if needed
    cv::Mat resized;
    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));

    // Copy the resized image to the top-left corner of the new image
    resized.copyTo(output(cv::Rect(0, 0, resized.cols, resized.rows)));

    return output;
}


void AiEngineInferenceOnnx::performInferenceSlot(cv::Mat frame)
{
    //qDebug() << "performInferenceSlot() in thread: " << QThread::currentThreadId();

    mActive = true;

    cv::Mat scaledImage = resizeAndPad(frame);

    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);

    AiEngineInferenceResult result;
    for (uint32_t i = 0; i < detections.size(); i++) {
        const Detection &detection = detections[i];

        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);
    }

    result.frame = mEngine->draw_labels(scaledImage.clone(), detections);
    emit resultsReady(result);

    mActive = false;
}