YOLOModel.hpp

Go to the documentation of this file.

 
#ifndef YOLO_MODEL_HPP
#define YOLO_MODEL_HPP
 
#include "../../AutonomyConstants.h"
 
#include <nlohmann/json.hpp>
#include <opencv2/opencv.hpp>
#include <torch/script.h>
#include <torch/torch.h>
 
 
 
namespace yolomodel
{
 
    struct Detection
    {
        public:
            // Define public struct attributes.
 
            int nClassID;               // The class index of the object. Dependent on class order when trained.
            std::string szClassName;    // The class name of the object. This is dependent on the class names used when training.
            float fConfidence;          // The detection confidence of the object.
            cv::Rect cvBoundingBox;     // An object used to access the dimensions and other properties of the objects bounding box.
    };
 
 
    inline void NonMaxSuppression(std::vector<Detection>& vObjects,
                                  std::vector<int>& vClassIDs,
                                  std::vector<float>& vClassConfidences,
                                  std::vector<cv::Rect>& vBoundingBoxes,
                                  float fMinObjectConfidence,
                                  float fNMSThreshold)
    {
        // Create instance variables.
        std::vector<int> vNMSValidIndices;
 
        // Perform Non-Max Suppression using OpenCV's implementation.
        cv::dnn::NMSBoxes(vBoundingBoxes, vClassConfidences, fMinObjectConfidence, fNMSThreshold, vNMSValidIndices);
 
        // Loop through each valid index.
        for (int nValidIndex : vNMSValidIndices)
        {
            // Create new Detection struct.
            Detection stNewDetection;
            // Repackage prediction data into easy-to-use struct.
            stNewDetection.nClassID      = vClassIDs[nValidIndex];
            stNewDetection.fConfidence   = vClassConfidences[nValidIndex];
            stNewDetection.cvBoundingBox = vBoundingBoxes[nValidIndex];
 
            // Append new object detection to objects vector.
            vObjects.emplace_back(stNewDetection);
        }
    }
 
 
    inline void DrawDetections(cv::Mat& cvInputFrame, std::vector<Detection>& vObjects)
    {
        // Loop through each detection.
        for (Detection stObject : vObjects)
        {
            // Calculate the hue value based on the class ID.
            int nHue = static_cast<int>(stObject.nClassID % 256);
            // Set saturation and value to 1.0 for full intensity colors.
            int nSaturation = 255;
            int nValue      = 255;
 
            // Convert HSV to RGB
            cv::Mat cvHSV(1, 1, CV_8UC3, cv::Scalar(nHue, nSaturation, nValue));
            cv::cvtColor(cvHSV, cvHSV, cv::COLOR_HSV2BGR);
            // Extract the RGB values
            cv::Vec3b cvConvertedValues = cvHSV.at<cv::Vec3b>(0, 0);
            cv::Scalar cvBoxColor(cvConvertedValues[2], cvConvertedValues[1], cvConvertedValues[0]);
 
            // Draw bounding box onto image.
            cv::rectangle(cvInputFrame, stObject.cvBoundingBox, cvBoxColor, 2);
            // Draw classID background box onto image.
            cv::rectangle(cvInputFrame,
                          cv::Point(stObject.cvBoundingBox.x, stObject.cvBoundingBox.y - 20),
                          cv::Point(stObject.cvBoundingBox.x + stObject.cvBoundingBox.width, stObject.cvBoundingBox.y),
                          cvBoxColor,
                          cv::FILLED);
            // Draw class text onto image.
            cv::putText(cvInputFrame,
                        std::to_string(stObject.nClassID) + " " + std::to_string(stObject.fConfidence),
                        cv::Point(stObject.cvBoundingBox.x, stObject.cvBoundingBox.y - 5),
                        cv::FONT_HERSHEY_SIMPLEX,
                        0.5,
                        cv::Scalar(255, 255, 255));
        }
    }
 
 
    namespace pytorch
    {
 
        class PyTorchInterpreter
        {
            public:
                // Declare public enums that are specific to and used within this class.
                enum class HardwareDevices
                {
                    eCPU,    // The CPU device.
                    eCUDA    // The CUDA device.
                };
 
