ANSCORE/modules/ANSODEngine/ANSONNXOBB.h

#ifndef ANSONNXOBB_H
#define ANSONNXOBB_H
#pragma once
#include "ANSEngineCommon.h"
#include "engine.h"

namespace ANSCENTER {
    static constexpr float EPS = 1e-7f;
    struct OrientedBoundingBox {
        float x;       // x-coordinate of the center
        float y;       // y-coordinate of the center
        float width;   // width of the box
        float height;  // height of the box
        float angle;   // rotation angle in radians
        OrientedBoundingBox() : x(0), y(0), width(0), height(0), angle(0) {}
        OrientedBoundingBox(float x_, float y_, float width_, float height_, float angle_)
            : x(x_), y(y_), width(width_), height(height_), angle(angle_) {
        }
    };
    struct Detection {
        OrientedBoundingBox box;  // Oriented bounding box in xywhr format
        float conf{};             // Confidence score
        int classId{};            // Class ID
        Detection() = default;
        Detection(const OrientedBoundingBox& box_, float conf_, int classId_)
            : box(box_), conf(conf_), classId(classId_) {
        }
    };
    class ANSENGINE_API ANSONNXOBB :public ANSODBase {
    public:
        virtual bool Initialize(std::string licenseKey, ModelConfig modelConfig, const std::string& modelZipFilePath, const std::string& modelZipPassword, std::string& labelMap) override;
        virtual bool LoadModel(const std::string& modelZipFilePath, const std::string& modelZipPassword)override;
        virtual bool LoadModelFromFolder(std::string licenseKey, ModelConfig modelConfig, std::string modelName, std::string className, const std::string& modelFolder, std::string& labelMap)override;
        virtual bool OptimizeModel(bool fp16, std::string& optimizedModelFolder);
        std::vector<Object> RunInference(const cv::Mat& input);
        std::vector<Object> RunInference(const cv::Mat& input, const std::string& camera_id);
        bool Destroy();
        ~ANSONNXOBB();
        int getInstanceId() const { return instanceId_; }

    private:
        std::string     _modelFilePath;
        bool            _modelLoadValid;
        bool            _fp16{ false };
        size_t vectorProduct(const std::vector<int64_t>& vector);
        void letterBox(const cv::Mat& image, cv::Mat& outImage,
            const cv::Size& newShape,
            const cv::Scalar& color = cv::Scalar(114, 114, 114),
            bool auto_ = true,
            bool scaleFill = false,
            bool scaleUp = true,
            int stride = 32);
        void NMSBoxes(const std::vector<BoundingBox>& boundingBoxes,
            const std::vector<float>& scores,
            float scoreThreshold,
            float nmsThreshold,
            std::vector<int>& indices);
        void drawBoundingBox(cv::Mat& image, const std::vector<Detection>& detections,
            const std::vector<std::string>& classNames, const std::vector<cv::Scalar>& colors);
		std::vector<cv::Point2f> OBBToPoints(const OrientedBoundingBox& obb);

        std::vector<cv::Scalar> generateColors(const std::vector<std::string>& classNames, int seed = 42);
        void getCovarianceComponents(const OrientedBoundingBox& box, float& out1, float& out2, float& out3);
        std::vector<std::vector<float>> batchProbiou(const std::vector<OrientedBoundingBox>& obb1, const std::vector<OrientedBoundingBox>& obb2, float eps = EPS);
        std::vector<int> nmsRotatedImpl(const std::vector<OrientedBoundingBox>& sorted_boxes, float iou_thres);
        std::vector<int> nmsRotated(const std::vector<OrientedBoundingBox>& boxes, const std::vector<float>& scores, float threshold = 0.75f);
        std::vector<Object> nonMaxSuppression(
            const std::vector<Detection>& input_detections,
            const std::string& camera_id,
            float conf_thres = 0.25f,
            float iou_thres = 0.75f,
            int max_det = 1000);

        void warmupModel();
        bool Init(const std::string& modelPath, bool useGPU=true, int deviceId = 0);
        cv::Mat preprocess(const cv::Mat& image, float*& blob, std::vector<int64_t>& inputTensorShape);
        std::vector<Object> postprocess(
            const cv::Size& originalImageSize,
            const cv::Size& resizedImageShape,
            const std::vector<Ort::Value>& outputTensors, int topk,
            const std::string& camera_id);
        std::vector<Object> detect(const cv::Mat& image, const std::string& camera_id);
    private:
        static std::atomic<int> instanceCounter_;  // Thread-safe counter
        int instanceId_;
        int deviceId_ = 0;
        
        Ort::Env env{ nullptr };                         // ONNX Runtime environment
        Ort::SessionOptions sessionOptions{ nullptr };   // Session options for ONNX Runtime
        Ort::Session session{ nullptr };                 // ONNX Runtime session for running inference
        bool isDynamicInputShape{};                      // Flag indicating if input shape is dynamic
        cv::Size inputImageShape;                        // Expected input image shape for the model

        // Vectors to hold allocated input and output node names
        std::vector<Ort::AllocatedStringPtr> inputNodeNameAllocatedStrings;
        std::vector<const char*> inputNames;
        std::vector<Ort::AllocatedStringPtr> outputNodeNameAllocatedStrings;
        std::vector<const char*> outputNames;

        size_t numInputNodes, numOutputNodes;           // Number of input and output nodes in the model
        std::vector<std::string> classNames;            // Vector of class names loaded from file
        std::vector<cv::Scalar> classColors;            // Vector of colors for each class

      
    };
}
#endif
#pragma once