Chiebot-Mirror
/
YOLOv8-TensorRT
mirror of https://github.com/triple-Mu/YOLOv8-TensorRT.git
9
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Format

Browse Source
pull/6/head
triple-Mu 2 years ago
parent 1f5cbaf779
commit c37f09006c
3 changed files with 482 additions and 0 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 84
      csrc/end2end/include/config.h
  2. 133
      csrc/end2end/include/utils.h
  3. 265
      csrc/end2end/include/yolov8.hpp

84
csrc/end2end/include/config.h
Unescape View File

@ -0,0 +1,84 @@
//
// Created by ubuntu on 1/10/23.
//
#ifndef YOLOV8_CSRC_END2END_INCLUDE_CONFIG_H
#define YOLOV8_CSRC_END2END_INCLUDE_CONFIG_H
#include "opencv2/opencv.hpp"
namespace det
{
const int DEVICE = 0;
static const int INPUT_W = 640;
static const int INPUT_H = 640;
static const int NUM_INPUT = 1;
static const int NUM_OUTPUT = 4;
static const int NUM_BINDINGS = NUM_INPUT + NUM_OUTPUT;
const cv::Scalar PAD_COLOR = { 114, 114, 114 };
const cv::Scalar RECT_COLOR = cv::Scalar(0, 0, 255);
const cv::Scalar TXT_COLOR = cv::Scalar(255, 255, 255);
const char* INPUT = "images";
const char* NUM_DETS = "num_dets";
const char* BBOXES = "bboxes";
const char* SCORES = "scores";
const char* LABELS = "labels";
const char* CLASS_NAMES[] = {
"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" };
const unsigned int COLORS[80][3] = {
{ 0, 114, 189 }, { 217, 83, 25 }, { 237, 177, 32 },
{ 126, 47, 142 }, { 119, 172, 48 }, { 77, 190, 238 },
{ 162, 20, 47 }, { 76, 76, 76 }, { 153, 153, 153 },
{ 255, 0, 0 }, { 255, 128, 0 }, { 191, 191, 0 },
{ 0, 255, 0 }, { 0, 0, 255 }, { 170, 0, 255 },
{ 85, 85, 0 }, { 85, 170, 0 }, { 85, 255, 0 },
{ 170, 85, 0 }, { 170, 170, 0 }, { 170, 255, 0 },
{ 255, 85, 0 }, { 255, 170, 0 }, { 255, 255, 0 },
{ 0, 85, 128 }, { 0, 170, 128 }, { 0, 255, 128 },
{ 85, 0, 128 }, { 85, 85, 128 }, { 85, 170, 128 },
{ 85, 255, 128 }, { 170, 0, 128 }, { 170, 85, 128 },
{ 170, 170, 128 }, { 170, 255, 128 }, { 255, 0, 128 },
{ 255, 85, 128 }, { 255, 170, 128 }, { 255, 255, 128 },
{ 0, 85, 255 }, { 0, 170, 255 }, { 0, 255, 255 },
{ 85, 0, 255 }, { 85, 85, 255 }, { 85, 170, 255 },
{ 85, 255, 255 }, { 170, 0, 255 }, { 170, 85, 255 },
{ 170, 170, 255 }, { 170, 255, 255 }, { 255, 0, 255 },
{ 255, 85, 255 }, { 255, 170, 255 }, { 85, 0, 0 },
{ 128, 0, 0 }, { 170, 0, 0 }, { 212, 0, 0 },
{ 255, 0, 0 }, { 0, 43, 0 }, { 0, 85, 0 },
{ 0, 128, 0 }, { 0, 170, 0 }, { 0, 212, 0 },
{ 0, 255, 0 }, { 0, 0, 43 }, { 0, 0, 85 },
{ 0, 0, 128 }, { 0, 0, 170 }, { 0, 0, 212 },
{ 0, 0, 255 }, { 0, 0, 0 }, { 36, 36, 36 },
{ 73, 73, 73 }, { 109, 109, 109 }, { 146, 146, 146 },
{ 182, 182, 182 }, { 219, 219, 219 }, { 0, 114, 189 },
{ 80, 183, 189 }, { 128, 128, 0 }
};
struct Object
{
cv::Rect_<float> rect;
int label = 0;
float prob = 0.0;
};
}
#endif //YOLOV8_CSRC_END2END_INCLUDE_CONFIG_H

