spacewalk01/tensorrt-yolov6

C++ implementation of YOLOv6 using TensorRT

TensorRT-YOLOv6

This repo provides the C++ implementation of YOLOv6 based on TensorRT for performing object detection in real-time.

Environment

• Windows 10
• Visual Studio 2017
• RTX 3070 GPU
• TensorRT 7.2.1.6
• CUDA 11.1, Cudnn 8
• OpenCV 4.5.1 with CUDA

Benchmark

The following result was obtained after performing inference 1000 iterations. It includes image pre-processing, model inference and post-processing (non-max suppression) times. The resolution of the test video was 1280x720.

Model	Device	Mode	Input Shape(HxW)	FPS
YOLOv6-n	RTX3070	FP16	640x640	214
YOLOv6-n	RTX3070	FP32	640x640	165
YOLOv6-tiny	RTX3070	FP16	640x640	193
YOLOv6-tiny	RTX3070	FP32	640x640	128
YOLOv6-s	RTX3070	FP16	640x640	130
YOLOv6-s	RTX3070	FP32	640x640	115

Tutorial

• Step1: Download a pretrained onnx model from the YOLOv6 release

• Step2: Open the solution with Visual Studio and select x64 and Release for the configuration

• Step3: Go to where function initializeSampleParams() is defined and modify the parameters

    // The engine file to generate or to load
    // The engine file does not exist:
    //     This program will try to load onnx file and convert onnx into engine
    // The engine file exists:
    //     This program will load the engine file directly
    params.engingFileName = "yolov6s.engine";

    // The onnx file to load
    params.onnxFileName = "yolov6s.onnx";
    
    // Number of classes (usually 80, but can be other values)
    params.outputClsSize = 80;
    
    // Input video filename
    params.inputVideoName = "test.mp4";
    // Model class filename
    params.cocoClassNamesFileName = "coco.names";
    
    // Threshold values
    params.confThreshold = 0.3;
    params.nmsThreshold = 0.5;
    
    // Floating point precision
    //     By default, fp32
    params.fp16 = true; 

Todo

• Benchmark Test with different yolov6 versions
• Documentation
• Fix arguments

Reference

• yolov4_deepstream
• yolov6

tensorrt-yolov6's Issues

[E] [TRT] INVALID_ARGUMENT: getPluginCreator could not find plugin EfficientNMS_TRT version 1

Hello, I have the following problem here. Can you help solve it?

I use Ubuntu 20.04, ICE CLion, TensorRT 7.2.2.3 and a custom Yolo model with its own weights

Unable to parse ONNX model file: yolov6s.onnx (ERROR: ModelImporter.cpp:296 In function importModel: [5] Assertion failed: tensors.count(input_name))

Hi,

I could compile the code, but following is the runtime error,

&&&& RUNNING TensorRT.sample_yolo # D:\Sushil_Projects\DL_Projects\YoloV6\YoloV6_TensorRT\YoloV6_TensorRT\x64\Release\YoloV6_TensorRT.exe
[07/11/2022-14:57:42] [I] Building and running a GPU inference engine for Yolo
[07/11/2022-14:57:42] [I] Parsing ONNX file: yolov6s.onnx
WARNING: ONNX model has a newer ir_version (0.0.6) than this parser was built against (0.0.3).
While parsing node number 0 [Conv]:
ERROR: ModelImporter.cpp:296 In function importModel:
[5] Assertion failed: tensors.count(input_name)
[07/11/2022-14:57:42] [E] Unable to parse ONNX model file: yolov6s.onnx
&&&& FAILED TensorRT.sample_yolo

I have downloaded the yolov6s.onnx from https://github.com/meituan/YOLOv6/releases/tag/0.1.0

In the code--> initializeSampleParams function is as follows,

---

SampleYoloParams initializeSampleParams(std::vectorstd::string args)
{
SampleYoloParams params;

// The engine file to generate or to load
// The engine file does not exist:
//     This program will try to load onnx file and convert onnx into engine
// The engine file exists:
//     This program will load the engine file directly
params.engingFileName = "yolov6s.engine";

// The onnx file to load
params.onnxFileName = "yolov6s.onnx";

// Input tensor name of ONNX file & engine file
params.inputTensorNames.push_back("image_arrays");

// Old batch configuration, it is zero if explicitBatch flag is true for the tensorrt engine
// May be deprecated in the future
params.batchSize = 0;
params.outputClsSize = 80;

// Threshold values
params.confThreshold = 0.3;
params.nmsThreshold = 0.5;

// Batch size, you can modify to other batch size values if needed
params.explicitBatchSize = 1;
params.width = 640;
params.height = 640;

params.inputVideoName = "test.mp4";
params.cocoClassNamesFileName = "coco.names";

for (auto& arg : args)
{
    params.demo = 1;
    params.outputImageName = "demo_out.jpg";

    if (arg == "--int8")
    {
        params.int8 = true;
    }
    else if (arg == "--fp16")
    {
        params.fp16 = true;
    }
}
specifyInputAndOutputNamesAndShapes(params);

return params;

}

Please let me know your comments,

Thanks

resize.h is missing

resize.h is missing which is included in SampleYolo.hpp. Could you please tell me where can I find resize.h.

could share the test video? I could not download it.

I also want to test

Problem with launch code

Hi, to run your project on Ubuntu and Clean, I had to write CMakeLists.txt . I did it, but when you start your project, the following error is displayed. Can you help me solve it?

I use Ubuntu 20.04, ICE CLion, TensorRT 7.2.2.3 and a custom Yolo model with its own weights

I attach a screenshot of CMakeLists (fisrt photo) and an error (second photo).