azaion/autopilot
1
0
Fork 0
mirror of https://github.com/azaion/autopilot.git synced 2026-04-22 22:46:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fully working opi_rtsp on PC with YOLOv8 ONNX models

Browse Source
This commit is contained in:
Tuomas Järvinen
2024-07-10 18:37:33 +02:00
parent 896307d296
commit 3d39d8fd99
5 changed files with 272 additions and 220 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tmp/opi_rtsp/src-onnx/inference.h
+36 -29
View File
@@ -1,44 +1,51 @@
#ifndef INFERENCE_H
#define INFERENCE_H
#pragma once
#include <onnxruntime_cxx_api.h>
#include <opencv2/opencv.hpp>
// Cpp native
#include <fstream>
#include <vector>
#include <string>
#include <random>
// OpenCV / DNN / Inference
#include <opencv2/imgproc.hpp>
#include <opencv2/opencv.hpp>
#include <opencv2/dnn.hpp>
struct Detection
{
float confidence;
cv::Rect bbox;
int class_id;
std::string class_name;
int class_id{0};
std::string className{};
float confidence{0.0};
cv::Scalar color{};
cv::Rect box{};
};
class InferenceEngine
class Inference
{
public:
InferenceEngine(const std::string &model_path);
~InferenceEngine();
Inference(const std::string &onnxModelPath, const cv::Size &modelInputShape = {640, 640}, const std::string &classesTxtFile = "", const bool &runWithCuda = false);
std::vector<Detection> runInference(const cv::Mat &input);
cv::Mat drawLabels(const cv::Mat &image, const std::vector<Detection> &detections);
std::vector<float> preprocessImage(const cv::Mat &image);
std::vector<Detection> filterDetections(const std::vector<float> &results, float confidence_threshold, int img_width, int img_height, int orig_width, int orig_height);
std::vector<float> runInference(const std::vector<float> &input_tensor_values);
cv::Mat draw_labels(const cv::Mat &image, const std::vector<Detection> &detections);
std::vector<int64_t> input_shape;
private:
Ort::Env env;
Ort::SessionOptions session_options;
Ort::Session session;
void loadClassesFromFile();
void loadOnnxNetwork();
std::string getInputName();
std::string getOutputName();
cv::Mat formatToSquare(const cv::Mat &source);
static const std::vector<std::string> CLASS_NAMES;
std::string modelPath{};
std::string classesPath{};
bool cudaEnabled{};
std::vector<std::string> classes{"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush"};
cv::Size2f modelShape{};
float modelConfidenceThreshold {0.25};
float modelScoreThreshold {0.45};
float modelNMSThreshold {0.50};
bool letterBoxForSquare = false;
cv::dnn::Net net;
};
#endif // INFERENCE_H