#ifndef CLASS_YOLO_DETECTOR_HPP_
|
#define CLASS_YOLO_DETECTOR_HPP_
|
|
#include <opencv2/opencv.hpp>
|
#include "ds_image.h"
|
#include "trt_utils.h"
|
#include "yolo.h"
|
#include "yolov2.h"
|
#include "yolov3.h"
|
#include "yolov4.h"
|
#include "yolov5.h"
|
|
#include <experimental/filesystem>
|
#include <fstream>
|
#include <string>
|
#include <chrono>
|
#include <stdio.h> /* defines FILENAME_MAX */
|
|
#include "class_detector.h"
|
#include "../utils/time_util.h"
|
#include "config.h"
|
|
class YoloDectector {
|
public:
|
YoloDectector() {
|
|
}
|
|
~YoloDectector() {
|
|
}
|
|
void init(const Config &config) {
|
_config = config;
|
// std::cout << "------------------------>config information start" << std::endl;
|
// std::cout << "net_type is: " << _config.net_type << std::endl;
|
// std::cout << "gpu_id is: " << _config.gpu_id << std::endl;
|
// std::cout << "detect_thresh is: " << _config.detect_thresh << std::endl;
|
// std::cout << "file_model_cfg is: " << _config.file_model_cfg << std::endl;
|
// std::cout << "file_model_weights is: " << _config.file_model_weights << std::endl;
|
// std::cout << "inference_precison is: " << _config.inference_precison << std::endl;
|
// std::cout << "calibration_image_list_file_txt is: " << _config.calibration_image_list_file_txt << std::endl;
|
// std::cout << "------------------------>config information end" << std::endl;
|
this->set_gpu_id(_config.gpu_id);
|
|
this->parse_config();
|
|
this->build_net();
|
}
|
|
void detect(const std::vector <cv::Mat> &vec_image, std::vector <BatchResult> &vec_batch_result) {
|
Timer timer;
|
std::vector <DsImage> vec_ds_images;
|
vec_batch_result.clear();
|
vec_batch_result.resize(vec_image.size());
|
for (const auto &img:vec_image) {
|
vec_ds_images.emplace_back(img, _vec_net_type[_config.net_type], _p_net->getInputH(), _p_net->getInputW());
|
}
|
cv::Mat trtInput = blobFromDsImages(vec_ds_images, _p_net->getInputH(), _p_net->getInputW());
|
timer.out("pre");
|
_p_net->doInference(trtInput.data, vec_ds_images.size());
|
timer.reset();
|
for (uint32_t i = 0; i < vec_ds_images.size(); ++i) {
|
auto curImage = vec_ds_images.at(i);
|
auto binfo = _p_net->decodeDetections(i, curImage.getImageHeight(), curImage.getImageWidth());
|
auto remaining = nmsAllClasses(_p_net->getNMSThresh(),
|
binfo,
|
_p_net->getNumClasses(),
|
_vec_net_type[_config.net_type]);
|
if (remaining.empty()) {
|
continue;
|
}
|
std::vector <Result> vec_result(0);
|
for (const auto &b : remaining) {
|
Result res;
|
res.id = b.label;
|
res.prob = b.prob;
|
const int x = b.box.x1;
|
const int y = b.box.y1;
|
const int w = b.box.x2 - b.box.x1;
|
const int h = b.box.y2 - b.box.y1;
|
res.rect = cv::Rect(x, y, w, h);
|
vec_result.push_back(res);
|
}
|
vec_batch_result[i] = vec_result;
|
}
|
timer.out("post");
|
}
|
|
private:
|
|
void set_gpu_id(const int id = 0) {
|
cudaError_t status = cudaSetDevice(id);
|
if (status != cudaSuccess) {
|
std::cout << "gpu id :" + std::to_string(id) + " not exist !" << std::endl;
|
assert(0);
|
}
|
}
|
|
void parse_config() {
|
_yolo_info.networkType = _vec_net_type[_config.net_type];
|
_yolo_info.configFilePath = m_staticStruct::model_cfg;
|
_yolo_info.wtsFilePath = m_staticStruct::model_wts;
|
_yolo_info.precision = _vec_precision[_config.inference_precison];
|
_yolo_info.deviceType = "kGPU";
|
auto npos = _yolo_info.wtsFilePath.find(".weights");
|
assert(npos != std::string::npos
|
&& "wts file file not recognised. File needs to be of '.weights' format");
|
_yolo_info.data_path = _yolo_info.wtsFilePath.substr(0, npos);
|
_yolo_info.calibrationTablePath = _yolo_info.data_path + "-calibration.table";
|
_yolo_info.inputBlobName = "data";
|
|
_infer_param.printPerfInfo = false;
|
_infer_param.printPredictionInfo = false;
|
_infer_param.calibImages = _config.calibration_image_list_file_txt;
|
_infer_param.calibImagesPath = "";
|
_infer_param.probThresh = _config.detect_thresh;
|
_infer_param.nmsThresh = 0.5;
|
std::cout << "------------------------>parse config complete" << std::endl;
|
}
|
|
void build_net() {
|
if (_config.net_type == YOLOV3) {
|
_p_net = std::unique_ptr < Yolo > {new YoloV3(_yolo_info, _infer_param)};
|
} else if (_config.net_type == YOLOV4) {
|
std::cout << "------------------------>build net start!" << std::endl;
|
_p_net = std::unique_ptr < Yolo > {new YoloV4(_yolo_info, _infer_param)};
|
std::cout << "------------------------>build net end" << std::endl;
|
} else if (_config.net_type == YOLOV5) {
|
_p_net = std::unique_ptr < Yolo > {new YoloV5(_yolo_info, _infer_param)};
|
} else {
|
assert(false && "Unrecognised network_type.");
|
}
|
}
|
|
private:
|
Config _config;
|
NetworkInfo _yolo_info;
|
InferParams _infer_param;
|
std::vector <std::string> _vec_net_type{"yolov3", "yolov4", "yolov5"};
|
std::vector <std::string> _vec_precision{"kINT8", "kHALF", "kFLOAT"};
|
std::unique_ptr <Yolo> _p_net = nullptr;
|
Timer _m_timer;
|
};
|
|
|
#endif
|
