From 5e85b515212b9b6590d55d5e0e7bc44ceb185905 Mon Sep 17 00:00:00 2001
From: triplemu <gpu@163.com>
Date: Mon, 27 May 2024 13:41:43 +0800
Subject: [PATCH] add new hwc->nchw

---
 csrc/detect/end2end/CMakeLists.txt     | 18 ++++++++++++------
 csrc/detect/end2end/include/yolov8.hpp | 21 +++++++++++++++++----
 2 files changed, 29 insertions(+), 10 deletions(-)

diff --git a/csrc/detect/end2end/CMakeLists.txt b/csrc/detect/end2end/CMakeLists.txt
index 1aafc18..96afcd7 100644
--- a/csrc/detect/end2end/CMakeLists.txt
+++ b/csrc/detect/end2end/CMakeLists.txt
@@ -15,6 +15,7 @@ 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")
+message(STATUS "CUDA Lib Dir: \n${CUDA_LIB_DIR}\n")
 
 # OpenCV
 find_package(OpenCV REQUIRED)
@@ -24,13 +25,15 @@ 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)
+set(TensorRT_LIBRARIES nvinfer nvinfer_plugin)
+set(TensorRT_LIB_DIR /usr/lib/x86_64-linux-gnu)
 
 
 message(STATUS "TensorRT Libs: \n${TensorRT_LIBRARIES}\n")
+message(STATUS "TensorRT Lib Dir: \n${TensorRT_LIB_DIR}\n")
 message(STATUS "TensorRT Headers: \n${TensorRT_INCLUDE_DIRS}\n")
 
-list(APPEND INCLUDE_DIRS
+list(APPEND ALL_INCLUDE_DIRS
         ${CUDA_INCLUDE_DIRS}
         ${OpenCV_INCLUDE_DIRS}
         ${TensorRT_INCLUDE_DIRS}
@@ -39,12 +42,14 @@ list(APPEND INCLUDE_DIRS
 
 list(APPEND ALL_LIBS
         ${CUDA_LIBRARIES}
-        ${CUDA_LIB_DIR}
         ${OpenCV_LIBRARIES}
         ${TensorRT_LIBRARIES}
         )
 
-include_directories(${INCLUDE_DIRS})
+list(APPEND ALL_LIB_DIRS
+        ${CUDA_LIB_DIR}
+        ${TensorRT_LIB_DIR}
+)
 
 add_executable(${PROJECT_NAME}
         main.cpp
@@ -52,5 +57,6 @@ add_executable(${PROJECT_NAME}
         include/common.hpp
         )
 
-target_link_directories(${PROJECT_NAME} PUBLIC ${ALL_LIBS})
-target_link_libraries(${PROJECT_NAME} PRIVATE nvinfer nvinfer_plugin cudart ${OpenCV_LIBS})
+target_include_directories(${PROJECT_NAME} PUBLIC ${ALL_INCLUDE_DIRS})
+target_link_directories(${PROJECT_NAME} PUBLIC ${ALL_LIB_DIRS})
+target_link_libraries(${PROJECT_NAME} PRIVATE nvinfer ${ALL_LIBS}})
diff --git a/csrc/detect/end2end/include/yolov8.hpp b/csrc/detect/end2end/include/yolov8.hpp
index 22ea548..659dba6 100644
--- a/csrc/detect/end2end/include/yolov8.hpp
+++ b/csrc/detect/end2end/include/yolov8.hpp
@@ -172,7 +172,19 @@ void YOLOv8::letterbox(const cv::Mat& image, cv::Mat& out, cv::Size& size)
 
     cv::copyMakeBorder(tmp, tmp, top, bottom, left, right, cv::BORDER_CONSTANT, {114, 114, 114});
 
-    cv::dnn::blobFromImage(tmp, out, 1 / 255.f, cv::Size(), cv::Scalar(0, 0, 0), true, false, CV_32F);
+    out.create({1, 3, (int)inp_h, (int)inp_w}, CV_32F);
+
+    std::vector<cv::Mat> channels;
+    cv::split(tmp, channels);
+
+    cv::Mat c0((int)inp_h, (int)inp_w, CV_32F, (float*)out.data);
+    cv::Mat c1((int)inp_h, (int)inp_w, CV_32F, (float*)out.data + (int)inp_h * (int)inp_w);
+    cv::Mat c2((int)inp_h, (int)inp_w, CV_32F, (float*)out.data + (int)inp_h * (int)inp_w * 2);
+
+    channels[0].convertTo(c2, CV_32F, 1 / 255.f);
+    channels[1].convertTo(c1, CV_32F, 1 / 255.f);
+    channels[2].convertTo(c0, CV_32F, 1 / 255.f);
+
     this->pparam.ratio  = 1 / r;
     this->pparam.dw     = dw;
     this->pparam.dh     = dh;
@@ -185,8 +197,8 @@ void YOLOv8::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];
+    int      width      = in_binding.dims.d[3];
+    int      height     = in_binding.dims.d[2];
     cv::Size size{width, height};
     this->letterbox(image, nchw, size);
 
@@ -272,8 +284,9 @@ void YOLOv8::draw_objects(const cv::Mat&                                image,
         int x = (int)obj.rect.x;
         int y = (int)obj.rect.y + 1;
 
-        if (y > res.rows)
+        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);