add new hwc->nchw

pull/219/head
triplemu 6 months ago
parent 94ec430a68
commit 5e85b51521
  1. 18
      csrc/detect/end2end/CMakeLists.txt
  2. 21
      csrc/detect/end2end/include/yolov8.hpp

@ -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}})

@ -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);

Loading…
Cancel
Save