109 lines
5.6 KiB
C++
109 lines
5.6 KiB
C++
#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
|