| #ifndef YOLO_INFERENCE_H_ |
| #define YOLO_INFERENCE_H_ |
|
|
| #include <string> |
| #include <vector> |
| #include <opencv2/imgproc.hpp> |
| #include <openvino/openvino.hpp> |
|
|
| namespace yolo { |
|
|
| struct Detection { |
| short class_id; |
| float confidence; |
| cv::Rect box; |
| }; |
|
|
| class Inference { |
| public: |
| Inference() {} |
| |
| Inference(const std::string &model_path, const float &model_confidence_threshold, const float &model_NMS_threshold); |
| |
| Inference(const std::string &model_path, const cv::Size model_input_shape, const float &model_confidence_threshold, const float &model_NMS_threshold); |
|
|
| void RunInference(cv::Mat &frame); |
|
|
| private: |
| void InitializeModel(const std::string &model_path); |
| void Preprocessing(const cv::Mat &frame); |
| void PostProcessing(cv::Mat &frame); |
| cv::Rect GetBoundingBox(const cv::Rect &src) const; |
| void DrawDetectedObject(cv::Mat &frame, const Detection &detections) const; |
|
|
| cv::Point2f scale_factor_; |
| cv::Size2f model_input_shape_; |
| cv::Size model_output_shape_; |
|
|
| ov::InferRequest inference_request_; |
| ov::CompiledModel compiled_model_; |
|
|
| float model_confidence_threshold_; |
| float model_NMS_threshold_; |
|
|
| 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" |
| }; |
| }; |
|
|
| } |
|
|
| #endif |
|
|