133
csrc/end2end/include/utils.h
Unescape View File

@ -0,0 +1,133 @@
//
// Created by ubuntu on 1/10/23.
//
#ifndef YOLOV8_CSRC_END2END_INCLUDE_UTILS_H
#define YOLOV8_CSRC_END2END_INCLUDE_UTILS_H
#include <sys/stat.h>
#include <iostream>
#include <string>
#include <assert.h>
#include <unistd.h>
#include "NvInfer.h"
#define CHECK(call) \
do \
{ \
const cudaError_t error_code = call; \
if (error_code != cudaSuccess) \
{ \
printf("CUDA Error:\n"); \
printf(" File: %s\n", __FILE__); \
printf(" Line: %d\n", __LINE__); \
printf(" Error code: %d\n", error_code); \
printf(" Error text: %s\n", \
cudaGetErrorString(error_code)); \
exit(1); \
} \
} while (0)
class Logger : public nvinfer1::ILogger
{
public:
nvinfer1::ILogger::Severity reportableSeverity;
explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO) :
reportableSeverity(severity)
{
}
void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override
{
if (severity > reportableSeverity)
{
return;
}
switch (severity)
{
case nvinfer1::ILogger::Severity::kINTERNAL_ERROR:
std::cerr << "INTERNAL_ERROR: ";
break;
case nvinfer1::ILogger::Severity::kERROR:
std::cerr << "ERROR: ";
break;
case nvinfer1::ILogger::Severity::kWARNING:
std::cerr << "WARNING: ";
break;
case nvinfer1::ILogger::Severity::kINFO:
std::cerr << "INFO: ";
break;
default:
std::cerr << "VERBOSE: ";
break;
}
std::cerr << msg << std::endl;
}
};
inline int get_size_by_dims(const nvinfer1::Dims& dims)
{
int size = 1;
for (int i = 0; i < dims.nbDims; i++)
{
size *= dims.d[i];
}
return size;
}
inline int DataTypeToSize(const nvinfer1::DataType& dataType)
{
switch (dataType)
{
case nvinfer1::DataType::kFLOAT:
return sizeof(float);
case nvinfer1::DataType::kHALF:
return 2;
case nvinfer1::DataType::kINT8:
return sizeof(int8_t);
case nvinfer1::DataType::kINT32:
return sizeof(int32_t);
case nvinfer1::DataType::kBOOL:
return sizeof(bool);
default:
return sizeof(float);
}
}
inline float clamp(const float val, const float minVal = 0.f, const float maxVal = 1280.f)
{
assert(minVal <= maxVal);
return std::min(maxVal, std::max(minVal, val));
}
inline bool IsPathExist(const std::string& path)
{
if (access(path.c_str(), 0) == F_OK)
{
return true;
}
return false;
}
inline bool IsFile(const std::string& path)
{
if (!IsPathExist(path))
{
printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str());
return false;
}
struct stat buffer;
return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode));
}
inline bool IsFolder(const std::string& path)
{
if (!IsPathExist(path))
{
return false;
}
struct stat buffer;
return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
}
#endif //YOLOV8_CSRC_END2END_INCLUDE_UTILS_H

265
csrc/end2end/include/yolov8.hpp
Unescape View File

