#ifndef DETECTOR_H
|
#define DETECTOR_H
|
#include <iostream>
|
#include <opencv2/opencv.hpp>
|
#include "config.h"
|
#include <experimental/filesystem>
|
#include <fstream>
|
#include <string>
|
#include <chrono>
|
#include <stdio.h>
|
#include "model.h"
|
|
|
struct Result
|
{
|
int id = -1;
|
float prob = 0.f;
|
cv::Rect rect;
|
};
|
|
using BatchResult = std::vector<Result>;
|
|
enum ModelType
|
{
|
COMMON,
|
SMALL
|
};
|
|
enum Precision
|
{
|
INT8 = 0,
|
FP16,
|
FP32
|
};
|
|
struct Config
|
{
|
float detect_thresh = 0.1;
|
|
Precision inference_precison = FP32;
|
|
ModelType net_type = COMMON;
|
|
int gpu_id = 0;
|
};
|
|
class Detector
|
{
|
private:
|
Config _config;
|
NetworkInfo _info;
|
InferParams _infer_param;
|
std::vector<std::string> _vec_net_type{ "common","small"};
|
std::vector<std::string> _vec_precision{ "kINT8","kHALF","kFLOAT" };
|
std::unique_ptr<Detecter> _p_net = nullptr;
|
Timer _m_timer;
|
|
|
|
public:
|
Detector();
|
~Detector();
|
|
void init(const Config &config);
|
void detect(const std::vector<cv::Mat> &vec_image,
|
std::vector<BatchResult> &vec_batch_result);
|
|
private:
|
void set_gpu_id(const int id=0);
|
|
void parse_config();
|
|
void build_net();
|
|
// std::vector<BBoxInfo> decodeTensor(const int imageIdx, const int imageH, const int imageW, const TensorInfo& tensor) override;
|
};
|
|
#endif
|