-
Notifications
You must be signed in to change notification settings - Fork 199
Description
#include"../utils/yolo.h"
#include"yolov8.h"
#include <Windows.h>
void setParameters(utils::InitParameter& initParameters)
{
initParameters.class_names = utils::dataSets::coco80;
//initParameters.class_names = utils::dataSets::voc20;
initParameters.num_class = 2; // for coco
//initParameters.num_class = 20; // for voc2012
initParameters.batch_size = 1;
initParameters.dst_h = 640;
initParameters.dst_w = 640;
initParameters.input_output_names = { "images", "output0" };
initParameters.conf_thresh = 0.25f;
initParameters.iou_thresh = 0.45f;
initParameters.save_path = "";
}
void task(YOLOV8& yolo, const utils::InitParameter& param, std::vectorcv::Mat& imgsBatch, const int& delayTime, const int& batchi,
const bool& isShow, const bool& isSave)
{
utils::DeviceTimer d_t0; yolo.copy(imgsBatch); float t0 = d_t0.getUsedTime();
utils::DeviceTimer d_t1; yolo.preprocess(imgsBatch); float t1 = d_t1.getUsedTime();
utils::DeviceTimer d_t2; yolo.infer(); float t2 = d_t2.getUsedTime();
utils::DeviceTimer d_t3; yolo.postprocess(imgsBatch); float t3 = d_t3.getUsedTime();
sample::gLogInfo <<
//"copy time = " << t0 / param.batch_size << "; "
"preprocess time = " << t1 / param.batch_size << "; "
"infer time = " << t2 / param.batch_size << "; "
"postprocess time = " << t3 / param.batch_size << std::endl;
if (isShow)
utils::show(yolo.getObjectss(), param.class_names, delayTime, imgsBatch);
if (isSave)
utils::save(yolo.getObjectss(), param.class_names, param.save_path, imgsBatch, param.batch_size, batchi);
yolo.reset();
}
int main(int argc, char** argv)
{
// parameters
utils::InitParameter param;
setParameters(param);
// path
//std::string model_path = "E:/app/TensorRT/TensorRT-8.6.1.6/bin/pubgv8p2.trt";
std::string model_path = "C:/Users/10420/pubgv8.trt";
std::string image_path = "C:/Users/10420/Pictures/0SbOi7H9.jpg";
cv::Mat img = cv::imread(image_path);
// camera' id
int camera_id = 0;
// get input
utils::InputStream source;
source = utils::InputStream::IMAGE;
//source = utils::InputStream::VIDEO;
//source = utils::InputStream::CAMERA;
// update params from command line parser
bool is_show = true;
bool is_save = false;
param.batch_size = 1;
int total_batches = 1;
int delay_time = 0;
param.src_h = img.rows;
param.src_w = img.cols;
YOLOV8 yolo(param);
// read model
std::vector<unsigned char> trt_file = utils::loadModel(model_path);
if (trt_file.empty())
{
sample::gLogError << "trt_file is empty!" << std::endl;
return -1;
}
// init model
if (!yolo.init(trt_file))
{
sample::gLogError << "initEngine() ocur errors!" << std::endl;
return -1;
}
yolo.check();
std::vector<cv::Mat> imgs_batch;
imgs_batch.reserve(param.batch_size);
sample::gLogInfo << imgs_batch.capacity() << std::endl;
int batchi = 0;
imgs_batch.emplace_back(img);
task(yolo, param, imgs_batch, delay_time, batchi, false, false);
Sleep(3000);
for (size_t i = 0; ; i++)
{
task(yolo, param, imgs_batch, delay_time, batchi, false, false);
}
return -1;
}