add YOLOv9 support

hhx518 · Mar 5, 2024 · 0315b15 · 0315b15
1 parent e56b38f
commit 0315b15
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 5 deletions.
diff --git a/src/application/app_yolo.cpp b/src/application/app_yolo.cpp
@@ -282,5 +282,6 @@ int app_yolo(){
     // test(Yolo::Type::V5, TRT::Mode::FP32, "yolov5s");
     // test(Yolo::Type::V6, TRT::Mode::FP32, "yolov6s");
     // test(Yolo::Type::V7, TRT::Mode::FP32, "yolov7");
+    // test(Yolo::Type::V9, TRT::Mode::FP32, "yolov9c");
     return 0;
 }
diff --git a/src/application/app_yolo/yolo.cpp b/src/application/app_yolo/yolo.cpp
@@ -22,6 +22,7 @@ namespace Yolo{
         case Type::V6: return "YoloV6";
         case Type::V8: return "YoloV8";
         case Type::X: return "YoloX";
+        case Type::V9: return "YoloV9";
         default: return "Unknow";
         }
     }
@@ -169,7 +170,7 @@ namespace Yolo{
             NMSMethod nms_method, int max_objects,
             bool use_multi_preprocess_stream
         ){
-            if(type == Type::V5 || type == Type::V3 || type == Type::V7 || type == Type::V8 || type == Type::V6){
+            if(type == Type::V5 || type == Type::V3 || type == Type::V7 || type == Type::V8 || type == Type::V6 || type == Type::V9){
                 normalize_ = CUDAKernel::Norm::alpha_beta(1 / 255.0f, 0.0f, CUDAKernel::ChannelType::Invert);
             }else if(type == Type::X){
                 //float mean[] = {0.485, 0.456, 0.406};
@@ -211,7 +212,7 @@ namespace Yolo{
             int max_batch_size = engine->get_max_batch_size();
             auto input         = engine->tensor("images");
             auto output        = engine->tensor("output");
-            int num_classes    = (type_ == Type::V8) ? output->size(2) - 4 : output->size(2) - 5;
+            int num_classes    = (type_ == Type::V8 || type_ == Type::V9) ? output->size(2) - 4 : output->size(2) - 5;
 
             input_width_       = input->size(3);
             input_height_      = input->size(2);
@@ -406,7 +407,7 @@ namespace Yolo{
     void image_to_tensor(const cv::Mat& image, shared_ptr<TRT::Tensor>& tensor, Type type, int ibatch){
 
         CUDAKernel::Norm normalize;
-        if(type == Type::V5 || type == Type::V3 || type == Type::V7 || type == Type::V8 || type == Type::V6){
+        if(type == Type::V5 || type == Type::V3 || type == Type::V7 || type == Type::V8 || type == Type::V6 || type == Type::V9){
             normalize = CUDAKernel::Norm::alpha_beta(1 / 255.0f, 0.0f, CUDAKernel::ChannelType::Invert);
         }else if(type == Type::X){
             //float mean[] = {0.485, 0.456, 0.406};

diff --git a/src/application/app_yolo/yolo.hpp b/src/application/app_yolo/yolo.hpp
@@ -24,7 +24,8 @@ namespace Yolo{
         V3 = 2,
         V7 = 3,
         V8 = 4,
-        V6 = 5
+        V6 = 5,
+        V9 = 6
     };
 
     enum class NMSMethod : int{

diff --git a/src/application/app_yolo/yolo_decode.cu b/src/application/app_yolo/yolo_decode.cu
@@ -156,7 +156,7 @@ namespace Yolo{
 
         auto grid = CUDATools::grid_dims(num_bboxes);
         auto block = CUDATools::block_dims(num_bboxes);
-        if(type == Type::V8){
+        if(type == Type::V8 || type == Type::V9){
             checkCudaKernel(decode_kernel_v8<<<grid, block, 0, stream>>>(predict, num_bboxes, num_classes, confidence_threshold, invert_affine_matrix, parray, max_objects));            
         }else{
             checkCudaKernel(decode_kernel<<<grid, block, 0, stream>>>(predict, num_bboxes, num_classes, confidence_threshold, invert_affine_matrix, parray, max_objects));
-Original file line number
+Diff line change
@@ Expand Up / @@ -24,7 +24,8 @@ namespace Yolo{ @@
             V3 = 2,
             V7 = 3,
             V8 = 4,
-            V6 = 5
+            V6 = 5,
+            V9 = 6
         };
         enum class NMSMethod : int{
@@ Expand Down @@