                // Declare public methods and member variables.
 
 
                PyTorchInterpreter(std::string szModelPath, HardwareDevices eHardwareDevice = HardwareDevices::eCUDA)
                {
                    // Initialize member variables.
                    m_szModelPath      = szModelPath;
                    m_bReady           = false;
                    m_cvModelInputSize = cv::Size(640, 640);
                    m_szModelTask      = "Unknown";
                    m_vClassLabels     = std::vector<std::string>();
 
                    // Translate the hardware device enum to a torch device.
                    switch (eHardwareDevice)
                    {
                        case HardwareDevices::eCPU: m_trDevice = torch::kCPU; break;
                        case HardwareDevices::eCUDA: m_trDevice = torch::kCUDA; break;
                        default: m_trDevice = torch::kCPU; break;
                    }
 
                    // Submit logger message.
                    LOG_INFO(logging::g_qSharedLogger, "Attempting to load model {} onto device {}", szModelPath, m_trDevice.str());
 
                    // Check if the model path is valid.
                    if (!std::filesystem::exists(szModelPath))
                    {
                        // Submit logger message.
                        LOG_ERROR(logging::g_qSharedLogger, "Model path {} does not exist!", szModelPath);
                        return;
                    }
                    // Check if the device is available.
                    if (!torch::cuda::is_available() && m_trDevice == torch::kCUDA)
                    {
                        // Submit logger message.
                        LOG_ERROR(logging::g_qSharedLogger, "CUDA device is not available, falling back to CPU.");
                        m_trDevice = torch::kCPU;
                        return;
                    }
                    else
                    {
                        // Submit logger message.
                        LOG_INFO(logging::g_qSharedLogger, "Using device: {}", m_trDevice.str());
                    }
 
                    // Finally, attempt to load the model.
                    try
                    {
                        // Load the model and set it to eval mode for inference.
                        torch::jit::ExtraFilesMap trExtraConfigFiles{{"config.txt", ""}};
                        m_trModel = torch::jit::load(szModelPath, m_trDevice, trExtraConfigFiles);
                        m_trModel.eval();
 
                        // Use nlohmann json to parse the config file.
                        nlohmann::json jConfig = nlohmann::json::parse(trExtraConfigFiles.at("config.txt"));
                        // Get the input image size for the model.
                        m_cvModelInputSize = cv::Size(jConfig["imgsz"][0], jConfig["imgsz"][1]);
                        m_szModelTask      = jConfig["task"];
                        for (const auto& item : jConfig["names"].items())
                        {
                            m_vClassLabels.push_back(item.value());
                        }
                        // Submit the config json as a debug message.
                        LOG_DEBUG(logging::g_qSharedLogger, "Model config: {}", jConfig.dump(4));
 
                        // Check if the model is empty.
                        if (m_trModel.get_methods().empty())
                        {
                            LOG_ERROR(logging::g_qSharedLogger, "Model is empty! Check if the correct model file was provided.");
                            return;
                        }
                        // Check if the model did not move to the expected device.
                        if (m_trModel.buffers().size() > 0)
                        {
                            // Get the device of the model.
                            torch::Device model_device = m_trModel.buffers().begin().operator->().device();
                            if (model_device != m_trDevice)
                            {
                                LOG_ERROR(logging::g_qSharedLogger, "Model did not move to the expected device! Model is on: {}", model_device.str());
                                return;
                            }
                        }
 
                        // Model is ready for inference.
                        LOG_INFO(logging::g_qSharedLogger,
                                 "Model successfully loaded and set to eval mode. The model is a {} model, and has {} classes.",
                                 m_szModelTask,
                                 m_vClassLabels.size());
 
                        // Set flag saying we are ready for inference.
                        m_bReady = true;
                    }
                    catch (const c10::Error& trError)
                    {
                        LOG_ERROR(logging::g_qSharedLogger, "Error loading model: {}", trError.what());
                    }
                }
 
 
                ~PyTorchInterpreter()
                {
                    // Nothing to destroy.
                }
 
 
                std::vector<Detection> Inference(const cv::Mat& cvInputFrame, const float fMinObjectConfidence = 0.85, const float fNMSThreshold = 0.6)
                {
                    // Force single-threaded execution (if acceptable for your workload)
                    torch::set_num_threads(1);
                    // Create instance variables.
                    std::vector<Detection> vObjects;
 
                    // Preprocess the given image and pack int into an image.
                    torch::Tensor trTensorImage = PreprocessImage(cvInputFrame, m_trDevice);
 
                    // Perform inference.
                    std::vector<torch::jit::IValue> vInputs;
                    vInputs.push_back(trTensorImage);
                    torch::Tensor trOutputTensor;
                    try
                    {
                        trOutputTensor = m_trModel.forward(vInputs).toTensor();
                    }
                    catch (const c10::Error& trError)
                    {
                        LOG_ERROR(logging::g_qSharedLogger, "Error running inference: {}", trError.what());
                        return vObjects;
                    }
 
                    // Calculate the general stride sizes for YOLO based on input tensor shape.
                    int nImgSize  = m_cvModelInputSize.height;
                    int nP3Stride = std::pow((nImgSize / 8), 2);
                    int nP4Stride = std::pow((nImgSize / 16), 2);
                    int nP5Stride = std::pow((nImgSize / 32), 2);
                    // Calculate the proper prediction length for different YOLO versions.
                    int nYOLOv5AnchorsPerGridPoint = 3;
                    int nYOLOv8AnchorsPerGridPoint = 1;
                    int nYOLOv5TotalPredictionLength =
                        (nP3Stride * nYOLOv5AnchorsPerGridPoint) + (nP4Stride * nYOLOv5AnchorsPerGridPoint) + (nP5Stride * nYOLOv5AnchorsPerGridPoint);
                    int nYOLOv8TotalPredictionLength =
                        (nP3Stride * nYOLOv8AnchorsPerGridPoint) + (nP4Stride * nYOLOv8AnchorsPerGridPoint) + (nP5Stride * nYOLOv8AnchorsPerGridPoint);
 
                    // Parse the output tensor.
                    std::vector<int> vClassIDs;
                    std::vector<std::string> vClassLabels;
                    std::vector<float> vClassConfidences;
                    std::vector<cv::Rect> vBoundingBoxes;
 
                    // Get the largest dimension of our output tensor.
                    int nLargestDimension = *std::max_element(trOutputTensor.sizes().begin(), trOutputTensor.sizes().end());
                    // Check if the output tensor is YOLOv5 format.
                    if (nLargestDimension == nYOLOv5TotalPredictionLength)
                    {
                        // Parse inferenced output from tensor.
                        this->ParseTensorOutputYOLOv5(trOutputTensor, vClassIDs, vClassConfidences, vBoundingBoxes, cvInputFrame.size(), fMinObjectConfidence);
                    }
                    // Check if the output tensor is YOLOv8 format.
                    else if (nLargestDimension == nYOLOv8TotalPredictionLength)
                    {
                        // Parse inferenced output from tensor.
                        this->ParseTensorOutputYOLOv8(trOutputTensor, vClassIDs, vClassConfidences, vBoundingBoxes, cvInputFrame.size(), fMinObjectConfidence);
                    }
 
                    // Perform NMS to filter out bad/duplicate detections.
                    NonMaxSuppression(vObjects, vClassIDs, vClassConfidences, vBoundingBoxes, fMinObjectConfidence, fNMSThreshold);
 
                    // Loop through the final detections and set the class names for each detection based on the class ID.
                    for (size_t nIter = 0; nIter < vObjects.size(); ++nIter)
                    {
                        // Check if the class ID is valid.
                        if (vClassIDs[nIter] >= 0 && vClassIDs[nIter] < static_cast<int>(m_vClassLabels.size()))
                        {
                            vObjects[nIter].szClassName = m_vClassLabels[vClassIDs[nIter]];
                        }
                        else
                        {
                            vObjects[nIter].szClassName = "UnknownClass";
                        }
                    }
 
                    return vObjects;
                }
 
 
                bool IsReadyForInference() const { return m_bReady; }
 
