84 lines
3.0 KiB
C
84 lines
3.0 KiB
C
|
//
|
||
|
|
// Created by DefTruth on 2021/5/23.
|
||
|
|
//
|
||
|
|
|
||
|
|
#ifndef LITE_AI_ORT_CV_YOLOV3_H
|
||
|
|
#define LITE_AI_ORT_CV_YOLOV3_H
|
||
|
|
|
||
|
|
#include "lite/ort/core/ort_core.h"
|
||
|
|
|
||
|
|
namespace ortcv
|
||
|
|
{
|
||
|
|
class LITE_EXPORTS YoloV3
|
||
|
|
{
|
||
|
|
private:
|
||
|
|
Ort::Env ort_env;
|
||
|
|
Ort::Session *ort_session = nullptr;
|
||
|
|
std::vector<const char *> input_node_names;
|
||
|
|
std::vector<std::vector<int64_t>> input_node_dims; // >=1 inputs.
|
||
|
|
std::vector<size_t> input_tensor_sizes;
|
||
|
|
Ort::MemoryInfo memory_info_handler = Ort::MemoryInfo::CreateCpu(
|
||
|
|
OrtArenaAllocator, OrtMemTypeDefault);
|
||
|
|
std::vector<const char *> output_node_names;
|
||
|
|
const LITEORT_CHAR *onnx_path = nullptr;
|
||
|
|
const char *log_id = nullptr;
|
||
|
|
unsigned int num_outputs = 1;
|
||
|
|
unsigned int num_inputs = 1;
|
||
|
|
std::vector<float> input_1_values_handler;
|
||
|
|
std::vector<float> image_shape_values_handler;
|
||
|
|
|
||
|
|
protected:
|
||
|
|
const unsigned int num_threads; // initialize at runtime.
|
||
|
|
|
||
|
|
public:
|
||
|
|
// yolov3 is an multi-inputs & multi-outputs & dynamic shape
|
||
|
|
// (dynamic input shape: batch,input_height,input_width)
|
||
|
|
// & (dynamic output shape: bactch, num_anchors, num_selected)
|
||
|
|
explicit YoloV3(const std::string &_onnx_path, unsigned int _num_threads = 1);
|
||
|
|
|
||
|
|
~YoloV3();
|
||
|
|
|
||
|
|
protected:
|
||
|
|
YoloV3(const YoloV3 &) = delete;
|
||
|
|
|
||
|
|
YoloV3(YoloV3 &&) = delete;
|
||
|
|
|
||
|
|
YoloV3 &operator=(const YoloV3 &) = delete;
|
||
|
|
|
||
|
|
YoloV3 &operator=(YoloV3 &&) = delete;
|
||
|
|
|
||
|
|
private:
|
||
|
|
static constexpr const float mean_val = 0.f;
|
||
|
|
static constexpr const float scale_val = 1.0 / 255.f;
|
||
|
|
static constexpr const unsigned int input_height = 416;
|
||
|
|
static constexpr const unsigned int input_width = 416;
|
||
|
|
static constexpr const unsigned int batch_size = 1;
|
||
|
|
|
||
|
|
const char *class_names[80] = {
|
||
|
|
"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"
|
||
|
|
};
|
||
|
|
|
||
|
|
private:
|
||
|
|
std::vector<Ort::Value> transform(const std::vector<cv::Mat> &mats);
|
||
|
|
|
||
|
|
void generate_bboxes(std::vector<types::Boxf> &bbox_collection,
|
||
|
|
std::vector<Ort::Value> &output_tensors); // rescale & exclude
|
||
|
|
|
||
|
|
void print_debug_string();
|
||
|
|
|
||
|
|
public:
|
||
|
|
void detect(const cv::Mat &mat, std::vector<types::Boxf> &detected_boxes);
|
||
|
|
|
||
|
|
};
|
||
|
|
}
|
||
|
|
|
||
|
|
#endif //LITE_AI_ORT_CV_YOLOV3_H
|