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 code style

Browse Source
pull/1/head
triple-Mu 2 years ago
parent 112a5d8f41
commit 8e16a15174
4 changed files with 5 additions and 518 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. 135
      csrc/config/detect_config.hpp
  2. 5
      csrc/end2end/CMakeLists.txt
  3. 2
      csrc/end2end/main.cpp
  4. 381
      csrc/end2end/yolov8.hpp

135
csrc/config/detect_config.hpp
Unescape View File

@ -1,135 +0,0 @@
//
// Created by ubuntu on 1/8/23.
//
#ifndef YOLOV8_TENSORRT_CSRC_DETECT_CONFIG_H
#define YOLOV8_TENSORRT_CSRC_DETECT_CONFIG_H
#include "NvInfer.h"
#include "iostream"
#include "string"
#include <sys/stat.h>
#include <unistd.h>
#include <assert.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)
const int DEVICE = 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_TENSORRT_CSRC_DETECT_CONFIG_H

5
csrc/end2end/CMakeLists.txt
Unescape View File

@ -33,7 +33,7 @@ list(APPEND INCLUDE_DIRS
${CUDA_INCLUDE_DIRS}
${OpenCV_INCLUDE_DIRS}
${TensorRT_INCLUDE_DIRS}
../config
./include
)
list(APPEND ALL_LIBS
@ -46,6 +46,9 @@ include_directories(${INCLUDE_DIRS})
add_executable(${PROJECT_NAME}
main.cpp
include/yolov8.hpp
include/config.h
include/utils.h
)
target_link_directories(${PROJECT_NAME} PUBLIC ${ALL_LIBS})

2
csrc/end2end/main.cpp
Unescape View File

@ -1,7 +1,7 @@
//
// Created by ubuntu on 1/8/23.
//
#include "yolov8.hpp"
#include "include/yolov8.hpp"
int main(int argc, char** argv)
{
cudaSetDevice(DEVICE);

381
csrc/end2end/yolov8.hpp
Unescape View File

@ -1,381 +0,0 @@
//
// Created by ubuntu on 1/8/23.
//
#include "detect_config.hpp"
#include <fstream>
#include <opencv2/opencv.hpp>
#include "NvInferPlugin.h"
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;
};
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