Skip to content

Commit

Permalink
YOLOv9 works with TensorRT backend
Browse files Browse the repository at this point in the history
  • Loading branch information
@Nuzhny007
Nuzhny007 committed Feb 25, 2024
1 parent 64bee78 commit baecd89
Show file tree
Hide file tree
Showing 10 changed files with 53 additions and 18 deletions.
6 changes: 2 additions & 4 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,14 +3,12 @@

# Last changes

* YOLOv9 detector worked with TensorRT! Export pretrained Pytorch models [here (WongKinYiu/yolov9)](https://github.com/WongKinYiu/yolov9) to onnx format and run Multitarget-tracker with -e=6 example

* YOLOv8 instance segmentation models worked with TensorRT! Export pretrained Pytorch models [here (ultralytics/ultralytics)](https://github.com/ultralytics/ultralytics) to onnx format and run Multitarget-tracker with -e=6 example

* Re-identification model osnet_x0_25_msmt17 from [mikel-brostrom/yolo_tracking](https://github.com/mikel-brostrom/yolo_tracking)

* YOLOv8 detector worked with TensorRT! Export pretrained Pytorch models [here (ultralytics/ultralytics)](https://github.com/ultralytics/ultralytics) to onnx format and run Multitarget-tracker with -e=6 example

* Some experiments with YOLOv7_mask and results with rotated rectangles: detector works tracker in progress

# New videos!

* YOLOv7 instance segmentation
Expand Down
17 changes: 15 additions & 2 deletions combined/combined.cpp
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
#include <iomanip>
#include <iomanip>
#include <ctime>
#include <future>

Expand Down Expand Up @@ -327,7 +327,8 @@ bool CombinedDetector::InitDetector(cv::UMat frame)
YOLOv7,
YOLOv7Mask,
YOLOv8,
YOLOv8Mask
YOLOv8Mask,
YOLOv9
};
YOLOModels usedModel = YOLOModels::YOLOv8;
switch (usedModel)
Expand Down Expand Up @@ -439,6 +440,18 @@ bool CombinedDetector::InitDetector(cv::UMat frame)
maxBatch = 1;
configDNN.emplace("maxCropRatio", "-1");
break;

case YOLOModels::YOLOv9:
configDNN.emplace("modelConfiguration", pathToModel + "yolov9-c.onnx");
configDNN.emplace("modelBinary", pathToModel + "yolov9-c.onnx");
configDNN.emplace("confidenceThreshold", "0.2");
configDNN.emplace("inference_precision", "FP16");
configDNN.emplace("net_type", "YOLOV9");
configDNN.emplace("inWidth", "640");
configDNN.emplace("inHeight", "640");
maxBatch = 1;
configDNN.emplace("maxCropRatio", "-1");
break;
}
configDNN.emplace("maxBatch", std::to_string(maxBatch));
configDNN.emplace("classNames", pathToModel + "coco.names");
Expand Down
15 changes: 13 additions & 2 deletions example/examples.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -834,9 +834,10 @@ class YoloTensorRTExample final : public VideoExample
YOLOv7,
YOLOv7Mask,
YOLOv8,
YOLOv8Mask
YOLOv8Mask,
YOLOv9
};
YOLOModels usedModel = YOLOModels::YOLOv4;
YOLOModels usedModel = YOLOModels::YOLOv9;
switch (usedModel)
{
case YOLOModels::TinyYOLOv3:
Expand Down Expand Up @@ -938,6 +939,16 @@ class YoloTensorRTExample final : public VideoExample
maxBatch = 1;
config.emplace("maxCropRatio", "-1");
break;

case YOLOModels::YOLOv9:
config.emplace("modelConfiguration", pathToModel + "yolov9-c.onnx");
config.emplace("modelBinary", pathToModel + "yolov9-c.onnx");
config.emplace("confidenceThreshold", "0.2");
config.emplace("inference_precision", "FP32");
config.emplace("net_type", "YOLOV9");
maxBatch = 1;
config.emplace("maxCropRatio", "-1");
break;
}
if (maxBatch < m_batchSize)
maxBatch = m_batchSize;
Expand Down
5 changes: 3 additions & 2 deletions src/Detector/OCVDNNDetector.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,7 @@ bool OCVDNNDetector::Init(const config_t& config)
dictNetType["YOLOV7Mask"] = ModelType::YOLOV7Mask;
dictNetType["YOLOV8"] = ModelType::YOLOV8;
dictNetType["YOLOV8Mask"] = ModelType::YOLOV8Mask;
dictNetType["YOLOV9"] = ModelType::YOLOV9;

auto netType = dictNetType.find(net_type->second);
if (netType != dictNetType.end())
Expand Down Expand Up @@ -345,7 +346,7 @@ void OCVDNNDetector::DetectInCrop(const cv::UMat& colorFrame, const cv::Rect& cr
}
else
{
if (m_netType == ModelType::YOLOV8 || m_netType == ModelType::YOLOV5)
if (m_netType == ModelType::YOLOV8 || m_netType == ModelType::YOLOV5 || m_netType == ModelType::YOLOV9)
{
int rows = detections[0].size[1];
int dimensions = detections[0].size[2];
Expand All @@ -367,7 +368,7 @@ void OCVDNNDetector::DetectInCrop(const cv::UMat& colorFrame, const cv::Rect& cr

for (int i = 0; i < rows; ++i)
{
if (m_netType == ModelType::YOLOV8)
if (m_netType == ModelType::YOLOV8 || m_netType == ModelType::YOLOV9)
{
float* classes_scores = data + 4;

Expand Down
3 changes: 2 additions & 1 deletion src/Detector/OCVDNNDetector.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,8 @@ class OCVDNNDetector final : public BaseDetector
YOLOV7,
YOLOV7Mask,
YOLOV8,
YOLOV8Mask
YOLOV8Mask,
YOLOV9
};

cv::dnn::Net m_net;
Expand Down
1 change: 1 addition & 0 deletions src/Detector/YoloTensorRTDetector.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,7 @@ bool YoloTensorRTDetector::Init(const config_t& config)
dictNetType["YOLOV7Mask"] = tensor_rt::YOLOV7Mask;
dictNetType["YOLOV8"] = tensor_rt::YOLOV8;
dictNetType["YOLOV8Mask"] = tensor_rt::YOLOV8Mask;
dictNetType["YOLOV9"] = tensor_rt::YOLOV9;

auto netType = dictNetType.find(net_type->second);
if (netType != dictNetType.end())
Expand Down
2 changes: 1 addition & 1 deletion src/Detector/tensorrt_yolo/YoloONNX.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1087,7 +1087,7 @@ void YoloONNX::ProcessBBoxesOutput(size_t imgIdx, const std::vector<float*>& out
}
else if (outputs.size() == 1)
{
if (m_params.m_netType == tensor_rt::ModelType::YOLOV8)
if (m_params.m_netType == tensor_rt::ModelType::YOLOV8 || m_params.m_netType == tensor_rt::ModelType::YOLOV9)
{
//0: name: images, size: 1x3x640x640
//1: name: output0, size: 1x84x8400
Expand Down
13 changes: 10 additions & 3 deletions src/Detector/tensorrt_yolo/class_detector.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,9 @@ namespace tensor_rt
// Input tensor name of ONNX file & engine file
if (config.net_type == ModelType::YOLOV6)
m_params.inputTensorNames.push_back("image_arrays");
else if (config.net_type == ModelType::YOLOV7 || config.net_type == ModelType::YOLOV7Mask || config.net_type == ModelType::YOLOV8 || config.net_type == ModelType::YOLOV8Mask)
else if (config.net_type == ModelType::YOLOV7 || config.net_type == ModelType::YOLOV7Mask ||
config.net_type == ModelType::YOLOV8 || config.net_type == ModelType::YOLOV8Mask ||
config.net_type == ModelType::YOLOV9)
m_params.inputTensorNames.push_back("images");

// Threshold values
Expand All @@ -76,7 +78,9 @@ namespace tensor_rt
{
m_params.outputTensorNames.push_back("outputs");
}
else if (config.net_type == ModelType::YOLOV7 || config.net_type == ModelType::YOLOV7Mask || config.net_type == ModelType::YOLOV8 || config.net_type == ModelType::YOLOV8Mask)
else if (config.net_type == ModelType::YOLOV7 || config.net_type == ModelType::YOLOV7Mask ||
config.net_type == ModelType::YOLOV8 || config.net_type == ModelType::YOLOV8Mask ||
config.net_type == ModelType::YOLOV9)
{
//if (config.batch_size == 1)
//{
Expand Down Expand Up @@ -161,7 +165,10 @@ namespace tensor_rt
if (m_impl)
delete m_impl;

if (config.net_type == ModelType::YOLOV6 || config.net_type == ModelType::YOLOV7 || config.net_type == ModelType::YOLOV7Mask || config.net_type == ModelType::YOLOV8 || config.net_type == ModelType::YOLOV8Mask)
if (config.net_type == ModelType::YOLOV6 ||
config.net_type == ModelType::YOLOV7 || config.net_type == ModelType::YOLOV7Mask ||
config.net_type == ModelType::YOLOV8 || config.net_type == ModelType::YOLOV8Mask ||
config.net_type == ModelType::YOLOV9)
m_impl = new YoloONNXImpl();
else
m_impl = new YoloDectectorImpl();
Expand Down
3 changes: 2 additions & 1 deletion src/Detector/tensorrt_yolo/class_detector.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,8 @@ namespace tensor_rt
YOLOV7,
YOLOV7Mask,
YOLOV8,
YOLOV8Mask
YOLOV8Mask,
YOLOV9
};

///
Expand Down
6 changes: 4 additions & 2 deletions src/Detector/tensorrt_yolo/ds_image.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,8 @@ DsImage::DsImage(const cv::Mat& mat_image_, tensor_rt::ModelType net_type, const
}
if (tensor_rt::ModelType::YOLOV5 == net_type || tensor_rt::ModelType::YOLOV6 == net_type ||
tensor_rt::ModelType::YOLOV7 == net_type || tensor_rt::ModelType::YOLOV7Mask == net_type ||
tensor_rt::ModelType::YOLOV8 == net_type || tensor_rt::ModelType::YOLOV8Mask == net_type)
tensor_rt::ModelType::YOLOV8 == net_type || tensor_rt::ModelType::YOLOV8Mask == net_type ||
tensor_rt::ModelType::YOLOV9 == net_type)
{
// resize the DsImage with scale
float r = std::min(static_cast<float>(inputH) / static_cast<float>(m_Height), static_cast<float>(inputW) / static_cast<float>(m_Width));
Expand Down Expand Up @@ -99,7 +100,8 @@ DsImage::DsImage(const std::string& path, tensor_rt::ModelType net_type, const i

if (tensor_rt::ModelType::YOLOV5 == net_type || tensor_rt::ModelType::YOLOV6 == net_type ||
tensor_rt::ModelType::YOLOV7 == net_type || tensor_rt::ModelType::YOLOV7Mask == net_type ||
tensor_rt::ModelType::YOLOV8 == net_type || tensor_rt::ModelType::YOLOV8Mask == net_type)
tensor_rt::ModelType::YOLOV8 == net_type || tensor_rt::ModelType::YOLOV8Mask == net_type ||
tensor_rt::ModelType::YOLOV9 == net_type)
{
// resize the DsImage with scale
float dim = std::max(m_Height, m_Width);
Expand Down

0 comments on commit baecd89

Please sign in to comment.