triplemu
7 months ago
parent
6c39635155
commit
19a1fc8ce6
13 changed files with 1951 additions and 26 deletions
@ -0,0 +1,56 @@ |
|||||||
|
cmake_minimum_required(VERSION 3.1) |
||||||
|
|
||||||
|
set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86 89 90) |
||||||
|
set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc) |
||||||
|
|
||||||
|
project(yolov8-cls LANGUAGES CXX CUDA) |
||||||
|
|
||||||
|
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -O3") |
||||||
|
set(CMAKE_CXX_STANDARD 14) |
||||||
|
set(CMAKE_BUILD_TYPE Release) |
||||||
|
option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) |
||||||
|
|
||||||
|
# CUDA |
||||||
|
find_package(CUDA REQUIRED) |
||||||
|
message(STATUS "CUDA Libs: \n${CUDA_LIBRARIES}\n") |
||||||
|
get_filename_component(CUDA_LIB_DIR ${CUDA_LIBRARIES} DIRECTORY) |
||||||
|
message(STATUS "CUDA Headers: \n${CUDA_INCLUDE_DIRS}\n") |
||||||
|
|
||||||
|
# OpenCV |
||||||
|
find_package(OpenCV REQUIRED) |
||||||
|
message(STATUS "OpenCV Libs: \n${OpenCV_LIBS}\n") |
||||||
|
message(STATUS "OpenCV Libraries: \n${OpenCV_LIBRARIES}\n") |
||||||
|
message(STATUS "OpenCV Headers: \n${OpenCV_INCLUDE_DIRS}\n") |
||||||
|
|
||||||
|
# TensorRT |
||||||
|
set(TensorRT_INCLUDE_DIRS /usr/include/x86_64-linux-gnu) |
||||||
|
set(TensorRT_LIBRARIES /usr/lib/x86_64-linux-gnu) |
||||||
|
|
||||||
|
|
||||||
|
message(STATUS "TensorRT Libs: \n${TensorRT_LIBRARIES}\n") |
||||||
|
message(STATUS "TensorRT Headers: \n${TensorRT_INCLUDE_DIRS}\n") |
||||||
|
|
||||||
|
list(APPEND INCLUDE_DIRS |
||||||
|
${CUDA_INCLUDE_DIRS} |
||||||
|
${OpenCV_INCLUDE_DIRS} |
||||||
|
${TensorRT_INCLUDE_DIRS} |
||||||
|
./include |
||||||
|
) |
||||||
|
|
||||||
|
list(APPEND ALL_LIBS |
||||||
|
${CUDA_LIBRARIES} |
||||||
|
${CUDA_LIB_DIR} |
||||||
|
${OpenCV_LIBRARIES} |
||||||
|
${TensorRT_LIBRARIES} |
||||||
|
) |
||||||
|
|
||||||
|
include_directories(${INCLUDE_DIRS}) |
||||||
|
|
||||||
|
add_executable(${PROJECT_NAME} |
||||||
|
main.cpp |
||||||
|
include/yolov8-cls.hpp |
||||||
|
include/common.hpp |
||||||
|
) |
||||||
|
|
||||||
|
target_link_directories(${PROJECT_NAME} PUBLIC ${ALL_LIBS}) |
||||||
|
target_link_libraries(${PROJECT_NAME} PRIVATE nvinfer nvinfer_plugin cudart ${OpenCV_LIBS}) |
||||||
@ -0,0 +1,132 @@ |
|||||||
|
//
|
||||||
|
// Created by ubuntu on 4/27/24.
|
||||||
|
//
|
||||||
|
|
||||||
|
#ifndef CLS_NORMAL_COMMON_HPP |
||||||
|
#define CLS_NORMAL_COMMON_HPP |
||||||
|
#include "NvInfer.h" |
||||||
|
#include "opencv2/opencv.hpp" |
||||||
|
#include <sys/stat.h> |
||||||
|
#include <unistd.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 type_to_size(const nvinfer1::DataType& dataType) |
||||||
|
{ |
||||||
|
switch (dataType) { |
||||||
|
case nvinfer1::DataType::kFLOAT: |
||||||
|
return 4; |
||||||
|
case nvinfer1::DataType::kHALF: |
||||||
|
return 2; |
||||||
|
case nvinfer1::DataType::kINT32: |
||||||
|
return 4; |
||||||
|
case nvinfer1::DataType::kINT8: |
||||||
|
return 1; |
||||||
|
case nvinfer1::DataType::kBOOL: |
||||||
|
return 1; |
||||||
|
default: |
||||||
|
return 4; |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
inline static float clamp(float val, float min, float max) |
||||||
|
{ |
||||||
|
return val > min ? (val < max ? val : max) : min; |
||||||
|
} |
||||||
|
|
||||||
|
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)); |
||||||
|
} |
||||||
|
|
||||||
|
namespace cls { |
||||||
|
struct Binding { |
||||||
|
size_t size = 1; |
||||||
|
size_t dsize = 1; |
||||||
|
nvinfer1::Dims dims; |
||||||
|
std::string name; |
||||||
|
}; |
||||||
|
|
||||||
|
struct Object { |
||||||
|
int label = 0; |
||||||
|
float prob = 0.0; |
||||||
|
}; |
||||||
|
} // namespace cls
|
||||||
|
#endif // CLS_NORMAL_COMMON_HPP
|
||||||
@ -0,0 +1,221 @@ |
|||||||
|
//
|
||||||
|
// Created by ubuntu on 4/27/24.
|
||||||
|
//
|
||||||
|
#ifndef CLS_NORMAL_YOLOv8_cls_HPP |
||||||
|
#define CLS_NORMAL_YOLOv8_cls_HPP |
||||||
|
|
||||||
|
#include "NvInferPlugin.h" |
||||||
|
#include "common.hpp" |
||||||
|
#include "fstream" |
||||||
|
|
||||||
|
using namespace cls; |
||||||
|
|
||||||
|
class YOLOv8_cls { |
||||||
|
public: |
||||||
|
explicit YOLOv8_cls(const std::string& engine_file_path); |
||||||
|
|
||||||
|
~YOLOv8_cls(); |
||||||
|
|
||||||
|
void make_pipe(bool warmup = true); |
||||||
|
|
||||||
|
void copy_from_Mat(const cv::Mat& image); |
||||||
|
|
||||||
|
void copy_from_Mat(const cv::Mat& image, cv::Size& size); |
||||||
|
|
||||||
|
void infer(); |
||||||
|
|
||||||
|
void postprocess(std::vector<Object>& objs); |
||||||
|
|
||||||
|
static void draw_objects(const cv::Mat& image, |
||||||
|
cv::Mat& res, |
||||||
|
const std::vector<Object>& objs, |
||||||
|
const std::vector<std::string>& CLASS_NAMES); |
||||||
|
|
||||||
|
int num_bindings; |
||||||
|
int num_inputs = 0; |
||||||
|
int num_outputs = 0; |
||||||
|
std::vector<Binding> input_bindings; |
||||||
|
std::vector<Binding> output_bindings; |
||||||
|
std::vector<void*> host_ptrs; |
||||||
|
std::vector<void*> device_ptrs; |
||||||
|
|
||||||
|
private: |
||||||
|
nvinfer1::ICudaEngine* engine = nullptr; |
||||||
|
nvinfer1::IRuntime* runtime = nullptr; |
||||||
|
nvinfer1::IExecutionContext* context = nullptr; |
||||||
|
cudaStream_t stream = nullptr; |
||||||
|
Logger gLogger{nvinfer1::ILogger::Severity::kERROR}; |
||||||
|
}; |
||||||
|
|
||||||
|
YOLOv8_cls::YOLOv8_cls(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(); |
||||||
|
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); |
||||||
|
delete[] trtModelStream; |
||||||
|
this->context = this->engine->createExecutionContext(); |
||||||
|
|
||||||
|
assert(this->context != nullptr); |
||||||
|
cudaStreamCreate(&this->stream); |
||||||
|
this->num_bindings = this->engine->getNbBindings(); |
||||||
|
|
||||||
|
for (int i = 0; i < this->num_bindings; ++i) { |
||||||
|
Binding binding; |
||||||
|
nvinfer1::Dims dims; |
||||||
|
nvinfer1::DataType dtype = this->engine->getBindingDataType(i); |
||||||
|
std::string name = this->engine->getBindingName(i); |
||||||
|
binding.name = name; |
||||||
|
binding.dsize = type_to_size(dtype); |
||||||
|
|
||||||
|
bool IsInput = engine->bindingIsInput(i); |
||||||
|
if (IsInput) { |
||||||
|
this->num_inputs += 1; |
||||||
|
dims = this->engine->getProfileDimensions(i, 0, nvinfer1::OptProfileSelector::kMAX); |
||||||
|
binding.size = get_size_by_dims(dims); |
||||||
|
binding.dims = dims; |
||||||
|
this->input_bindings.push_back(binding); |
||||||
|
// set max opt shape
|
||||||
|
this->context->setBindingDimensions(i, dims); |
||||||
|
} |
||||||
|
else { |
||||||
|
dims = this->context->getBindingDimensions(i); |
||||||
|
binding.size = get_size_by_dims(dims); |
||||||
|
binding.dims = dims; |
||||||
|
this->output_bindings.push_back(binding); |
||||||
|
this->num_outputs += 1; |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
YOLOv8_cls::~YOLOv8_cls() |
||||||
|
{ |
||||||
|
this->context->destroy(); |
||||||
|
this->engine->destroy(); |
||||||
|
this->runtime->destroy(); |
||||||
|
cudaStreamDestroy(this->stream); |
||||||
|
for (auto& ptr : this->device_ptrs) { |
||||||
|
CHECK(cudaFree(ptr)); |
||||||
|
} |
||||||
|
|
||||||
|
for (auto& ptr : this->host_ptrs) { |
||||||
|
CHECK(cudaFreeHost(ptr)); |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
void YOLOv8_cls::make_pipe(bool warmup) |
||||||
|
{ |
||||||
|
|
||||||
|
for (auto& bindings : this->input_bindings) { |
||||||
|
void* d_ptr; |
||||||
|
CHECK(cudaMallocAsync(&d_ptr, bindings.size * bindings.dsize, this->stream)); |
||||||
|
this->device_ptrs.push_back(d_ptr); |
||||||
|
} |
||||||
|
|
||||||
|
for (auto& bindings : this->output_bindings) { |
||||||
|
void * d_ptr, *h_ptr; |
||||||
|
size_t size = bindings.size * bindings.dsize; |
||||||
|
CHECK(cudaMallocAsync(&d_ptr, size, this->stream)); |
||||||
|
CHECK(cudaHostAlloc(&h_ptr, size, 0)); |
||||||
|
this->device_ptrs.push_back(d_ptr); |
||||||
|
this->host_ptrs.push_back(h_ptr); |
||||||
|
} |
||||||
|
|
||||||
|
if (warmup) { |
||||||
|
for (int i = 0; i < 10; i++) { |
||||||
|
for (auto& bindings : this->input_bindings) { |
||||||
|
size_t size = bindings.size * bindings.dsize; |
||||||
|
void* h_ptr = malloc(size); |
||||||
|
memset(h_ptr, 0, size); |
||||||
|
CHECK(cudaMemcpyAsync(this->device_ptrs[0], h_ptr, size, cudaMemcpyHostToDevice, this->stream)); |
||||||
|
free(h_ptr); |
||||||
|
} |
||||||
|
this->infer(); |
||||||
|
} |
||||||
|
printf("model warmup 10 times\n"); |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
void YOLOv8_cls::copy_from_Mat(const cv::Mat& image) |
||||||
|
{ |
||||||
|
cv::Mat nchw; |
||||||
|
auto& in_binding = this->input_bindings[0]; |
||||||
|
auto width = in_binding.dims.d[3]; |
||||||
|
auto height = in_binding.dims.d[2]; |
||||||
|
|
||||||
|
cv::dnn::blobFromImage(image, nchw, 1 / 255.f, cv::Size(width, height), cv::Scalar(0, 0, 0), true, false, CV_32F); |
||||||
|
|
||||||
|
this->context->setBindingDimensions(0, nvinfer1::Dims{4, {1, 3, height, width}}); |
||||||
|
|
||||||
|
CHECK(cudaMemcpyAsync( |
||||||
|
this->device_ptrs[0], nchw.ptr<float>(), nchw.total() * nchw.elemSize(), cudaMemcpyHostToDevice, this->stream)); |
||||||
|
} |
||||||
|
|
||||||
|
void YOLOv8_cls::copy_from_Mat(const cv::Mat& image, cv::Size& size) |
||||||
|
{ |
||||||
|
cv::Mat nchw; |
||||||
|
cv::dnn::blobFromImage(image, nchw, 1 / 255.f, size, cv::Scalar(0, 0, 0), true, false, CV_32F); |
||||||
|
this->context->setBindingDimensions(0, nvinfer1::Dims{4, {1, 3, size.height, size.width}}); |
||||||
|
CHECK(cudaMemcpyAsync( |
||||||
|
this->device_ptrs[0], nchw.ptr<float>(), nchw.total() * nchw.elemSize(), cudaMemcpyHostToDevice, this->stream)); |
||||||
|
} |
||||||
|
|
||||||
|
void YOLOv8_cls::infer() |
||||||
|
{ |
||||||
|
|
||||||
|
this->context->enqueueV2(this->device_ptrs.data(), this->stream, nullptr); |
||||||
|
for (int i = 0; i < this->num_outputs; i++) { |
||||||
|
size_t osize = this->output_bindings[i].size * this->output_bindings[i].dsize; |
||||||
|
CHECK(cudaMemcpyAsync( |
||||||
|
this->host_ptrs[i], this->device_ptrs[i + this->num_inputs], osize, cudaMemcpyDeviceToHost, this->stream)); |
||||||
|
} |
||||||
|
cudaStreamSynchronize(this->stream); |
||||||
|
} |
||||||
|
|
||||||
|
void YOLOv8_cls::postprocess(std::vector<Object>& objs) |
||||||
|
{ |
||||||
|
objs.clear(); |
||||||
|
auto num_cls = this->output_bindings[0].dims.d[1]; |
||||||
|
|
||||||
|
float* max_ptr = |
||||||
|
std::max_element(static_cast<float*>(this->host_ptrs[0]), static_cast<float*>(this->host_ptrs[0]) + num_cls); |
||||||
|
Object obj; |
||||||
|
obj.label = std::distance(static_cast<float*>(this->host_ptrs[0]), max_ptr); |
||||||
|
obj.prob = *max_ptr; |
||||||
|
objs.push_back(obj); |
||||||
|
} |
||||||
|
|
||||||
|
void YOLOv8_cls::draw_objects(const cv::Mat& image, |
||||||
|
cv::Mat& res, |
||||||
|
const std::vector<Object>& objs, |
||||||
|
const std::vector<std::string>& CLASS_NAMES) |
||||||
|
{ |
||||||
|
res = image.clone(); |
||||||
|
char text[256]; |
||||||
|
Object obj = objs[0]; |
||||||
|
sprintf(text, "%s %.1f%%", CLASS_NAMES[obj.label].c_str(), obj.prob * 100); |
||||||
|
|
||||||
|
int baseLine = 0; |
||||||
|
cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.4, 1, &baseLine); |
||||||
|
int x = 10; |
||||||
|
int y = 10; |
||||||
|
|
||||||
|
if (y > res.rows) |
||||||
|
y = res.rows; |
||||||
|
|
||||||
|
cv::rectangle(res, cv::Rect(x, y, label_size.width, label_size.height + baseLine), {0, 0, 255}, -1); |
||||||
|
cv::putText(res, text, cv::Point(x, y + label_size.height), cv::FONT_HERSHEY_SIMPLEX, 0.4, {255, 255, 255}, 1); |
||||||
|
} |
||||||
|
|
||||||
|
#endif // CLS_NORMAL_YOLOv8_cls_HPP
|
||||||
File diff suppressed because it is too large
Load Diff
@ -0,0 +1,129 @@ |
|||||||
|
# YOLOv8-cls Model with TensorRT |
||||||
|
|
||||||
|
The yolov8-cls model conversion route is : |
||||||
|
YOLOv8 PyTorch model -> ONNX -> TensorRT Engine |
||||||
|
|
||||||
|
***Notice !!!*** We don't support TensorRT API building !!! |
||||||
|
|
||||||
|
# Export Orin ONNX model by ultralytics |
||||||
|
|
||||||
|
You can leave this repo and use the original `ultralytics` repo for onnx export. |
||||||
|
|
||||||
|
### 1. ONNX -> TensorRT |
||||||
|
|
||||||
|
You can export your onnx model by `ultralytics` API. |
||||||
|
|
||||||
|
``` shell |
||||||
|
yolo export model=yolov8s-cls.pt format=onnx opset=11 simplify=True |
||||||
|
``` |
||||||
|
|
||||||
|
or run this python script: |
||||||
|
|
||||||
|
```python |
||||||
|
from ultralytics import YOLO |
||||||
|
|
||||||
|
# Load a model |
||||||
|
model = YOLO("yolov8s-cls.pt") # load a pretrained model (recommended for training) |
||||||
|
success = model.export(format="onnx", opset=11, simplify=True) # export the model to onnx format |
||||||
|
assert success |
||||||
|
``` |
||||||
|
|
||||||
|
Then build engine by Trtexec Tools. |
||||||
|
|
||||||
|
You can export TensorRT engine by [`trtexec`](https://github.com/NVIDIA/TensorRT/tree/main/samples/trtexec) tools. |
||||||
|
|
||||||
|
Usage: |
||||||
|
|
||||||
|
``` shell |
||||||
|
/usr/src/tensorrt/bin/trtexec \ |
||||||
|
--onnx=yolov8s-cls.onnx \ |
||||||
|
--saveEngine=yolov8s-cls.engine \ |
||||||
|
--fp16 |
||||||
|
``` |
||||||
|
|
||||||
|
### 2. Direct to TensorRT (NOT RECOMMAND!!) |
||||||
|
|
||||||
|
Usage: |
||||||
|
|
||||||
|
```shell |
||||||
|
yolo export model=yolov8s-cls.pt format=engine device=0 |
||||||
|
``` |
||||||
|
|
||||||
|
or run python script: |
||||||
|
|
||||||
|
```python |
||||||
|
from ultralytics import YOLO |
||||||
|
|
||||||
|
# Load a model |
||||||
|
model = YOLO("yolov8s-cls.pt") # load a pretrained model (recommended for training) |
||||||
|
success = model.export(format="engine", device=0) # export the model to engine format |
||||||
|
assert success |
||||||
|
``` |
||||||
|
|
||||||
|
After executing the above script, you will get an engine named `yolov8s-cls.engine` . |
||||||
|
|
||||||
|
# Inference |
||||||
|
|
||||||
|
## Infer with python script |
||||||
|
|
||||||
|
You can infer images with the engine by [`infer-cls.py`](../infer-cls.py) . |
||||||
|
|
||||||
|
Usage: |
||||||
|
|
||||||
|
``` shell |
||||||
|
python3 infer-cls.py \ |
||||||
|
--engine yolov8s-cls.engine \ |
||||||
|
--imgs data \ |
||||||
|
--show \ |
||||||
|
--out-dir outputs \ |
||||||
|
--device cuda:0 |
||||||
|
``` |
||||||
|
|
||||||
|
#### Description of all arguments |
||||||
|
|
||||||
|
- `--engine` : The Engine you export. |
||||||
|
- `--imgs` : The images path you want to detect. |
||||||
|
- `--show` : Whether to show detection results. |
||||||
|
- `--out-dir` : Where to save detection results images. It will not work when use `--show` flag. |
||||||
|
- `--device` : The CUDA deivce you use. |
||||||
|
|
||||||
|
## Inference with c++ |
||||||
|
|
||||||
|
You can infer with c++ in [`csrc/cls/normal`](../csrc/cls/normal) . |
||||||
|
|
||||||
|
### Build: |
||||||
|
|
||||||
|
Please set you own librarys in [`CMakeLists.txt`](../csrc/cls/normal/CMakeLists.txt) and modify `KPS_COLORS` |
||||||
|
and `SKELETON` and `LIMB_COLORS` in [`main.cpp`](../csrc/cls/normal/main.cpp). |
||||||
|
|
||||||
|
Besides, you can modify the postprocess parameters such as `score_thres` and `iou_thres` and `topk` |
||||||
|
in [`main.cpp`](../csrc/cls/normal/main.cpp). |
||||||
|
|
||||||
|
```c++ |
||||||
|
int topk = 100; |
||||||
|
float score_thres = 0.25f; |
||||||
|
float iou_thres = 0.65f; |
||||||
|
``` |
||||||
|
|
||||||
|
And build: |
||||||
|
|
||||||
|
``` shell |
||||||
|
export root=${PWD} |
||||||
|
cd src/cls/normal |
||||||
|
mkdir build |
||||||
|
cmake .. |
||||||
|
make |
||||||
|
mv yolov8-cls ${root} |
||||||
|
cd ${root} |
||||||
|
``` |
||||||
|
|
||||||
|
Usage: |
||||||
|
|
||||||
|
``` shell |
||||||
|
# infer image |
||||||
|
./yolov8-cls yolov8s-cls.engine data/bus.jpg |
||||||
|
# infer images |
||||||
|
./yolov8-cls yolov8s-cls.engine data |
||||||
|
# infer video |
||||||
|
./yolov8-cls yolov8s-cls.engine data/test.mp4 # the video path |
||||||
|
``` |
||||||
@ -0,0 +1,79 @@ |
|||||||
|
import argparse |
||||||
|
from pathlib import Path |
||||||
|
|
||||||
|
import cv2 |
||||||
|
import numpy as np |
||||||
|
|
||||||
|
from config import CLASSES_CLS |
||||||
|
from models.utils import blob, path_to_list |
||||||
|
|
||||||
|
|
||||||
|
def main(args: argparse.Namespace) -> None: |
||||||
|
if args.method == 'cudart': |
||||||
|
from models.cudart_api import TRTEngine |
||||||
|
elif args.method == 'pycuda': |
||||||
|
from models.pycuda_api import TRTEngine |
||||||
|
else: |
||||||
|
raise NotImplementedError |
||||||
|
|
||||||
|
Engine = TRTEngine(args.engine) |
||||||
|
H, W = Engine.inp_info[0].shape[-2:] |
||||||
|
|
||||||
|
images = path_to_list(args.imgs) |
||||||
|
save_path = Path(args.out_dir) |
||||||
|
|
||||||
|
if not args.show and not save_path.exists(): |
||||||
|
save_path.mkdir(parents=True, exist_ok=True) |
||||||
|
|
||||||
|
for image in images: |
||||||
|
save_image = save_path / image.name |
||||||
|
bgr = cv2.imread(str(image)) |
||||||
|
draw = bgr.copy() |
||||||
|
bgr = cv2.resize(bgr, (W, H)) |
||||||
|
rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) |
||||||
|
tensor = blob(rgb, return_seg=False) |
||||||
|
tensor = np.ascontiguousarray(tensor) |
||||||
|
# inference |
||||||
|
data = Engine(tensor) |
||||||
|
|
||||||
|
data = data[0] |
||||||
|
score = data.max().item() |
||||||
|
cls_id = data.argmax().item() |
||||||
|
cls = CLASSES_CLS[cls_id] |
||||||
|
|
||||||
|
text = f'{cls}:{score:.3f}' |
||||||
|
(_w, _h), _bl = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.8, 1) |
||||||
|
_y1 = min(10, draw.shape[0]) |
||||||
|
|
||||||
|
cv2.rectangle(draw, (10, _y1), (10 + _w, _y1 + _h + _bl), (0, 0, 255), -1) |
||||||
|
cv2.putText(draw, text, (10, _y1 + _h), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (255, 255, 255), 2) |
||||||
|
|
||||||
|
if args.show: |
||||||
|
cv2.imshow('result', draw) |
||||||
|
cv2.waitKey(0) |
||||||
|
else: |
||||||
|
cv2.imwrite(str(save_image), draw) |
||||||
|
|
||||||
|
|
||||||
|
def parse_args(): |
||||||
|
parser = argparse.ArgumentParser() |
||||||
|
parser.add_argument('--engine', type=str, help='Engine file') |
||||||
|
parser.add_argument('--imgs', type=str, help='Images file') |
||||||
|
parser.add_argument('--show', |
||||||
|
action='store_true', |
||||||
|
help='Show the detection results') |
||||||
|
parser.add_argument('--out-dir', |
||||||
|
type=str, |
||||||
|
default='./output', |
||||||
|
help='Path to output file') |
||||||
|
parser.add_argument('--method', |
||||||
|
type=str, |
||||||
|
default='cudart', |
||||||
|
help='CUDART pipeline') |
||||||
|
args = parser.parse_args() |
||||||
|
return args |
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__': |
||||||
|
args = parse_args() |
||||||
|
main(args) |
||||||
@ -0,0 +1,74 @@ |
|||||||
|
from models import TRTModule # isort:skip |
||||||
|
import argparse |
||||||
|
from pathlib import Path |
||||||
|
|
||||||
|
import cv2 |
||||||
|
import torch |
||||||
|
|
||||||
|
from config import CLASSES_CLS |
||||||
|
from models.utils import blob, path_to_list |
||||||
|
|
||||||
|
|
||||||
|
def main(args: argparse.Namespace) -> None: |
||||||
|
device = torch.device(args.device) |
||||||
|
Engine = TRTModule(args.engine, device) |
||||||
|
H, W = Engine.inp_info[0].shape[-2:] |
||||||
|
|
||||||
|
images = path_to_list(args.imgs) |
||||||
|
save_path = Path(args.out_dir) |
||||||
|
|
||||||
|
if not args.show and not save_path.exists(): |
||||||
|
save_path.mkdir(parents=True, exist_ok=True) |
||||||
|
|
||||||
|
for image in images: |
||||||
|
save_image = save_path / image.name |
||||||
|
bgr = cv2.imread(str(image)) |
||||||
|
draw = bgr.copy() |
||||||
|
bgr = cv2.resize(bgr, (W, H)) |
||||||
|
rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) |
||||||
|
tensor = blob(rgb, return_seg=False) |
||||||
|
tensor = torch.asarray(tensor, device=device) |
||||||
|
# inference |
||||||
|
data = Engine(tensor) |
||||||
|
|
||||||
|
score, cls_id = data[0].max(0) |
||||||
|
score = float(score) |
||||||
|
cls_id = int(cls_id) |
||||||
|
cls = CLASSES_CLS[cls_id] |
||||||
|
|
||||||
|
text = f'{cls}:{score:.3f}' |
||||||
|
(_w, _h), _bl = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.8, 1) |
||||||
|
_y1 = min(10, draw.shape[0]) |
||||||
|
|
||||||
|
cv2.rectangle(draw, (10, _y1), (10 + _w, _y1 + _h + _bl), (0, 0, 255), -1) |
||||||
|
cv2.putText(draw, text, (10, _y1 + _h), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (255, 255, 255), 2) |
||||||
|
|
||||||
|
if args.show: |
||||||
|
cv2.imshow('result', draw) |
||||||
|
cv2.waitKey(0) |
||||||
|
else: |
||||||
|
cv2.imwrite(str(save_image), draw) |
||||||
|
|
||||||
|
|
||||||
|
def parse_args() -> argparse.Namespace: |
||||||
|
parser = argparse.ArgumentParser() |
||||||
|
parser.add_argument('--engine', type=str, help='Engine file') |
||||||
|
parser.add_argument('--imgs', type=str, help='Images file') |
||||||
|
parser.add_argument('--show', |
||||||
|
action='store_true', |
||||||
|
help='Show the detection results') |
||||||
|
parser.add_argument('--out-dir', |
||||||
|
type=str, |
||||||
|
default='./output', |
||||||
|
help='Path to output file') |
||||||
|
parser.add_argument('--device', |
||||||
|
type=str, |
||||||
|
default='cuda:0', |
||||||
|
help='TensorRT infer device') |
||||||
|
args = parser.parse_args() |
||||||
|
return args |
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__': |
||||||
|
args = parse_args() |
||||||
|
main(args) |
||||||
Loading…
Reference in new issue