80 lines
2.8 KiB
C
80 lines
2.8 KiB
C
|
//
|
||
|
|
// Created by DefTruth on 2021/8/7.
|
||
|
|
//
|
||
|
|
|
||
|
|
#ifndef LITE_AI_ORT_CV_YOLOR_H
|
||
|
|
#define LITE_AI_ORT_CV_YOLOR_H
|
||
|
|
|
||
|
|
#include "lite/ort/core/ort_core.h"
|
||
|
|
|
||
|
|
namespace ortcv
|
||
|
|
{
|
||
|
|
class LITE_EXPORTS YoloR : public BasicOrtHandler
|
||
|
|
{
|
||
|
|
public:
|
||
|
|
explicit YoloR(const std::string &_onnx_path, unsigned int _num_threads = 1) :
|
||
|
|
BasicOrtHandler(_onnx_path, _num_threads)
|
||
|
|
{};
|
||
|
|
|
||
|
|
~YoloR() override = default;
|
||
|
|
|
||
|
|
private:
|
||
|
|
// nested classes
|
||
|
|
typedef struct
|
||
|
|
{
|
||
|
|
float r;
|
||
|
|
int dw;
|
||
|
|
int dh;
|
||
|
|
int new_unpad_w;
|
||
|
|
int new_unpad_h;
|
||
|
|
bool flag;
|
||
|
|
} YoloRScaleParams;
|
||
|
|
|
||
|
|
private:
|
||
|
|
static constexpr const float mean_val = 0.f;
|
||
|
|
static constexpr const float scale_val = 1.0 / 255.f;
|
||
|
|
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"
|
||
|
|
};
|
||
|
|
enum NMS
|
||
|
|
{
|
||
|
|
HARD = 0, BLEND = 1, OFFSET = 2
|
||
|
|
};
|
||
|
|
static constexpr const unsigned int max_nms = 30000;
|
||
|
|
|
||
|
|
private:
|
||
|
|
Ort::Value transform(const cv::Mat &mat_rs) override; // without resize
|
||
|
|
|
||
|
|
void resize_unscale(const cv::Mat &mat,
|
||
|
|
cv::Mat &mat_rs,
|
||
|
|
int target_height,
|
||
|
|
int target_width,
|
||
|
|
YoloRScaleParams &scale_params);
|
||
|
|
|
||
|
|
void generate_bboxes(const YoloRScaleParams &scale_params,
|
||
|
|
std::vector<types::Boxf> &bbox_collection,
|
||
|
|
std::vector<Ort::Value> &output_tensors,
|
||
|
|
float score_threshold, float img_height,
|
||
|
|
float img_width); // rescale & exclude
|
||
|
|
|
||
|
|
void nms(std::vector<types::Boxf> &input, std::vector<types::Boxf> &output,
|
||
|
|
float iou_threshold, unsigned int topk, unsigned int nms_type);
|
||
|
|
|
||
|
|
public:
|
||
|
|
void detect(const cv::Mat &mat, std::vector<types::Boxf> &detected_boxes,
|
||
|
|
float score_threshold = 0.25f, float iou_threshold = 0.45f,
|
||
|
|
unsigned int topk = 100, unsigned int nms_type = NMS::OFFSET);
|
||
|
|
|
||
|
|
};
|
||
|
|
}
|
||
|
|
|
||
|
|
#endif //LITE_AI_ORT_CV_YOLOR_H
|