89 lines
4.8 KiB
C
89 lines
4.8 KiB
C
|
#ifndef ANSYOLO12OD_H
|
||
|
|
#define ANSYOL12OOD_H
|
||
|
|
#pragma once
|
||
|
|
#include "ANSEngineCommon.h"
|
||
|
|
#include <onnxruntime_cxx_api.h>
|
||
|
|
#include <algorithm>
|
||
|
|
#include <fstream>
|
||
|
|
#include <iostream>
|
||
|
|
#include <numeric>
|
||
|
|
#include <string>
|
||
|
|
#include <vector>
|
||
|
|
#include <memory>
|
||
|
|
#include <chrono>
|
||
|
|
#include <random>
|
||
|
|
#include <unordered_map>
|
||
|
|
#include <thread>
|
||
|
|
namespace ANSCENTER {
|
||
|
|
// Standard Yolo engine class: yolo12
|
||
|
|
class ANSENGINE_API YOLO12OD :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();
|
||
|
|
~YOLO12OD();
|
||
|
|
private:
|
||
|
|
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
|
||
|
|
std::string _modelFilePath;
|
||
|
|
|
||
|
|
// 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
|
||
|
|
float m_imgWidth = 0;
|
||
|
|
float m_imgHeight = 0;
|
||
|
|
protected:
|
||
|
|
bool loadModel(const std::string& modelPath, bool useGPU = true);
|
||
|
|
std::vector<Object> detect(const cv::Mat& image, float confThreshold = 0.4f, float iouThreshold = 0.45f);
|
||
|
|
//cv::Mat preprocess(const cv::Mat& image, float*& blob, std::vector<int64_t>& inputTensorShape);
|
||
|
|
cv::Mat preprocess(const cv::Mat& image, std::vector<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,
|
||
|
|
float confThreshold, float iouThreshold);
|
||
|
|
|
||
|
|
private:
|
||
|
|
template <typename T>
|
||
|
|
typename std::enable_if<std::is_arithmetic<T>::value, T>::type
|
||
|
|
inline clamp(const T& value, const T& low, const T& high)
|
||
|
|
{
|
||
|
|
// Ensure the range [low, high] is valid; swap if necessary
|
||
|
|
T validLow = low < high ? low : high;
|
||
|
|
T validHigh = low < high ? high : low;
|
||
|
|
|
||
|
|
// Clamp the value to the range [validLow, validHigh]
|
||
|
|
if (value < validLow)
|
||
|
|
return validLow;
|
||
|
|
if (value > validHigh)
|
||
|
|
return validHigh;
|
||
|
|
return value;
|
||
|
|
}
|
||
|
|
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);
|
||
|
|
BoundingBox scaleCoords(const cv::Size& imageShape, BoundingBox coords,
|
||
|
|
const cv::Size& imageOriginalShape, bool p_Clip);
|
||
|
|
void NMSBoxes(const std::vector<BoundingBox>& boundingBoxes,
|
||
|
|
const std::vector<float>& scores,
|
||
|
|
float scoreThreshold,
|
||
|
|
float nmsThreshold,
|
||
|
|
std::vector<int>& indices);
|
||
|
|
};
|
||
|
|
}
|
||
|
|
#endif
|