diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml new file mode 100644 index 0000000..5dcb6b6 --- /dev/null +++ b/.github/workflows/main.yml @@ -0,0 +1,13 @@ +name: CI + +on: [push, pull_request, workflow_dispatch] + +jobs: + tue-ci: + name: TUe CI - ${{ github.event_name }} + runs-on: ubuntu-latest + steps: + - name: TUe CI + uses: tue-robotics/tue-env/ci/main@master + with: + package: ${{ github.event.repository.name }} diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..eb55d1a --- /dev/null +++ b/.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 diff --git a/.vscode/c_cpp_properties_json b/.vscode/c_cpp_properties_json new file mode 100644 index 0000000..07812f5 --- /dev/null +++ b/.vscode/c_cpp_properties_json @@ -0,0 +1,19 @@ +{ + "configurations": [ + { + "name": "Linux", + "includePath": [ + "${workspaceFolder}/**", + "/usr/include/opencv4", + "${workspaceFolder}/sam_inference/inc", + "${workspaceFolder}/sam_inference/build" + ], + "defines": [], + "compilerPath": "/usr/bin/gcc", + "cStandard": "c11", + "cppStandard": "c++17", + "intelliSenseMode": "linux-gcc-x64" + } + ], + "version": 4 +} \ No newline at end of file diff --git a/.vscode/launch.json b/.vscode/launch.json new file mode 100644 index 0000000..b36e8ac --- /dev/null +++ b/.vscode/launch.json @@ -0,0 +1,48 @@ +{ + "version": "0.2.0", + "configurations": [ + { + "name": "Debug Yolo Model", + "type": "cppdbg", + "request": "launch", + "program": "/home/amigo/ros/noetic/system/devel/lib/yolo_onnx_ros/test_yolo_onnx_ros", // Path to the executable + "args": [ + "/home/amigo/ros/noetic/repos/github.com/tue-robotics/yolo_onnx_ros/data/yolo11m.onnx", + "/home/amigo/Documents/repos/hero_sam.bak/pipeline/build/images" + ], // Add any command-line arguments for your program here + "stopAtEntry": false, + "cwd": "${workspaceFolder}/build", // Set the working directory + "environment": [], + "externalConsole": false, + "MIMode": "gdb", + "setupCommands": [ + { + "description": "Enable pretty-printing for gdb", + "text": "-enable-pretty-printing", + "ignoreFailures": true + } + ], + //"preLaunchTask": "build-yolo-project" // Ensure the project is built before debugging + }, + { + "name": "Debug Yolo Test", + "type": "cppdbg", + "request": "launch", + "program": "/home/amigo/ros/noetic/system/devel/lib/yolo_onnx_ros/yolo_test", // Path to the executable + "args": [], // Add any command-line arguments for your program here + "stopAtEntry": false, + "cwd": "${workspaceFolder}/build", // Set the working directory + "environment": [], + "externalConsole": false, + "MIMode": "gdb", + "setupCommands": [ + { + "description": "Enable pretty-printing for gdb", + "text": "-enable-pretty-printing", + "ignoreFailures": true + } + ], + //"preLaunchTask": "build-yolo-project" // Ensure the project is built before debugging + } + ] +} diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..57bd363 --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,77 @@ +{ + "files.associations": { + "iostream": "cpp", + "cctype": "cpp", + "clocale": "cpp", + "cmath": "cpp", + "cstdarg": "cpp", + "cstddef": "cpp", + "cstdio": "cpp", + "cstdlib": "cpp", + "cstring": "cpp", + "ctime": "cpp", + "cwchar": "cpp", + "cwctype": "cpp", + "array": "cpp", + "atomic": "cpp", + "strstream": "cpp", + "bit": "cpp", + "*.tcc": "cpp", + "bitset": "cpp", + "chrono": "cpp", + "codecvt": "cpp", + "complex": "cpp", + "condition_variable": "cpp", + "cstdint": "cpp", + "deque": "cpp", + "list": "cpp", + "map": "cpp", + "set": "cpp", + "unordered_map": "cpp", + "unordered_set": "cpp", + "vector": "cpp", + "exception": "cpp", + "algorithm": "cpp", + "functional": "cpp", + "iterator": "cpp", + "memory": "cpp", + "memory_resource": "cpp", + "numeric": "cpp", + "optional": "cpp", + "random": "cpp", + "ratio": "cpp", + "regex": "cpp", + "string": "cpp", + "string_view": "cpp", + "system_error": "cpp", + "tuple": "cpp", + "type_traits": "cpp", + "utility": "cpp", + "fstream": "cpp", + "initializer_list": "cpp", + "iomanip": "cpp", + "iosfwd": "cpp", + "istream": "cpp", + "limits": "cpp", + "mutex": "cpp", + "new": "cpp", + "ostream": "cpp", + "shared_mutex": "cpp", + "sstream": "cpp", + "stdexcept": "cpp", + "streambuf": "cpp", + "thread": "cpp", + "cfenv": "cpp", + "cinttypes": "cpp", + "typeindex": "cpp", + "typeinfo": "cpp", + "variant": "cpp", + "compare": "cpp", + "concepts": "cpp", + "numbers": "cpp", + "semaphore": "cpp", + "stop_token": "cpp", + "*.txx": "cpp", + "filesystem": "cpp" + } +} \ No newline at end of file diff --git a/.vscode/tasks.json b/.vscode/tasks.json new file mode 100644 index 0000000..5b58f2b --- /dev/null +++ b/.vscode/tasks.json @@ -0,0 +1,94 @@ +{ + "version": "2.0.0", + "tasks": [ + { + "label": "build-sam-project", + "dependsOrder": "sequence", + "dependsOn": ["configure-sam", "build-sam-binary"], + "problemMatcher": ["$gcc"], + "group": "build" + }, + { + "label": "build-pipeline-project", + "dependsOrder": "sequence", + "dependsOn": ["configure-pipeline", "build-pipeline-binary"], + "problemMatcher": ["$gcc"], + "group": "build" + }, + { + "label": "build-yolo-project", + "dependsOrder": "sequence", + "dependsOn": ["configure-yolo", "build-yolo-binary"], + "problemMatcher": ["$gcc"], + "group": "build" + }, + { + "label": "configure-sam", + "type": "shell", + "command": "cmake", + "args": [ + "-DCMAKE_BUILD_TYPE=Debug", + "-S", "${workspaceFolder}/sam_inference", + "-B", "${workspaceFolder}/sam_inference/build" + ], + "problemMatcher": ["$gcc"] + }, + { + "label": "configure-pipeline", + "type": "shell", + "command": "cmake", + "args": [ + "-DCMAKE_BUILD_TYPE=Debug", + "-S", "${workspaceFolder}/pipeline", + "-B", "${workspaceFolder}/pipeline/build" + ], + "problemMatcher": ["$gcc"] + }, + { + "label": "configure-yolo", + "type": "shell", + "command": "cmake", + "args": [ + "-DCMAKE_BUILD_TYPE=Debug", + "-DCMAKE_POLICY_VERSION_MINIMUM=3.5", + "-S", "${workspaceFolder}/", + "-B", "${workspaceFolder}/build" + ], + "problemMatcher": ["$gcc"] + }, + { + "label": "build-sam-binary", + "type": "shell", + "command": "cmake", + "args": [ + "--build", + "${workspaceFolder}/sam_inference/build", + "--config", "Debug" + ], + "problemMatcher": ["$gcc"] + }, + { + "label": "build-pipeline-binary", + "type": "shell", + "command": "cmake", + "args": [ + "--build", + "${workspaceFolder}/pipeline/build", + "--config", "Debug" + ], + "problemMatcher": ["$gcc"] + }, + { + "label": "build-yolo-binary", + "type": "shell", + "command": "cmake", + "args": [ + "--build", + "${workspaceFolder}/build", + "--config", "Debug" + ], + "problemMatcher": ["$gcc"] + } + + ] +} diff --git a/CMakeLists.txt b/CMakeLists.txt index 575e2e9..09520c9 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,54 +1,117 @@ cmake_minimum_required(VERSION 3.5) -set(PROJECT_NAME Yolov8OnnxRuntimeCPPInference) -project(${PROJECT_NAME} VERSION 0.0.1 LANGUAGES CXX) +project(yolo_onnx_ros) + +add_compile_options(-Wall -Werror=all) +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(catkin REQUIRED + COMPONENTS + onnxruntime_ros +) + + +set(${PROJECT_NAME}_CUDA_ENABLED ${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() -# -------------- ONNXRuntime ------------------# -set(ONNXRUNTIME_VERSION 1.21.0) -set(ONNXRUNTIME_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/../onnxruntime-linux-x64-gpu-1.21.1") -include_directories(${ONNXRUNTIME_ROOT}/include) +configure_file(include/${PROJECT_NAME}/config.hpp.in ${CATKIN_DEVEL_PREFIX}/${CATKIN_GLOBAL_INCLUDE_DESTINATION}/${PROJECT_NAME}/config.hpp) +# add_custom_target(generate_config_hpp +# DEPENDS ${CATKIN_DEVEL_PREFIX}/${CATKIN_GLOBAL_INCLUDE_DESTINATION}/${PROJECT_NAME}/config.hpp +# ) -# -------------- Cuda ------------------# -add_definitions(-DUSE_CUDA=1) -include_directories(/usr/local/cuda/include) +# ------------------------------------------------------------------------------------------------ +# CATKIN EXPORT +# ------------------------------------------------------------------------------------------------ -set(PROJECT_SOURCES - src/main.cpp - src/yolo_inference.cpp +catkin_package( + INCLUDE_DIRS include ${CATKIN_DEVEL_PREFIX}/${CATKIN_GLOBAL_INCLUDE_DESTINATION} + LIBRARIES ${PROJECT_NAME} + CATKIN_DEPENDS + DEPENDS ${${PROJECT_NAME}_CUDA_CATKIN_DEPENDS} OpenCV ) -add_executable(${PROJECT_NAME} ${PROJECT_SOURCES}) -include_directories(${CMAKE_CURRENT_SOURCE_DIR}/inc) +# ------------------------------------------------------------------------------------------------ +# BUILD +# ------------------------------------------------------------------------------------------------ -# Link OpenCV libraries along with ONNX Runtime -target_link_libraries(${PROJECT_NAME} ${OpenCV_LIBS} ${ONNXRUNTIME_ROOT}/lib/libonnxruntime.so) +include_directories( + include + ${CATKIN_DEVEL_PREFIX}/${CATKIN_GLOBAL_INCLUDE_DESTINATION} + SYSTEM + ${${PROJECT_NAME}_CUDA_INCLUDE_DIRS} + ${OpenCV_INCLUDE_DIRS} + ${catkin_INCLUDE_DIRS} +) -# 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" - $) -endif () +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) -# 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) +add_executable(test_${PROJECT_NAME} + src/main.cpp +) +target_link_libraries(test_${PROJECT_NAME} ${PROJECT_NAME} ${OpenCV_LIBRARIES} ${catkin_LIBRARIES}) -# Copy yolov8n.onnx file to the same folder of the executable file -configure_file(model/yolo11m.onnx ${CMAKE_CURRENT_BINARY_DIR}/yolo11m.onnx COPYONLY) +# ------------------------------------------------------------------------------------------------ +# INSTALL +# ------------------------------------------------------------------------------------------------ -# 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 +install(FILES + ${CATKIN_DEVEL_PREFIX}/${CATKIN_PACKAGE_INCLUDE_DESTINATION}/config.hpp + DESTINATION ${CATKIN_PACKAGE_INCLUDE_DESTINATION} ) + +install( + DIRECTORY include/${PROJECT_NAME}/ + DESTINATION ${CATKIN_PACKAGE_INCLUDE_DESTINATION} +) + +install( + TARGETS + ${PROJECT_NAME} + ARCHIVE DESTINATION ${CATKIN_PACKAGE_LIB_DESTINATION} + LIBRARY DESTINATION ${CATKIN_PACKAGE_LIB_DESTINATION} + RUNTIME DESTINATION ${CATKIN_GLOBAL_BIN_DESTINATION} +) + +install( + TARGETS + ${PROJECT_NAME} + DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION} +) + +# ------------------------------------------------------------------------------------------------ +# Testing +# ------------------------------------------------------------------------------------------------ +if (CATKIN_ENABLE_TESTING) + find_package(catkin_lint_cmake REQUIRED) + catkin_add_catkin_lint_test("-W2 --ignore HEADER_OUTSIDE_PACKAGE_INCLUDE_PATH") + + # Utils unit tests (no models needed) + catkin_add_gtest(yolo_test test/yolo_test.cpp) + if(TARGET yolo_test) + target_link_libraries( + yolo_test + ${PROJECT_NAME}_lib + ${catkin_LIBRARIES} + GTest::gtest + GTest::gtest_main + ) + #target_include_directories(yolo_test PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include) + endif() + +endif() diff --git a/README.md b/README.md index ed9d936..37f7eea 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,9 @@ C++ Onnx-runtime -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 results; +std::unique_ptr 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; diff --git a/include/yolo_onnx_ros/config.hpp.in b/include/yolo_onnx_ros/config.hpp.in new file mode 100644 index 0000000..7c7f2ff --- /dev/null +++ b/include/yolo_onnx_ros/config.hpp.in @@ -0,0 +1,8 @@ +#ifndef YOLO_ONNX_ROS_CONFIG_HPP_ +#define YOLO_ONNX_ROS_CONFIG_HPP_ + +#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_ diff --git a/include/yolo_onnx_ros/detection.hpp b/include/yolo_onnx_ros/detection.hpp new file mode 100644 index 0000000..b0d4cbc --- /dev/null +++ b/include/yolo_onnx_ros/detection.hpp @@ -0,0 +1,10 @@ +#include "yolo_onnx_ros/yolo_inference.hpp" + +#include + +// #define LOGGING +std::tuple, DL_INIT_PARAM> Initialize(const std::filesystem::path& model_filename); + +std::vector Detector(std::unique_ptr& p, const cv::Mat& img); + +int ReadCocoYaml(const std::filesystem::path& filename, std::unique_ptr& p); diff --git a/inc/yolo_inference.h b/include/yolo_onnx_ros/yolo_inference.hpp similarity index 66% rename from inc/yolo_inference.h rename to include/yolo_onnx_ros/yolo_inference.hpp index c35aeb6..4ef6b5e 100644 --- a/inc/yolo_inference.h +++ b/include/yolo_onnx_ros/yolo_inference.hpp @@ -1,12 +1,6 @@ #pragma once -#define RET_OK nullptr - -#ifdef _WIN32 -#include -#include -#include -#endif +#define RET_OK nullptr #include #include @@ -14,7 +8,9 @@ #include #include "onnxruntime_cxx_api.h" -#ifdef USE_CUDA +#include + +#if defined(YOLO_ONNX_ROS_CUDA_ENABLED) && YOLO_ONNX_ROS_CUDA_ENABLED #include #endif @@ -59,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& oResult); + // imgSize is [width, height] + char* PreProcess(const cv::Mat& iImg, const std::vector& iImgSize, cv::Mat& oImg); std::vector classes{}; @@ -76,20 +74,19 @@ class YOLO_V8 // Note: The logic is on the .cpp file since its a private method. template char* TensorProcess(clock_t& starttime_1, const cv::Mat& iImg, N& blob, std::vector& inputNodeDims, - std::vector& oResult); + std::vector& oResult); - char* PreProcess(const cv::Mat& iImg, std::vector iImgSize, cv::Mat& oImg); - Ort::Env env; - std::unique_ptr session; - bool cudaEnable; + Ort::Env env_; + std::unique_ptr session_; + bool cudaEnable_; Ort::RunOptions options; - std::vector inputNodeNames; - std::vector outputNodeNames; - - MODEL_TYPE modelType; - std::vector imgSize; - float rectConfidenceThreshold; - float iouThreshold; - float resizeScales;//letterbox scale -}; \ No newline at end of file + std::vector inputNodeNames_; + std::vector outputNodeNames_; + + MODEL_TYPE modelType_; + std::vector imgSize_; + float rectConfidenceThreshold_; + float iouThreshold_; + float resizeScales_; //letterbox scale +}; diff --git a/package.xml b/package.xml new file mode 100644 index 0000000..8106bc0 --- /dev/null +++ b/package.xml @@ -0,0 +1,29 @@ + + + + yolo_onnx_ros + 0.0.0 + Yolo inference + + Iason Theodorou + + ToDo + + catkin + + libopencv-dev + libopencv-dev + onnxruntime_ros + onnxruntime_ros + + catkin_lint_cmake + + doxygen + + + + + + diff --git a/src/detection.cpp b/src/detection.cpp new file mode 100644 index 0000000..8200467 --- /dev/null +++ b/src/detection.cpp @@ -0,0 +1,178 @@ +#include "yolo_onnx_ros/detection.hpp" +#include + +#include +#include +#include + +std::vector Detector(std::unique_ptr& p, const cv::Mat& img) { + + std::vector res; + p->RunSession(img, res); + #ifdef LOGGING + for (auto& re : res) + { + cv::RNG rng(cv::getTickCount()); + cv::Scalar color(rng.uniform(0, 256), rng.uniform(0, 256), rng.uniform(0, 256)); + + cv::rectangle(img, re.box, color, 3); + + float confidence = floor(100 * re.confidence) / 100; + std::cout << std::fixed << std::setprecision(2); + std::string label = p->classes[re.classId] + " " + + std::to_string(confidence).substr(0, std::to_string(confidence).size() - 4); + + cv::rectangle( + img, + cv::Point(re.box.x, re.box.y - 25), + cv::Point(re.box.x + label.length() * 15, re.box.y), + color, + cv::FILLED + ); + + cv::putText( + img, + label, + cv::Point(re.box.x, re.box.y - 5), + cv::FONT_HERSHEY_SIMPLEX, + 0.75, + cv::Scalar(0, 0, 0), + 2 + ); + } + std::cout << "Press any key to exit" << std::endl; + cv::imshow("Result of Detection", img); + cv::waitKey(0); + cv::destroyAllWindows(); + #endif + return res; + +} + + + +// void Classifier(std::unique_ptr& p) +// { +// std::filesystem::path current_path = std::filesystem::current_path(); +// std::filesystem::path imgs_path = current_path;// / "images" +// std::random_device rd; +// std::mt19937 gen(rd()); +// std::uniform_int_distribution dis(0, 255); +// for (auto& i : std::filesystem::directory_iterator(imgs_path)) +// { +// if (i.path().extension() == ".jpg" || i.path().extension() == ".png") +// { +// std::string img_path = i.path().string(); +// //std::cout << img_path << std::endl; +// cv::Mat img = cv::imread(img_path); +// std::vector res; +// const char* ret = p->RunSession(img, res); + +// float positionY = 50; +// for (int i = 0; i < res.size(); i++) +// { +// int r = dis(gen); +// int g = dis(gen); +// int b = dis(gen); +// cv::putText(img, std::to_string(i) + ":", cv::Point(10, positionY), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(b, g, r), 2); +// cv::putText(img, std::to_string(res.at(i).confidence), cv::Point(70, positionY), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(b, g, r), 2); +// positionY += 50; +// } + +// cv::imshow("TEST_CLS", img); +// cv::waitKey(0); +// cv::destroyAllWindows(); +// } + +// } +// } + + + +int ReadCocoYaml(const std::filesystem::path& filename, std::unique_ptr& p) +{ + // Open the YAML file + std::ifstream file(filename); + if (!file.is_open()) + { + std::cerr << "Failed to open file" << std::endl; + return 1; + } + + // Read the file line by line + std::string line; + std::vector lines; + while (std::getline(file, line)) + { + lines.push_back(line); + } + + // Find the start and end of the names section + std::size_t start = 0; + std::size_t end = 0; + for (std::size_t i = 0; i < lines.size(); i++) + { + if (lines[i].find("names:") != std::string::npos) + { + start = i + 1; + } + else if (start > 0 && lines[i].find(':') == std::string::npos) + { + end = i; + break; + } + } + + // Extract the names + std::vector names; + for (std::size_t i = start; i < end; i++) + { + std::stringstream ss(lines[i]); + std::string name; + std::getline(ss, name, ':'); // Extract the number before the delimiter + std::getline(ss, name); // Extract the string after the delimiter + names.push_back(name); + } + + p->classes = names; + return 0; +} + +std::tuple, DL_INIT_PARAM> Initialize(const std::filesystem::path& model_filename) +{ + std::unique_ptr yoloDetector = std::make_unique(); + ReadCocoYaml(model_filename.parent_path() / "coco.yaml", yoloDetector); + DL_INIT_PARAM params; + params.rectConfidenceThreshold = 0.1; + params.iouThreshold = 0.5; + params.modelPath = model_filename; + params.imgSize = { 640, 640 }; + #if defined(YOLO_ONNX_ROS_CUDA_ENABLED) && YOLO_ONNX_ROS_CUDA_ENABLED + params.cudaEnable = true; + + // GPU FP32 inference + params.modelType = YOLO_DETECT_V8; + // GPU FP16 inference + //Note: change fp16 onnx model + //params.modelType = YOLO_DETECT_V8_HALF; + + #else + // CPU inference + params.modelType = YOLO_DETECT_V8; + params.cudaEnable = false; + + #endif + yoloDetector->CreateSession(params); + return std::make_tuple(std::move(yoloDetector), std::move(params)); +} + + +// void ClsTest() +// { +// std::unique_ptr yoloDetector = std::make_unique(); +// std::string model_path = "cls.onnx"; +// ReadCocoYaml(yoloDetector); +// DL_INIT_PARAM params{ model_path, YOLO_CLS, {224, 224} }; +// yoloDetector->CreateSession(params); +// Classifier(yoloDetector); +// } diff --git a/src/main.cpp b/src/main.cpp index 091b185..a637b05 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -1,195 +1,45 @@ #include #include -#include "yolo_inference.h" +#include "yolo_onnx_ros/detection.hpp" #include #include #include -void Detector(std::unique_ptr& p) { - std::filesystem::path current_path = std::filesystem::current_path(); - std::filesystem::path imgs_path = current_path / "images"; - for (auto& i : std::filesystem::directory_iterator(imgs_path)) - { - if (i.path().extension() == ".jpg" || i.path().extension() == ".png" || i.path().extension() == ".jpeg") - { - std::string img_path = i.path().string(); - cv::Mat img = cv::imread(img_path); - std::vector res; - p->RunSession(img, res); - - for (auto& re : res) - { - cv::RNG rng(cv::getTickCount()); - cv::Scalar color(rng.uniform(0, 256), rng.uniform(0, 256), rng.uniform(0, 256)); - - cv::rectangle(img, re.box, color, 3); - - float confidence = floor(100 * re.confidence) / 100; - std::cout << std::fixed << std::setprecision(2); - std::string label = p->classes[re.classId] + " " + - std::to_string(confidence).substr(0, std::to_string(confidence).size() - 4); - - cv::rectangle( - img, - cv::Point(re.box.x, re.box.y - 25), - cv::Point(re.box.x + label.length() * 15, re.box.y), - color, - cv::FILLED - ); - - cv::putText( - img, - label, - cv::Point(re.box.x, re.box.y - 5), - cv::FONT_HERSHEY_SIMPLEX, - 0.75, - cv::Scalar(0, 0, 0), - 2 - ); - +int main(int argc, char *argv[]) +{ + std::unique_ptr yoloDetector; + DL_INIT_PARAM params; - } - std::cout << "Press any key to exit" << std::endl; - cv::imshow("Result of Detection", img); - cv::waitKey(0); - cv::destroyAllWindows(); - } + if (argc < 2) + { + std::cerr << "Not enough args provided" << std::endl; + return 1; } -} + const std::filesystem::path model_name = argv[1]; -void Classifier(std::unique_ptr& p) -{ - std::filesystem::path current_path = std::filesystem::current_path(); - std::filesystem::path imgs_path = current_path;// / "images" - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_int_distribution dis(0, 255); + std::tie(yoloDetector, params) = Initialize(model_name); + + std::filesystem::path imgs_path = argv[2]; for (auto& i : std::filesystem::directory_iterator(imgs_path)) { - if (i.path().extension() == ".jpg" || i.path().extension() == ".png") + if (i.path().extension() == ".jpg" || i.path().extension() == ".png" || i.path().extension() == ".jpeg") { std::string img_path = i.path().string(); - //std::cout << img_path << std::endl; cv::Mat img = cv::imread(img_path); - std::vector res; - const char* ret = p->RunSession(img, res); - - float positionY = 50; - for (int i = 0; i < res.size(); i++) + std::vector results; + results = Detector(yoloDetector, img); + #ifdef LOGGING + for (const auto& result : results) { - int r = dis(gen); - int g = dis(gen); - int b = dis(gen); - cv::putText(img, std::to_string(i) + ":", cv::Point(10, positionY), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(b, g, r), 2); - cv::putText(img, std::to_string(res.at(i).confidence), cv::Point(70, positionY), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(b, g, r), 2); - positionY += 50; + std::cout << "Image path: " << img_path << "\n" + << "class id: " << result.classId << "\n" + << "confidence: " << result.confidence << "\n"; } - - cv::imshow("TEST_CLS", img); - cv::waitKey(0); - cv::destroyAllWindows(); - //cv::imwrite("E:\\output\\" + std::to_string(k) + ".png", img); + #endif } - } -} - -int ReadCocoYaml(std::unique_ptr& p) { - // Open the YAML file - std::ifstream file("coco.yaml"); - if (!file.is_open()) - { - std::cerr << "Failed to open file" << std::endl; - return 1; - } - - // Read the file line by line - std::string line; - std::vector lines; - while (std::getline(file, line)) - { - lines.push_back(line); - } - - // Find the start and end of the names section - std::size_t start = 0; - std::size_t end = 0; - for (std::size_t i = 0; i < lines.size(); i++) - { - if (lines[i].find("names:") != std::string::npos) - { - start = i + 1; - } - else if (start > 0 && lines[i].find(':') == std::string::npos) - { - end = i; - break; - } - } - - // Extract the names - std::vector names; - for (std::size_t i = start; i < end; i++) - { - std::stringstream ss(lines[i]); - std::string name; - std::getline(ss, name, ':'); // Extract the number before the delimiter - std::getline(ss, name); // Extract the string after the delimiter - names.push_back(name); - } - - p->classes = names; return 0; } - - -void DetectTest() -{ - //YOLO_V8* yoloDetector = new YOLO_V8; - std::unique_ptr yoloDetector = std::make_unique(); - - ReadCocoYaml(yoloDetector); - DL_INIT_PARAM params; - params.rectConfidenceThreshold = 0.1; - params.iouThreshold = 0.5; - params.modelPath = "yolo11m.onnx"; - params.imgSize = { 640, 640 }; -#ifdef USE_CUDA - params.cudaEnable = true; - - // GPU FP32 inference - params.modelType = YOLO_DETECT_V8; - // GPU FP16 inference - //Note: change fp16 onnx model - //params.modelType = YOLO_DETECT_V8_HALF; - -#else - // CPU inference - params.modelType = YOLO_DETECT_V8; - params.cudaEnable = false; - -#endif - yoloDetector->CreateSession(params); - Detector(yoloDetector); -} - - -void ClsTest() -{ - std::unique_ptr yoloDetector = std::make_unique(); - std::string model_path = "cls.onnx"; - ReadCocoYaml(yoloDetector); - DL_INIT_PARAM params{ model_path, YOLO_CLS, {224, 224} }; - yoloDetector->CreateSession(params); - Classifier(yoloDetector); -} - - -int main() -{ - DetectTest(); - //ClsTest(); -} \ No newline at end of file diff --git a/src/yolo_inference.cpp b/src/yolo_inference.cpp index 9f2f6bc..ebda07b 100644 --- a/src/yolo_inference.cpp +++ b/src/yolo_inference.cpp @@ -1,18 +1,10 @@ -#include "yolo_inference.h" +#include "yolo_onnx_ros/yolo_inference.hpp" #include #define benchmark #define min(a, b) (((a) < (b)) ? (a) : (b)) -YOLO_V8::YOLO_V8() -{ -} - -YOLO_V8::~YOLO_V8() -{ -} - -#ifdef USE_CUDA +#if defined(YOLO_ONNX_ROS_CUDA_ENABLED) && YOLO_ONNX_ROS_CUDA_ENABLED namespace Ort { template<> @@ -49,8 +41,13 @@ char* BlobFromImage(const cv::Mat& iImg, T& iBlob) { } -char* YOLO_V8::PreProcess(const cv::Mat& iImg, std::vector iImgSize, cv::Mat& oImg) +char* YOLO_V8::PreProcess(const cv::Mat& iImg, const std::vector& iImgSize, cv::Mat& oImg) { + // Basic validation + if (iImg.empty()) return (char*)"[YOLO_V8]: Empty input image."; + if (iImgSize.size() != 2 || iImgSize[0] <= 0 || iImgSize[1] <= 0) + return (char*)"[YOLO_V8]: Invalid target image size."; + if (iImg.channels() == 3) { oImg = iImg.clone(); @@ -61,7 +58,7 @@ char* YOLO_V8::PreProcess(const cv::Mat& iImg, std::vector iImgSize, cv::Ma cv::cvtColor(iImg, oImg, cv::COLOR_GRAY2RGB); } - switch (modelType) + switch (modelType_) { case YOLO_DETECT_V8: case YOLO_POSE: @@ -70,20 +67,24 @@ char* YOLO_V8::PreProcess(const cv::Mat& iImg, std::vector iImgSize, cv::Ma { if (iImg.cols >= iImg.rows) { - resizeScales = iImg.cols / (float)iImgSize.at(0); - cv::resize(oImg, oImg, cv::Size(iImgSize.at(0), int(iImg.rows / resizeScales))); + // scale by target width + resizeScales_ = iImg.cols / (float)iImgSize.at(0); + cv::resize(oImg, oImg, cv::Size(iImgSize.at(0), int(std::round(iImg.rows / resizeScales_)))); } else { - resizeScales = iImg.rows / (float)iImgSize.at(0); - cv::resize(oImg, oImg, cv::Size(int(iImg.cols / resizeScales), iImgSize.at(1))); + // scale by target height + resizeScales_ = iImg.rows / (float)iImgSize.at(1); + cv::resize(oImg, oImg, cv::Size(int(std::round(iImg.cols / resizeScales_)), iImgSize.at(1))); } - cv::Mat tempImg = cv::Mat::zeros(iImgSize.at(0), iImgSize.at(1), CV_8UC3); + // Note: cv::Mat takes (rows, cols) = (height, width) + cv::Mat tempImg = cv::Mat::zeros(iImgSize.at(1), iImgSize.at(0), CV_8UC3); oImg.copyTo(tempImg(cv::Rect(0, 0, oImg.cols, oImg.rows))); oImg = tempImg; break; } case YOLO_CLS: // CenterCrop + case YOLO_CLS_HALF: { int h = iImg.rows; int w = iImg.cols; @@ -100,18 +101,18 @@ char* YOLO_V8::PreProcess(const cv::Mat& iImg, std::vector iImgSize, cv::Ma const char* YOLO_V8::CreateSession(DL_INIT_PARAM& iParams) { const char* Ret = RET_OK; - if (session) + if (session_) { // Clear node names from previous declaration - for (auto& name : inputNodeNames) { + for (auto& name : inputNodeNames_) { delete[] name; } - inputNodeNames.clear(); + inputNodeNames_.clear(); - for (auto& name : outputNodeNames) { + for (auto& name : outputNodeNames_) { delete[] name; } - outputNodeNames.clear(); + outputNodeNames_.clear(); } std::regex pattern("[\u4e00-\u9fa5]"); bool result = std::regex_search(iParams.modelPath, pattern); @@ -123,12 +124,12 @@ const char* YOLO_V8::CreateSession(DL_INIT_PARAM& iParams) { } try { - rectConfidenceThreshold = iParams.rectConfidenceThreshold; - iouThreshold = iParams.iouThreshold; - imgSize = iParams.imgSize; - modelType = iParams.modelType; - cudaEnable = iParams.cudaEnable; - env = Ort::Env(ORT_LOGGING_LEVEL_WARNING, "Yolo"); + rectConfidenceThreshold_ = iParams.rectConfidenceThreshold; + iouThreshold_ = iParams.iouThreshold; + imgSize_ = iParams.imgSize; + modelType_ = iParams.modelType; + cudaEnable_ = iParams.cudaEnable; + env_ = Ort::Env(ORT_LOGGING_LEVEL_WARNING, "Yolo"); Ort::SessionOptions sessionOption; if (iParams.cudaEnable) { @@ -141,23 +142,23 @@ const char* YOLO_V8::CreateSession(DL_INIT_PARAM& iParams) { sessionOption.SetLogSeverityLevel(iParams.logSeverityLevel); const char *modelPath = iParams.modelPath.c_str(); - session = std::make_unique(env, modelPath, sessionOption); + session_ = std::make_unique(env_, modelPath, sessionOption); Ort::AllocatorWithDefaultOptions allocator; - size_t inputNodesNum = session->GetInputCount(); + size_t inputNodesNum = session_->GetInputCount(); for (size_t i = 0; i < inputNodesNum; i++) { - Ort::AllocatedStringPtr input_node_name = session->GetInputNameAllocated(i, allocator); + Ort::AllocatedStringPtr input_node_name = session_->GetInputNameAllocated(i, allocator); char* temp_buf = new char[50]; strcpy(temp_buf, input_node_name.get()); - inputNodeNames.push_back(temp_buf); + inputNodeNames_.push_back(temp_buf); } - size_t OutputNodesNum = session->GetOutputCount(); + size_t OutputNodesNum = session_->GetOutputCount(); for (size_t i = 0; i < OutputNodesNum; i++) { - Ort::AllocatedStringPtr output_node_name = session->GetOutputNameAllocated(i, allocator); + Ort::AllocatedStringPtr output_node_name = session_->GetOutputNameAllocated(i, allocator); char* temp_buf = new char[10]; strcpy(temp_buf, output_node_name.get()); - outputNodeNames.push_back(temp_buf); + outputNodeNames_.push_back(temp_buf); } options = Ort::RunOptions{ nullptr }; WarmUpSession(); @@ -181,24 +182,24 @@ const char* YOLO_V8::CreateSession(DL_INIT_PARAM& iParams) { const char* YOLO_V8::RunSession(const cv::Mat& iImg, std::vector& oResult) { #ifdef benchmark clock_t starttime_1 = clock(); -#endif // benchmark - +#endif const char* Ret = RET_OK; cv::Mat processedImg; - PreProcess(iImg, imgSize, processedImg); - if (modelType < 4) + PreProcess(iImg, imgSize_, processedImg); + if (modelType_ < 4) { float* blob = new float[processedImg.total() * 3]; BlobFromImage(processedImg, blob); - std::vector inputNodeDims = { 1, 3, imgSize.at(0), imgSize.at(1) }; + // ONNX expects {N, C, H, W} = {1, 3, height, width} + std::vector inputNodeDims = { 1, 3, imgSize_.at(1), imgSize_.at(0) }; TensorProcess(starttime_1, iImg, blob, inputNodeDims, oResult); } else { -#ifdef USE_CUDA +#if defined(YOLO_ONNX_ROS_CUDA_ENABLED) && YOLO_ONNX_ROS_CUDA_ENABLED half* blob = new half[processedImg.total() * 3]; BlobFromImage(processedImg, blob); - std::vector inputNodeDims = { 1,3,imgSize.at(0),imgSize.at(1) }; + std::vector inputNodeDims = { 1, 3, imgSize_.at(1), imgSize_.at(0) }; TensorProcess(starttime_1, iImg, blob, inputNodeDims, oResult); #endif } @@ -208,16 +209,16 @@ const char* YOLO_V8::RunSession(const cv::Mat& iImg, std::vector& oRe template -char* YOLO_V8::TensorProcess(clock_t& starttime_1, const cv::Mat& iImg, N& blob, std::vector& inputNodeDims, +char* YOLO_V8::TensorProcess(clock_t& starttime_1, const cv::Mat&, N& blob, std::vector& inputNodeDims, std::vector& oResult) { Ort::Value inputTensor = Ort::Value::CreateTensor::type>( - Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU), blob, 3 * imgSize.at(0) * imgSize.at(1), + Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU), blob, 3 * imgSize_.at(0) * imgSize_.at(1), inputNodeDims.data(), inputNodeDims.size()); #ifdef benchmark clock_t starttime_2 = clock(); #endif // benchmark - auto outputTensor = session->Run(options, inputNodeNames.data(), &inputTensor, 1, outputNodeNames.data(), - outputNodeNames.size()); + auto outputTensor = session_->Run(options, inputNodeNames_.data(), &inputTensor, 1, outputNodeNames_.data(), + outputNodeNames_.size()); #ifdef benchmark clock_t starttime_3 = clock(); #endif // benchmark @@ -227,7 +228,7 @@ char* YOLO_V8::TensorProcess(clock_t& starttime_1, const cv::Mat& iImg, N& blob, std::vector outputNodeDims = tensor_info.GetShape(); auto output = outputTensor.front().GetTensorMutableData::type>(); delete[] blob; - switch (modelType) + switch (modelType_) { case YOLO_DETECT_V8: case YOLO_DETECT_V8_HALF: @@ -238,7 +239,7 @@ char* YOLO_V8::TensorProcess(clock_t& starttime_1, const cv::Mat& iImg, N& blob, std::vector confidences; std::vector boxes; cv::Mat rawData; - if (modelType == YOLO_DETECT_V8) + if (modelType_ == YOLO_DETECT_V8) { // FP32 rawData = cv::Mat(signalResultNum, strideNum, CV_32F, output); @@ -263,7 +264,7 @@ char* YOLO_V8::TensorProcess(clock_t& starttime_1, const cv::Mat& iImg, N& blob, cv::Point class_id; double maxClassScore; cv::minMaxLoc(scores, 0, &maxClassScore, 0, &class_id); - if (maxClassScore > rectConfidenceThreshold) + if (maxClassScore > rectConfidenceThreshold_) { confidences.push_back(maxClassScore); class_ids.push_back(class_id.x); @@ -272,19 +273,19 @@ char* YOLO_V8::TensorProcess(clock_t& starttime_1, const cv::Mat& iImg, N& blob, float w = data[2]; float h = data[3]; - int left = int((x - 0.5 * w) * resizeScales); - int top = int((y - 0.5 * h) * resizeScales); + int left = int((x - 0.5 * w) * resizeScales_); + int top = int((y - 0.5 * h) * resizeScales_); - int width = int(w * resizeScales); - int height = int(h * resizeScales); + int width = int(w * resizeScales_); + int height = int(h * resizeScales_); boxes.push_back(cv::Rect(left, top, width, height)); } data += signalResultNum; } std::vector nmsResult; - cv::dnn::NMSBoxes(boxes, confidences, rectConfidenceThreshold, iouThreshold, nmsResult); - for (int i = 0; i < nmsResult.size(); ++i) + cv::dnn::NMSBoxes(boxes, confidences, rectConfidenceThreshold_, iouThreshold_, nmsResult); + for (uint i = 0; i < nmsResult.size(); ++i) { int idx = nmsResult[i]; DL_RESULT result; @@ -299,7 +300,7 @@ char* YOLO_V8::TensorProcess(clock_t& starttime_1, const cv::Mat& iImg, N& blob, double pre_process_time = (double)(starttime_2 - starttime_1) / CLOCKS_PER_SEC * 1000; double process_time = (double)(starttime_3 - starttime_2) / CLOCKS_PER_SEC * 1000; double post_process_time = (double)(starttime_4 - starttime_3) / CLOCKS_PER_SEC * 1000; - if (cudaEnable) + if (cudaEnable_) { std::cout << "[YOLO_V8(CUDA)]: " << pre_process_time << "ms pre-process, " << process_time << "ms inference, " << post_process_time << "ms post-process." << std::endl; } @@ -315,7 +316,7 @@ char* YOLO_V8::TensorProcess(clock_t& starttime_1, const cv::Mat& iImg, N& blob, case YOLO_CLS_HALF: { cv::Mat rawData; - if (modelType == YOLO_CLS) { + if (modelType_ == YOLO_CLS) { // FP32 rawData = cv::Mat(1, this->classes.size(), CV_32F, output); } else { @@ -326,7 +327,7 @@ char* YOLO_V8::TensorProcess(clock_t& starttime_1, const cv::Mat& iImg, N& blob, float *data = (float *) rawData.data; DL_RESULT result; - for (int i = 0; i < this->classes.size(); i++) + for (uint i = 0; i < this->classes.size(); i++) { result.classId = i; result.confidence = data[i]; @@ -344,43 +345,44 @@ char* YOLO_V8::TensorProcess(clock_t& starttime_1, const cv::Mat& iImg, N& blob, char* YOLO_V8::WarmUpSession() { clock_t starttime_1 = clock(); - cv::Mat iImg = cv::Mat(cv::Size(imgSize.at(0), imgSize.at(1)), CV_8UC3); + // cv::Size takes (width, height) + cv::Mat iImg = cv::Mat(cv::Size(imgSize_.at(0), imgSize_.at(1)), CV_8UC3); cv::Mat processedImg; - PreProcess(iImg, imgSize, processedImg); - if (modelType < 4) + PreProcess(iImg, imgSize_, processedImg); + if (modelType_ < 4) { float* blob = new float[iImg.total() * 3]; BlobFromImage(processedImg, blob); - std::vector YOLO_input_node_dims = { 1, 3, imgSize.at(0), imgSize.at(1) }; + std::vector YOLO_input_node_dims = { 1, 3, imgSize_.at(1), imgSize_.at(0) }; Ort::Value input_tensor = Ort::Value::CreateTensor( - Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU), blob, 3 * imgSize.at(0) * imgSize.at(1), + Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU), blob, 3 * imgSize_.at(0) * imgSize_.at(1), YOLO_input_node_dims.data(), YOLO_input_node_dims.size()); - auto output_tensors = session->Run(options, inputNodeNames.data(), &input_tensor, 1, outputNodeNames.data(), - outputNodeNames.size()); + auto output_tensors = session_->Run(options, inputNodeNames_.data(), &input_tensor, 1, outputNodeNames_.data(), + outputNodeNames_.size()); delete[] blob; clock_t starttime_4 = clock(); double post_process_time = (double)(starttime_4 - starttime_1) / CLOCKS_PER_SEC * 1000; - if (cudaEnable) + if (cudaEnable_) { std::cout << "[YOLO_V8(CUDA)]: " << "Cuda warm-up cost " << post_process_time << " ms. " << std::endl; } } else { -#ifdef USE_CUDA +#if defined(YOLO_ONNX_ROS_CUDA_ENABLED) && YOLO_ONNX_ROS_CUDA_ENABLED half* blob = new half[iImg.total() * 3]; BlobFromImage(processedImg, blob); - std::vector YOLO_input_node_dims = { 1,3,imgSize.at(0),imgSize.at(1) }; - Ort::Value input_tensor = Ort::Value::CreateTensor(Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU), blob, 3 * imgSize.at(0) * imgSize.at(1), YOLO_input_node_dims.data(), YOLO_input_node_dims.size()); - auto output_tensors = session->Run(options, inputNodeNames.data(), &input_tensor, 1, outputNodeNames.data(), outputNodeNames.size()); + std::vector YOLO_input_node_dims = { 1, 3, imgSize_.at(1), imgSize_.at(0) }; + Ort::Value input_tensor = Ort::Value::CreateTensor(Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU), blob, 3 * imgSize_.at(0) * imgSize_.at(1), YOLO_input_node_dims.data(), YOLO_input_node_dims.size()); + auto output_tensors = session_->Run(options, inputNodeNames_.data(), &input_tensor, 1, outputNodeNames_.data(), outputNodeNames_.size()); delete[] blob; clock_t starttime_4 = clock(); double post_process_time = (double)(starttime_4 - starttime_1) / CLOCKS_PER_SEC * 1000; - if (cudaEnable) + if (cudaEnable_) { std::cout << "[YOLO_V8(CUDA)]: " << "Cuda warm-up cost " << post_process_time << " ms. " << std::endl; } #endif } return RET_OK; -} \ No newline at end of file +} diff --git a/test/yolo_test.cpp b/test/yolo_test.cpp new file mode 100644 index 0000000..5e2a4f5 --- /dev/null +++ b/test/yolo_test.cpp @@ -0,0 +1,103 @@ +#include "yolo_onnx_ros/yolo_inference.hpp" +#include "yolo_onnx_ros/detection.hpp" +#include +#include +#include + +using namespace testing; + +class YoloInferenceTest : public testing::Test +{ +protected: + void SetUp() override + { + // Create test images with different characteristics + testImage_640x640 = cv::Mat::ones(640, 640, CV_8UC3) * 255; + testImage_800x600 = cv::Mat::ones(600, 800, CV_8UC3) * 128; + + // Create a more realistic test image with some patterns + testImage_realistic = cv::Mat(640, 640, CV_8UC3); + cv::randu(testImage_realistic, cv::Scalar(0, 0, 0), cv::Scalar(255, 255, 255)); + + // Setup common parameters + std::tie(yolo, params) = Initialize(); + + + NonSquareImgSize = { testImage_800x600.cols, testImage_800x600.rows }; + } + + void TearDown() override + { + // Clean up if needed + yolo.reset(); + } + + // Test data + cv::Mat testImage_640x640; + cv::Mat testImage_800x600; + cv::Mat testImage_realistic; + DL_INIT_PARAM params; + std::vector NonSquareImgSize; + std::unique_ptr yolo; + std::vector results; +}; + +TEST_F(YoloInferenceTest, ObjectCreation) +{ + EXPECT_NO_THROW({ + YOLO_V8 localYolo; + }); +} + +TEST_F(YoloInferenceTest, PreProcessSquareImage) +{ + cv::Mat processedImg; + char* result = yolo->PreProcess(testImage_640x640, params.imgSize, processedImg); + + EXPECT_EQ(result, nullptr) << "PreProcess should succeed"; + EXPECT_EQ(processedImg.size(), cv::Size(640, 640)) << "Output should be 640x640"; + EXPECT_FALSE(processedImg.empty()) << "Processed image should not be empty"; +} + +TEST_F(YoloInferenceTest, PreProcessRectangularImage) +{ + cv::Mat processedImg; + char* result = yolo->PreProcess(testImage_800x600, NonSquareImgSize, processedImg); + + EXPECT_EQ(result, nullptr) << "PreProcess should succeed"; + EXPECT_EQ(processedImg.size(), cv::Size(800, 600)) << "Output should be letterboxed to 800x600"; + EXPECT_FALSE(processedImg.empty()) << "Processed image should not be empty"; +} + +TEST_F(YoloInferenceTest, CreateSessionWithValidModel) +{ + const char* result = yolo->CreateSession(params); + EXPECT_EQ(result, nullptr) << "CreateSession should succeed with valid parameters"; +} + +TEST_F(YoloInferenceTest, CreateSessionWithInvalidModel) +{ + params.modelPath = "nonexistent_model.onnx"; + const char* result = yolo->CreateSession(params); + EXPECT_NE(result, nullptr) << "CreateSession should fail with invalid model path"; +} + +TEST_F(YoloInferenceTest, FullInferencePipeline) +{ + // First create session + const char* createResult = yolo->CreateSession(params); + ASSERT_EQ(createResult, nullptr) << "Session creation must succeed for inference test"; + + const char* runResult = yolo->RunSession(testImage_realistic, results); + + EXPECT_EQ(runResult, nullptr) << "RunSession should succeed"; + // Note: results might be empty for random test image, that's okay + EXPECT_TRUE(results.size() >= 0) << "Results should be a valid vector"; +} + +// Run all tests +int main(int argc, char **argv) +{ + testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +}