#include "aiengineinference.h"


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)
{
    return mActive;
}


cv::Mat AiEngineInference::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 > INFERENCE_SQUARE_WIDTH || src.rows > INFERENCE_SQUARE_HEIGHT) {
        if (aspectRatio > 1)
        {
            // Width is greater than height
            newWidth = INFERENCE_SQUARE_WIDTH;
            newHeight = static_cast<int>(INFERENCE_SQUARE_WIDTH / aspectRatio);
        }
        else {
            // Height is greater than or equal to width
            newHeight = INFERENCE_SQUARE_HEIGHT;
            newWidth = static_cast<int>(INFERENCE_SQUARE_HEIGHT * 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(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)));

    return output;
}