Feature/refactoring yolo (#4)

Author: MatthijsBurgh

.github/workflows/main.yml

@@ -1,6 +1,6 @@
 name: CI
 
-on: [push, pull_request]
+on: [push, pull_request, workflow_dispatch]
 
 jobs:
   tue-ci:

.gitignore

@@ -0,0 +1,17 @@
+build/
+images/*
+onnxruntime*/
+onnxruntime/*
+docker/*
+CMakefile
+CMakeCache.txt
+CMakeFiles/*
+cmake_install.cmake
+Makefile
+SPEED-SAM-C-TENSORRT/
+sam_inference/model/FastSAM-x.onnx
+mask*
+segmentation_results*
+
+# Models
+*.onnx

CMakeLists.txt

@@ -1,4 +1,5 @@
-cmake_minimum_required(VERSION 3.5)
+cmake_minimum_required(VERSION 3.14)
+
 project(yolo_onnx_ros)
 
 add_compile_options(-Wall -Werror=all)
@@ -7,22 +8,21 @@ add_compile_options(-Wextra -Werror=extra)
 # -------------- Support C++17 for using filesystem  ------------------#
 set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
-set(CMAKE_CXX_EXTENSIONS ON)
-#set(CMAKE_INCLUDE_CURRENT_DIR ON)
 
-# -------------- OpenCV  ------------------#
-find_package(OpenCV REQUIRED)
-include_directories(${OpenCV_INCLUDE_DIRS})
 
+find_package(OpenCV REQUIRED)
 find_package(catkin REQUIRED
   COMPONENTS
   onnxruntime_ros
 )
 
 
 set(${PROJECT_NAME}_CUDA_ENABLED ${onnxruntime_ros_CUDA_ENABLED})
-if(onnxruntime_ros_CUDA_ENABLED)
+if(${PROJECT_NAME}_CUDA_ENABLED)
   find_package(CUDAToolkit REQUIRED)
+  set(${PROJECT_NAME}_CUDA_CATKIN_DEPENDS "CUDAToolkit")
+  set(${PROJECT_NAME}_CUDA_INCLUDE_DIRS "${CUDAToolkit_INCLUDE_DIRS}")
+  set(${PROJECT_NAME}_CUDA_TARGET_LINK_LIBRARIES "CUDA::cudart")
 endif()
 
 configure_file(include/${PROJECT_NAME}/config.hpp.in ${CATKIN_DEVEL_PREFIX}/${CATKIN_GLOBAL_INCLUDE_DESTINATION}/${PROJECT_NAME}/config.hpp)
@@ -38,10 +38,9 @@ catkin_package(
   INCLUDE_DIRS include ${CATKIN_DEVEL_PREFIX}/${CATKIN_GLOBAL_INCLUDE_DESTINATION}
   LIBRARIES ${PROJECT_NAME}
   CATKIN_DEPENDS
-  DEPENDS CUDAToolkit OpenCV
+  DEPENDS ${${PROJECT_NAME}_CUDA_CATKIN_DEPENDS} OpenCV
 )
 
-
 # ------------------------------------------------------------------------------------------------
 #                                           BUILD
 # ------------------------------------------------------------------------------------------------
@@ -50,33 +49,23 @@ include_directories(
   include
   ${CATKIN_DEVEL_PREFIX}/${CATKIN_GLOBAL_INCLUDE_DESTINATION}
   SYSTEM
-  ${CUDAToolkit_INCLUDE_DIRS}
+  ${${PROJECT_NAME}_CUDA_INCLUDE_DIRS}
   ${OpenCV_INCLUDE_DIRS}
   ${catkin_INCLUDE_DIRS}
 )
 
-# -------------- ONNXRuntime  ------------------#
-
 add_library(${PROJECT_NAME}
+  src/detection.cpp
   src/yolo_inference.cpp
 )
+target_link_libraries(${PROJECT_NAME} ${${PROJECT_NAME}_CUDA_TARGET_LINK_LIBRARIES} ${OpenCV_LIBRARIES} ${catkin_LIBRARIES})
 # add_dependencies(${PROJECT_NAME} generate_config_hpp)
-target_link_libraries(${PROJECT_NAME} CUDA::toolkit ${OpenCV_LIBRARIES} ${catkin_LIBRARIES})
 
 add_executable(test_${PROJECT_NAME}
   src/main.cpp
 )
 target_link_libraries(test_${PROJECT_NAME} ${PROJECT_NAME} ${OpenCV_LIBRARIES} ${catkin_LIBRARIES})
 
-# For Windows system, copy onnxruntime.dll to the same folder of the executable file
-if (WIN32)
-    add_custom_command(TARGET ${PROJECT_NAME} POST_BUILD
-            COMMAND ${CMAKE_COMMAND} -E copy_if_different
-            "${ONNXRUNTIME_ROOT}/lib/onnxruntime.dll"
-            $<TARGET_FILE_DIR:${PROJECT_NAME}>)
-endif ()
-
-
 # ------------------------------------------------------------------------------------------------
 #                                               INSTALL
 # ------------------------------------------------------------------------------------------------
@@ -104,15 +93,3 @@ install(
     ${PROJECT_NAME}
   DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION}
 )
-
-# Download https://raw.githubusercontent.com/ultralytics/ultralytics/main/ultralytics/cfg/datasets/coco.yaml
-# and put it in the same folder of the executable file
-configure_file(data/coco.yaml ${CMAKE_CURRENT_BINARY_DIR}/coco.yaml COPYONLY)
-
-# Copy yolov8n.onnx file to the same folder of the executable file
-# configure_file(model/yolo11m.onnx ${CMAKE_CURRENT_BINARY_DIR}/yolo11m.onnx COPYONLY)
-
-# # Create folder name images in the same folder of the executable file
-# add_custom_command(TARGET ${PROJECT_NAME} POST_BUILD
-#     COMMAND ${CMAKE_COMMAND} -E make_directory ${CMAKE_CURRENT_BINARY_DIR}/images
-# )

README.md

@@ -2,7 +2,9 @@
 
 <img alt="C++" src="https://img.shields.io/badge/C++-17-blue.svg?style=flat&logo=c%2B%2B"> <img alt="Onnx-runtime" src="https://img.shields.io/badge/OnnxRuntime-717272.svg?logo=Onnx&logoColor=white">
 
-This example demonstrates how to perform inference using YOLOv8 in C++ with ONNX Runtime and OpenCV's API.
+
+
+This algorithm is inspired by [Ultralitics](https://github.com/ultralytics/ultralytics/tree/main/examples/YOLOv8-ONNXRuntime-CPP) implementation to perform inference using YOLOv8 (we also supports v11) in C++ with ONNX Runtime and OpenCV's API.
 
 ## Benefits ✨
 
@@ -57,12 +59,12 @@ In order to run example, you also need to download coco.yaml. You can download t
 | OpenCV                           | >=4.0.0       |
 | C++ Standard                     | >=17          |
 | Cmake                            | >=3.5         |
-| Cuda (Optional)                  | >=11.4 \<12.0 |
-| cuDNN (Cuda required)            | =8            |
+| Cuda (Optional)                  |  =12.8        |
+| cuDNN (Cuda required)            | =9            |
 
 Note: The dependency on C++17 is due to the usage of the C++17 filesystem feature.
 
-Note (2): Due to ONNX Runtime, we need to use CUDA 11 and cuDNN 8. Keep in mind that this requirement might change in the future.
+Note (2): Due to ONNX Runtime, we need to use CUDA 12.8 and cuDNN 9. Keep in mind that this requirement might change in the future.
 
 ## Build 🛠️
 
@@ -87,12 +89,8 @@ Note (2): Due to ONNX Runtime, we need to use CUDA 11 and cuDNN 8. Keep in mind
    If you encounter an error indicating that the `ONNXRUNTIME_ROOT` variable is not set correctly, you can resolve this by building the project using the appropriate command tailored to your system.
 
    ```console
-   # compiled in a win32 system
-   cmake -D WIN32=TRUE ..
    # compiled in a linux system
    cmake -D LINUX=TRUE ..
-   # compiled in an apple system
-   cmake -D APPLE=TRUE ..
    ```
 
 5. Build the project:
@@ -104,9 +102,15 @@ Note (2): Due to ONNX Runtime, we need to use CUDA 11 and cuDNN 8. Keep in mind
 6. The built executable should now be located in the `build` directory.
 
 ## Usage 🚀
-
+To run from main just run the executable.
+To run the detector on you C++ application:
 ```c++
-//change your param as you like
+//To run the detector add on you C++ application:
+std::vector<DL_RESULT> results;
+std::unique_ptr<YOLO_V8> yoloDetector = Initialize();
+results = DetectObjects(yoloDetector, img);
+
+//You can change your param as you like (inside the Initialize() function)
 //Pay attention to your device and the onnx model type(fp32 or fp16)
 DL_INIT_PARAM params;
 params.rectConfidenceThreshold = 0.1;

include/yolo_onnx_ros/config.hpp.in

@@ -1,7 +1,8 @@
 #ifndef YOLO_ONNX_ROS_CONFIG_HPP_
 #define YOLO_ONNX_ROS_CONFIG_HPP_
 
-//Set which version of the Tree Interface to use
-#define YOLO_ONNX_ROS_CUDA_ENABLED @onnx_yolo_ros_CUDA_ENABLED@
+#define YOLO_ONNX_FALSE 0
+#define YOLO_ONNX_TRUE 1
+#define YOLO_ONNX_ROS_CUDA_ENABLED YOLO_ONNX_@yolo_onnx_ros_CUDA_ENABLED@
 
 #endif //#define YOLO_ONNX_ROS_CONFIG_HPP_

include/yolo_onnx_ros/detection.hpp

@@ -0,0 +1,5 @@
+#include "yolo_onnx_ros/yolo_inference.hpp"
+
+std::tuple<std::unique_ptr<YOLO_V8>, DL_INIT_PARAM> Initialize();
+std::vector<DL_RESULT> Detector(std::unique_ptr<YOLO_V8>& p, const cv::Mat& img);
+int ReadCocoYaml(std::unique_ptr<YOLO_V8>& p);

include/yolo_onnx_ros/yolo_inference.h renamed to include/yolo_onnx_ros/yolo_inference.hpp

@@ -1,22 +1,16 @@
 #pragma once
 
-#define    RET_OK nullptr
-
-#ifdef _WIN32
-#include <Windows.h>
-#include <direct.h>
-#include <io.h>
-#endif
+#define RET_OK nullptr
 
 #include <string>
 #include <vector>
 #include <cstdio>
 #include <opencv2/opencv.hpp>
-#include <onnxruntime_cxx_api.h>
+#include "onnxruntime_cxx_api.h"
 
 #include <yolo_onnx_ros/config.hpp>
 
-#ifdef YOLO_ONNX_ROS_CUDA_ENABLED
+#if defined(YOLO_ONNX_ROS_CUDA_ENABLED) && YOLO_ONNX_ROS_CUDA_ENABLED
 #include <cuda_fp16.h>
 #endif
 
@@ -61,14 +55,16 @@ typedef struct _DL_RESULT
 class YOLO_V8
 {
 public:
-    YOLO_V8();
+    YOLO_V8() = default;
 
-    ~YOLO_V8();
+    ~YOLO_V8() = default;
 
 public:
     const char* CreateSession(DL_INIT_PARAM& iParams);
 
     const char* RunSession(const cv::Mat& iImg, std::vector<DL_RESULT>& oResult);
+    // imgSize is [width, height]
+    char* PreProcess(const cv::Mat& iImg, const std::vector<int>& iImgSize, cv::Mat& oImg);
 
     std::vector<std::string> classes{};
 
@@ -77,15 +73,14 @@ class YOLO_V8
 
     // Note: The logic is on the .cpp file since its a private method.
     template<typename N>
-    char* TensorProcess(clock_t& starttime_1, N& blob, std::vector<int64_t>& inputNodeDims,
-        std::vector<DL_RESULT>& oResult);
+    char* TensorProcess(clock_t& starttime_1, const cv::Mat& iImg, N& blob, std::vector<int64_t>& inputNodeDims,
+                        std::vector<DL_RESULT>& oResult);
 
-    char* PreProcess(const cv::Mat& iImg, std::vector<int> iImgSize, cv::Mat& oImg);
 
     Ort::Env env_;
     std::unique_ptr<Ort::Session> session_;
     bool cudaEnable_;
-    Ort::RunOptions options_;
+    Ort::RunOptions options;
     std::vector<const char*> inputNodeNames_;
     std::vector<const char*> outputNodeNames_;