            private:
                // Declare private methods.
 
 
                torch::Tensor PreprocessImage(const cv::Mat& cvInputFrame, const torch::Device& trDevice)
                {
                    // Resize the input image to match model and normalize it to 0-1.
                    cv::Mat cvResizedImage;
                    cv::resize(cvInputFrame, cvResizedImage, cv::Size(m_cvModelInputSize.width, m_cvModelInputSize.height), cv::INTER_LINEAR);
                    cvResizedImage.convertTo(cvResizedImage, CV_32FC3, 1.0 / 255.0);
 
                    // Convert OpenCV mat to a tensor.
                    torch::Tensor trTensorImage = torch::from_blob(cvResizedImage.data, {1, cvResizedImage.rows, cvResizedImage.cols, 3}, torch::kFloat);
                    trTensorImage               = trTensorImage.permute({0, 3, 1, 2});    // Convert to CxHxW format.
                    trTensorImage               = trTensorImage.to(trDevice);             // Move tensor to the specified hardware device.
 
                    return trTensorImage;
                }
 
 
                void ParseTensorOutputYOLOv5(const torch::Tensor& trOutput,
                                             std::vector<int>& vClassIDs,
                                             std::vector<float>& vClassConfidences,
                                             std::vector<cv::Rect>& vBoundingBoxes,
                                             const cv::Size& cvInputFrameSize,
                                             const float fMinObjectConfidence)
                {
                    /*
                     * For YOLOv5, you divide your image size, i.e. 640 by the P3, P4, P5 output strides of 8, 16, 32 to arrive at grid sizes
                     * of 80x80, 40x40, 20x20. Each grid point has 3 anchors by default (anchor box values: small, medium, large), and each anchor contains a vector 5 +
                     * nc long, where nc is the number of classes the model has. So for a 640 image, the output tensor will be [1, 25200, 85]
                     */
                    // Squeeze the batch dimension from the output tensor.
                    torch::Tensor trSqueezedOutput = trOutput.squeeze(0);
 
                    // Move the tensor to CPU if necessary. If we're using GPU and we don't move the tensor to CPU, we will get an error and it will be slow.
                    if (trSqueezedOutput.device().is_cuda())
                    {
                        trSqueezedOutput = trSqueezedOutput.to(torch::kCPU);
                    }
                    // Convert tensor to float if necessary.
                    if (trSqueezedOutput.scalar_type() != torch::kFloat32)
                    {
                        trSqueezedOutput = trSqueezedOutput.to(torch::kFloat32);
                    }
                    // Ensure tensor is contiguous in memory.
                    if (!trSqueezedOutput.is_contiguous())
                    {
                        trSqueezedOutput = trSqueezedOutput.contiguous();
                    }
 
                    // Create an accessor for fast element-wise access.
                    at::TensorAccessor trAccessor = trSqueezedOutput.accessor<float, 2>();
                    const int nNumDetections      = trSqueezedOutput.size(0);
                    const int nTotalValues        = trSqueezedOutput.size(1);    // equals 5 + number_of_classes
 
                    // Loop through each detection.
                    for (int i = 0; i < nNumDetections; i++)
                    {
                        // Get the objectness confidence. This is the 5th value for each grid/anchor prediction. (4th index)
                        float fObjectnessConfidence = trAccessor[i][4];
 
                        // Check if the object confidence is greater than or equal to the threshold.
                        if (fObjectnessConfidence < fMinObjectConfidence)
                        {
                            continue;
                        }
 
                        // Retrieve bounding box data.
                        float fCenterX = trAccessor[i][0];
                        float fCenterY = trAccessor[i][1];
                        float fWidth   = trAccessor[i][2];
                        float fHeight  = trAccessor[i][3];
 
                        // Scale bounding box to original image size.
                        int nLeft           = static_cast<int>((fCenterX - (0.5 * fWidth)) * cvInputFrameSize.width);
                        int nTop            = static_cast<int>((fCenterY - (0.5 * fHeight)) * cvInputFrameSize.height);
                        int nBoundingWidth  = static_cast<int>(fWidth * cvInputFrameSize.width);
                        int nBoundingHeight = static_cast<int>(fHeight * cvInputFrameSize.height);
 
                        // Repackaged bounding box data to be more readable.
                        cv::Rect cvBoundingBox(nLeft, nTop, nBoundingWidth, nBoundingHeight);
 
                        // Loop over class confidence values and find the class ID with the highest confidence.
                        float fClassConfidence = -1.0f;
                        int nClassID           = -1;
                        for (int j = 5; j < nTotalValues; j++)
                        {
                            float fConfidence = trAccessor[i][j];
                            if (fConfidence > fClassConfidence)
                            {
                                fClassConfidence = fConfidence;
                                nClassID         = j - 5;
                            }
                        }
 
                        // Only process detections that meet the minimum confidence.
                        if (fClassConfidence < fMinObjectConfidence)
                        {
                            continue;
                        }
 
                        // Add data to vectors.
                        vClassIDs.emplace_back(nClassID);
                        vClassConfidences.emplace_back(fClassConfidence);
                        vBoundingBoxes.emplace_back(cvBoundingBox);
                    }
                }
 
 
                void ParseTensorOutputYOLOv8(const torch::Tensor& trOutput,
                                             std::vector<int>& vClassIDs,
                                             std::vector<float>& vClassConfidences,
                                             std::vector<cv::Rect>& vBoundingBoxes,
                                             const cv::Size& cvInputFrameSize,
                                             const float fMinObjectConfidence)
                {
                    /*
                     * Permute the output tensor shape to match the expected format of the model. If the model is YOLOv8, the output
                     * shape for a 640x640 image will be [1, 4 + nc, 8400] (nc = number of classes). Notice how the larger dimensions is swapped
                     * when compared to YOLOv5. We will permute the tensor to [1, 8400, 4 + nc] to make it easier to parse. Then squeeze the
                     * tensor to remove the batch dimension so the final shape will be [8400, 4 + nc]. Thanks pytorch for being cool with the
                     * permute function.
                     */
                    // Permute the tensor shape from [1, 4 + nc, 8400] to [1, 8400, 4 + nc]
                    // and then squeeze to remove the batch dimension, resulting in [8400, 4 + nc]
                    torch::Tensor trPermuteOutput = trOutput.permute({0, 2, 1}).squeeze(0);
 
                    // Move tensor to CPU if necessary. If we're using GPU and we don't move the tensor to CPU, we will get an error and it will be slow.
                    if (trPermuteOutput.device().is_cuda())
                    {
                        trPermuteOutput = trPermuteOutput.to(torch::kCPU);
                    }
                    // Convert tensor to float if necessary.
                    if (trPermuteOutput.scalar_type() != torch::kFloat32)
                    {
                        trPermuteOutput = trPermuteOutput.to(torch::kFloat32);
                    }
                    // Ensure tensor is contiguous in memory.
                    if (!trPermuteOutput.is_contiguous())
                    {
                        trPermuteOutput = trPermuteOutput.contiguous();
                    }
 
                    // Create an accessor for fast element-wise access.
                    at::TensorAccessor trAccessor = trPermuteOutput.accessor<float, 2>();
                    const int nNumDetections      = trPermuteOutput.size(0);
                    const int nTotalValues        = trPermuteOutput.size(1);    // equals 4 + number_of_classes
 
                    // Loop through each detection.
                    for (int i = 0; i < nNumDetections; i++)
                    {
                        float fClassConfidence = -1.0f;
                        int nClassID           = -1;
 
                        // Loop over class confidence values.
                        for (int j = 4; j < nTotalValues; j++)
                        {
                            float fConfidence = trAccessor[i][j];
                            if (fConfidence > fClassConfidence)
                            {
                                fClassConfidence = fConfidence;
                                nClassID         = j - 4;
                            }
                        }
 
                        // Only process detections that meet the minimum confidence.
                        if (fClassConfidence < fMinObjectConfidence)
                        {
                            continue;
                        }
 
                        // Retrieve bounding box data.
                        float fCenterX = trAccessor[i][0];
                        float fCenterY = trAccessor[i][1];
                        float fWidth   = trAccessor[i][2];
                        float fHeight  = trAccessor[i][3];
 
                        // Scale bounding box to original image size.
                        int nLeft      = static_cast<int>(fCenterX * cvInputFrameSize.width / 640.0f - (0.5f * fWidth * cvInputFrameSize.width / 640.0f));
                        int nTop       = static_cast<int>(fCenterY * cvInputFrameSize.height / 640.0f - (0.5f * fHeight * cvInputFrameSize.height / 640.0f));
                        int nBoxWidth  = static_cast<int>(fWidth * cvInputFrameSize.width / 640.0f);
                        int nBoxHeight = static_cast<int>(fHeight * cvInputFrameSize.height / 640.0f);
                        cv::Rect cvBoundingBox(nLeft, nTop, nBoxWidth, nBoxHeight);
 
                        // Append results.
                        vClassIDs.push_back(nClassID);
                        vClassConfidences.push_back(fClassConfidence);
                        vBoundingBoxes.push_back(cvBoundingBox);
                    }
                }
 
                // Declare private member variables.
                torch::jit::script::Module m_trModel;
                torch::Device m_trDevice = torch::kCPU;
                std::string m_szModelPath;
                bool m_bReady;
                std::string m_szModelTask;
                cv::Size m_cvModelInputSize;
                std::vector<std::string> m_vClassLabels;
        };
    }    // namespace pytorch
}    // namespace yolomodel
 
#endif