@ -0,0 +1,265 @@
//
// Created by ubuntu on 1/8/23.
//
#include "config.h"
#include "utils.h"
#include <fstream>
#include "NvInferPlugin.h"
using namespace det;
class YOLOv8
{
public:
explicit YOLOv8(const std::string& engine_file_path);
~YOLOv8();
void make_pipe(bool warmup = true);
void copy_from_Mat(const cv::Mat& image);
void infer();
void postprocess(std::vector<Object>& objs);
size_t in_size = 1 * 3 * INPUT_W * INPUT_H;
float w = INPUT_W;
float h = INPUT_H;
float ratio = 1.0f;
float dw = 0.f;
float dh = 0.f;
std::array<std::pair<int, int>, NUM_OUTPUT> out_sizes{};
std::array<void*, NUM_OUTPUT> outputs{};
private:
nvinfer1::ICudaEngine* engine = nullptr;
nvinfer1::IRuntime* runtime = nullptr;
nvinfer1::IExecutionContext* context = nullptr;
cudaStream_t stream = nullptr;
std::array<void*, NUM_BINDINGS> buffs{};
Logger gLogger{ nvinfer1::ILogger::Severity::kERROR };
};
YOLOv8::YOLOv8(const std::string& engine_file_path)
{
std::ifstream file(engine_file_path, std::ios::binary);
assert(file.good());
file.seekg(0, std::ios::end);
auto size = file.tellg();
std::ostringstream fmt;
file.seekg(0, std::ios::beg);
char* trtModelStream = new char[size];
assert(trtModelStream);
file.read(trtModelStream, size);
file.close();
initLibNvInferPlugins(&this->gLogger, "");
this->runtime = nvinfer1::createInferRuntime(this->gLogger);
assert(this->runtime != nullptr);
this->engine = this->runtime->deserializeCudaEngine(trtModelStream, size);
assert(this->engine != nullptr);
this->context = this->engine->createExecutionContext();
assert(this->context != nullptr);
cudaStreamCreate(&this->stream);
}
YOLOv8::~YOLOv8()
{
this->context->destroy();
this->engine->destroy();
this->runtime->destroy();
cudaStreamDestroy(this->stream);
for (auto& ptr : this->buffs)
{
CHECK(cudaFree(ptr));
}
for (auto& ptr : this->outputs)
{
CHECK(cudaFreeHost(ptr));
}
}
void YOLOv8::make_pipe(bool warmup)
{
const nvinfer1::Dims input_dims = this->engine->getBindingDimensions(
this->engine->getBindingIndex(INPUT)
);
this->in_size = get_size_by_dims(input_dims);
CHECK(cudaMalloc(&this->buffs[0], this->in_size * sizeof(float)));
this->context->setBindingDimensions(0, input_dims);
const int32_t num_dets_idx = this->engine->getBindingIndex(NUM_DETS);
const nvinfer1::Dims num_dets_dims = this->context->getBindingDimensions(num_dets_idx);
this->out_sizes[num_dets_idx - NUM_INPUT].first = get_size_by_dims(num_dets_dims);
this->out_sizes[num_dets_idx - NUM_INPUT].second = DataTypeToSize(
this->engine->getBindingDataType(num_dets_idx));
const int32_t bboxes_idx = this->engine->getBindingIndex(BBOXES);
const nvinfer1::Dims bboxes_dims = this->context->getBindingDimensions(bboxes_idx);
this->out_sizes[bboxes_idx - NUM_INPUT].first = get_size_by_dims(bboxes_dims);
this->out_sizes[bboxes_idx - NUM_INPUT].second = DataTypeToSize(
this->engine->getBindingDataType(bboxes_idx));
const int32_t scores_idx = this->engine->getBindingIndex(SCORES);
const nvinfer1::Dims scores_dims = this->context->getBindingDimensions(scores_idx);
this->out_sizes[scores_idx - NUM_INPUT].first = get_size_by_dims(scores_dims);
this->out_sizes[scores_idx - NUM_INPUT].second = DataTypeToSize(
this->engine->getBindingDataType(scores_idx));
const int32_t labels_idx = this->engine->getBindingIndex(LABELS);
const nvinfer1::Dims labels_dims = this->context->getBindingDimensions(labels_idx);
this->out_sizes[labels_idx - NUM_INPUT].first = get_size_by_dims(labels_dims);
this->out_sizes[labels_idx - NUM_INPUT].second = DataTypeToSize(
this->engine->getBindingDataType(labels_idx));
for (int i = 0; i < NUM_OUTPUT; i++)
{
const int osize = this->out_sizes[i].first * out_sizes[i].second;
CHECK(cudaHostAlloc(&this->outputs[i], osize, 0));
CHECK(cudaMalloc(&this->buffs[NUM_INPUT + i], osize));
}
if (warmup)
{
for (int i = 0; i < 10; i++)
{
size_t isize = this->in_size * sizeof(float);
auto* tmp = new float[isize];
CHECK(cudaMemcpyAsync(this->buffs[0],
tmp,
isize,
cudaMemcpyHostToDevice,
this->stream));
this->infer();
}
printf("model warmup 10 times\n");
}
}
void YOLOv8::copy_from_Mat(const cv::Mat& image)
{
float height = (float)image.rows;
float width = (float)image.cols;
float r = std::min(INPUT_H / height, INPUT_W / width);
int padw = (int)std::round(width * r);
int padh = (int)std::round(height * r);
cv::Mat tmp;
if ((int)width != padw || (int)height != padh)
{
cv::resize(image, tmp, cv::Size(padw, padh));
}
else
{
tmp = image.clone();
}
float _dw = INPUT_W - padw;
float _dh = INPUT_H - padh;
_dw /= 2.0f;
_dh /= 2.0f;
int top = int(std::round(_dh - 0.1f));
int bottom = int(std::round(_dh + 0.1f));
int left = int(std::round(_dw - 0.1f));
int right = int(std::round(_dw + 0.1f));
cv::copyMakeBorder(tmp, tmp, top, bottom, left, right, cv::BORDER_CONSTANT, PAD_COLOR);
cv::dnn::blobFromImage(tmp,
tmp,
1 / 255.f,
cv::Size(),
cv::Scalar(0, 0, 0),
true,
false,
CV_32F);
CHECK(cudaMemcpyAsync(this->buffs[0],
tmp.ptr<float>(),
this->in_size * sizeof(float),
cudaMemcpyHostToDevice,
this->stream));
this->ratio = 1 / r;
this->dw = _dw;
this->dh = _dh;
this->w = width;
this->h = height;
}
void YOLOv8::infer()
{
this->context->enqueueV2(buffs.data(), this->stream, nullptr);
for (int i = 0; i < NUM_OUTPUT; i++)
{
const int osize = this->out_sizes[i].first * out_sizes[i].second;
CHECK(cudaMemcpyAsync(this->outputs[i],
this->buffs[NUM_INPUT + i],
osize,
cudaMemcpyDeviceToHost,
this->stream));
}
cudaStreamSynchronize(this->stream);
}
void YOLOv8::postprocess(std::vector<Object>& objs)
{
int* num_dets = static_cast<int*>(this->outputs[0]);
auto* boxes = static_cast<float*>(this->outputs[1]);
auto* scores = static_cast<float*>(this->outputs[2]);
int* labels = static_cast<int*>(this->outputs[3]);
for (int i = 0; i < num_dets[0]; i++)
{
Object obj;
float x0 = (boxes[i * 4]) - this->dw;
float y0 = (boxes[i * 4 + 1]) - this->dh;
float x1 = (boxes[i * 4 + 2]) - this->dw;
float y1 = (boxes[i * 4 + 3]) - this->dh;
x0 = clamp(x0 * this->ratio, 0.f, this->w);
y0 = clamp(y0 * this->ratio, 0.f, this->h);
x1 = clamp(x1 * this->ratio, 0.f, this->w);
y1 = clamp(y1 * this->ratio, 0.f, this->h);
obj.rect.x = x0;
obj.rect.y = y0;
obj.rect.width = x1 - x0;
obj.rect.height = y1 - y0;
obj.prob = scores[i];
obj.label = labels[i];
objs.push_back(obj);
}
}
static void draw_objects(const cv::Mat& image, cv::Mat& res, const std::vector<Object>& objs)
{
res = image.clone();
for (auto& obj : objs)
{
cv::Scalar color = cv::Scalar(COLORS[obj.label][0], COLORS[obj.label][1], COLORS[obj.label][2]);
cv::rectangle(res, obj.rect, color, 2);
char text[256];
sprintf(text, "%s %.1f%%", CLASS_NAMES[obj.label], obj.prob * 100);
int baseLine = 0;
cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.4, 1, &baseLine);
int x = (int)obj.rect.x;
int y = (int)obj.rect.y + 1;
if (y > res.rows)
y = res.rows;
cv::rectangle(res, cv::Rect(x, y, label_size.width, label_size.height + baseLine), RECT_COLOR, -1);
cv::putText(res, text, cv::Point(x, y + label_size.height),
cv::FONT_HERSHEY_SIMPLEX, 0.4, TXT_COLOR, 1);
}
}
Write Preview
Loading…
Cancel
Save