55 lines
1.2 KiB
C++
55 lines
1.2 KiB
C++
#pragma once
|
|
|
|
#include <opencv2/opencv.hpp>
|
|
#include <onnxruntime_cxx_api.h>
|
|
#include <vector>
|
|
#include <string>
|
|
#include <memory>
|
|
#include <stdexcept>
|
|
|
|
|
|
struct Detection
|
|
{
|
|
int class_id;
|
|
float score;
|
|
int x;
|
|
int y;
|
|
int width;
|
|
int height;
|
|
};
|
|
|
|
|
|
cv::Mat preprocess(const cv::Mat& img, int target_width, int target_height, int& pad_w, int& pad_h, float& scale);
|
|
|
|
std::vector<Detection> postprocess(Ort::Value& output_tensor, float scale, int pad_w, int pad_h, int img_w, int img_h, float conf_threshold, float iou_threshold);
|
|
|
|
|
|
class YoloDetector {
|
|
public:
|
|
|
|
Ort::Env env;
|
|
std::unique_ptr<Ort::Session> session;
|
|
|
|
|
|
int input_width = 0;
|
|
int input_height = 0;
|
|
|
|
|
|
Ort::AllocatorWithDefaultOptions allocator;
|
|
std::string input_name_str;
|
|
std::string output_name_str;
|
|
std::vector<const char*> input_node_names;
|
|
std::vector<const char*> output_node_names;
|
|
|
|
public:
|
|
/**
|
|
* */
|
|
YoloDetector(const ORTCHAR_T* model_path, int in_width, int in_height);
|
|
|
|
std::vector<Detection> detect(
|
|
unsigned char* image_bytes,
|
|
int image_width,
|
|
int image_height,
|
|
float conf_threshold,
|
|
float iou_threshold);
|
|
}; |