|
| 1 | +cmake_minimum_required(VERSION 3.5) |
| 2 | +project(tensorrt_yolo) |
| 3 | + |
| 4 | +if(NOT CMAKE_CXX_STANDARD) |
| 5 | + set(CMAKE_CXX_STANDARD 17) |
| 6 | + set(CMAKE_CXX_STANDARD_REQUIRED ON) |
| 7 | + set(CMAKE_CXX_EXTENSIONS OFF) |
| 8 | +endif() |
| 9 | +if(CMAKE_COMPILER_IS_GNUCXX OR CMAKE_CXX_COMPILER_ID MATCHES "Clang") |
| 10 | + add_compile_options(-Wall -Wextra -Wpedantic -Wno-deprecated-declarations -Werror) |
| 11 | +endif() |
| 12 | + |
| 13 | +option(CUDA_VERBOSE "Verbose output of CUDA modules" OFF) |
| 14 | + |
| 15 | +find_package(ament_cmake_auto REQUIRED) |
| 16 | +ament_auto_find_build_dependencies() |
| 17 | + |
| 18 | +find_package(OpenCV REQUIRED) |
| 19 | + |
| 20 | +# set flags for CUDA availability |
| 21 | +option(CUDA_AVAIL "CUDA available" OFF) |
| 22 | +find_package(CUDA) |
| 23 | +if(CUDA_FOUND) |
| 24 | + find_library(CUBLAS_LIBRARIES cublas HINTS |
| 25 | + ${CUDA_TOOLKIT_ROOT_DIR}/lib64 |
| 26 | + ${CUDA_TOOLKIT_ROOT_DIR}/lib |
| 27 | + ) |
| 28 | + if(CUDA_VERBOSE) |
| 29 | + message(STATUS "CUDA is available!") |
| 30 | + message(STATUS "CUDA Libs: ${CUDA_LIBRARIES}") |
| 31 | + message(STATUS "CUDA Headers: ${CUDA_INCLUDE_DIRS}") |
| 32 | + endif() |
| 33 | + set(CUDA_AVAIL ON) |
| 34 | +else() |
| 35 | + message("CUDA NOT FOUND") |
| 36 | + set(CUDA_AVAIL OFF) |
| 37 | +endif() |
| 38 | + |
| 39 | +# set flags for TensorRT availability |
| 40 | +option(TRT_AVAIL "TensorRT available" OFF) |
| 41 | +# try to find the tensorRT modules |
| 42 | +find_library(NVINFER NAMES nvinfer) |
| 43 | +find_library(NVONNXPARSER nvonnxparser) |
| 44 | +find_library(NVINFER_PLUGIN NAMES nvinfer_plugin) |
| 45 | +if(NVINFER AND NVONNXPARSER AND NVINFER_PLUGIN) |
| 46 | + if(CUDA_VERBOSE) |
| 47 | + message(STATUS "TensorRT is available!") |
| 48 | + message(STATUS "NVINFER: ${NVINFER}") |
| 49 | + message(STATUS "NVPARSERS: ${NVPARSERS}") |
| 50 | + message(STATUS "NVINFER_PLUGIN: ${NVINFER_PLUGIN}") |
| 51 | + message(STATUS "NVONNXPARSER: ${NVONNXPARSER}") |
| 52 | + endif() |
| 53 | + set(TRT_AVAIL ON) |
| 54 | +else() |
| 55 | + message("TensorRT is NOT Available") |
| 56 | + set(TRT_AVAIL OFF) |
| 57 | +endif() |
| 58 | + |
| 59 | +# set flags for CUDNN availability |
| 60 | +option(CUDNN_AVAIL "CUDNN available" OFF) |
| 61 | +# try to find the CUDNN module |
| 62 | +find_library(CUDNN_LIBRARY |
| 63 | +NAMES libcudnn.so${__cudnn_ver_suffix} libcudnn${__cudnn_ver_suffix}.dylib ${__cudnn_lib_win_name} |
| 64 | +PATHS $ENV{LD_LIBRARY_PATH} ${__libpath_cudart} ${CUDNN_ROOT_DIR} ${PC_CUDNN_LIBRARY_DIRS} ${CMAKE_INSTALL_PREFIX} |
| 65 | +PATH_SUFFIXES lib lib64 bin |
| 66 | +DOC "CUDNN library." |
| 67 | +) |
| 68 | +if(CUDNN_LIBRARY) |
| 69 | + if(CUDA_VERBOSE) |
| 70 | + message(STATUS "CUDNN is available!") |
| 71 | + message(STATUS "CUDNN_LIBRARY: ${CUDNN_LIBRARY}") |
| 72 | + endif() |
| 73 | + set(CUDNN_AVAIL ON) |
| 74 | +else() |
| 75 | + message("CUDNN is NOT Available") |
| 76 | + set(CUDNN_AVAIL OFF) |
| 77 | +endif() |
| 78 | + |
| 79 | +# Download onnx |
| 80 | +find_program(GDOWN_AVAIL "gdown") |
| 81 | +if(NOT GDOWN_AVAIL) |
| 82 | + message("gdown: command not found. External files could not be downloaded.") |
| 83 | +endif() |
| 84 | +set(PATH "${CMAKE_CURRENT_SOURCE_DIR}/data") |
| 85 | +if(NOT EXISTS "${PATH}") |
| 86 | + execute_process(COMMAND mkdir -p ${PATH}) |
| 87 | +endif() |
| 88 | +# yolov3 |
| 89 | +set(FILE "${PATH}/yolov3.onnx") |
| 90 | +message(STATUS "Checking and downloading yolov3.onnx") |
| 91 | +if(NOT EXISTS "${FILE}") |
| 92 | + message(STATUS "... file does not exist. Downloading now ...") |
| 93 | + execute_process( |
| 94 | + COMMAND gdown "https://drive.google.com/uc?id=1ZYoBqVynmO5kKntyN56GEbELRpuXG8Ym" -O ${PATH}/yolov3.onnx |
| 95 | + ERROR_QUIET |
| 96 | + ) |
| 97 | +endif() |
| 98 | + |
| 99 | +# yolov4 |
| 100 | +set(FILE "${PATH}/yolov4.onnx") |
| 101 | +message(STATUS "Checking and downloading yolov4.onnx") |
| 102 | +if(NOT EXISTS "${FILE}") |
| 103 | + message(STATUS "... file does not exist. Downloading now ...") |
| 104 | + execute_process( |
| 105 | + COMMAND gdown "https://drive.google.com//uc?id=1vkNmSwcIpTkJ_-BrKhxtit0PBJeJYTVX" -O ${PATH}/yolov4.onnx |
| 106 | + ERROR_QUIET |
| 107 | + ) |
| 108 | +endif() |
| 109 | + |
| 110 | +# yolov4-tiny |
| 111 | +set(FILE "${PATH}/yolov4-tiny.onnx") |
| 112 | +message(STATUS "Checking and downloading yolov4.onnx") |
| 113 | +if(NOT EXISTS "${FILE}") |
| 114 | + message(STATUS "... file does not exist. Downloading now ...") |
| 115 | + execute_process( |
| 116 | + COMMAND gdown "https://drive.google.com/uc?id=1rUHjV_JfkmlVFgb20XXrOMWo_HZAjrFh" -O ${PATH}/yolov4-tiny.onnx |
| 117 | + ERROR_QUIET |
| 118 | + ) |
| 119 | +endif() |
| 120 | + |
| 121 | +# yolov5s |
| 122 | +set(FILE "${PATH}/yolov5s.onnx") |
| 123 | +message(STATUS "Checking and downloading yolov5s.onnx") |
| 124 | +if(NOT EXISTS "${FILE}") |
| 125 | + message(STATUS "... file does not exist. Downloading now ...") |
| 126 | + execute_process( |
| 127 | + COMMAND gdown "https://drive.google.com//uc?id=1CF21nQWigwCPTr5psQZXg6cBQIOYKbad" -O ${PATH}/yolov5s.onnx |
| 128 | + ERROR_QUIET |
| 129 | + ) |
| 130 | +endif() |
| 131 | + |
| 132 | +# yolov5m |
| 133 | +set(FILE "${PATH}/yolov5m.onnx") |
| 134 | +message(STATUS "Checking and downloading yolov5m.onnx") |
| 135 | +if(NOT EXISTS "${FILE}") |
| 136 | + message(STATUS "... file does not exist. Downloading now ...") |
| 137 | + execute_process( |
| 138 | + COMMAND gdown "https://drive.google.com//uc?id=1a1h50KJH6slwmjKZpPlS-errukF-BrgG" -O ${PATH}/yolov5m.onnx |
| 139 | + ERROR_QUIET |
| 140 | + ) |
| 141 | +endif() |
| 142 | + |
| 143 | +# yolov5l |
| 144 | +set(FILE "${PATH}/yolov5l.onnx") |
| 145 | +message(STATUS "Checking and downloading yolov5l.onnx") |
| 146 | +if(NOT EXISTS "${FILE}") |
| 147 | + message(STATUS "... file does not exist. Downloading now ...") |
| 148 | + execute_process( |
| 149 | + COMMAND gdown "https://drive.google.com/uc?id=1xO8S92Cq7qrmx93UHHyA7Cd7-dJsBDP8" -O ${PATH}/yolov5l.onnx |
| 150 | + ERROR_QUIET |
| 151 | + ) |
| 152 | +endif() |
| 153 | + |
| 154 | +# yolov5x |
| 155 | +set(FILE "${PATH}/yolov5x.onnx") |
| 156 | +message(STATUS "Checking and downloading yolov5x.onnx") |
| 157 | +if(NOT EXISTS "${FILE}") |
| 158 | + message(STATUS "... file does not exist. Downloading now ...") |
| 159 | + execute_process( |
| 160 | + COMMAND gdown "https://drive.google.com/uc?id=1kAHuNJUCxpD-yWrS6t95H3zbAPfClLxI" -O ${PATH}/yolov5x.onnx |
| 161 | + ERROR_QUIET |
| 162 | + ) |
| 163 | +endif() |
| 164 | + |
| 165 | +set(FILE "${PATH}/coco.names") |
| 166 | +message(STATUS "Checking and downloading coco.names") |
| 167 | +if(NOT EXISTS "${FILE}") |
| 168 | + message(STATUS "... file does not exist. Downloading now ...") |
| 169 | + execute_process( |
| 170 | + COMMAND gdown "https://drive.google.com/uc?id=19wXD23PE16kJDkZ7j2W3Ijvx5I1kO9kJ" -O ${PATH}/coco.names |
| 171 | + ERROR_QUIET |
| 172 | + ) |
| 173 | +endif() |
| 174 | + |
| 175 | +# create calib image directory for int8 calibration |
| 176 | +set(CALIB_IMAGE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/calib_image") |
| 177 | +if(NOT EXISTS "${CALIB_IMAGE_PATH}") |
| 178 | + execute_process(COMMAND mkdir -p ${CALIB_IMAGE_PATH}) |
| 179 | +endif() |
| 180 | + |
| 181 | +if(TRT_AVAIL AND CUDA_AVAIL AND CUDNN_AVAIL) |
| 182 | + include_directories( |
| 183 | + include |
| 184 | + lib/include |
| 185 | + ${OpenCV_INCLUDE_DIRS} |
| 186 | + ${CUDA_INCLUDE_DIRS} |
| 187 | + ) |
| 188 | + |
| 189 | + |
| 190 | + ### yolo ### |
| 191 | + cuda_add_library(mish_plugin SHARED |
| 192 | + lib/src/plugins/mish.cu |
| 193 | + lib/src/plugins/mish_plugin.cpp |
| 194 | + ) |
| 195 | + |
| 196 | + cuda_add_library(yolo_layer_plugin SHARED |
| 197 | + lib/src/plugins/yolo_layer.cu |
| 198 | + lib/src/plugins/yolo_layer_plugin.cpp |
| 199 | + ) |
| 200 | + |
| 201 | + cuda_add_library(nms_plugin SHARED |
| 202 | + lib/src/plugins/nms.cu |
| 203 | + lib/src/plugins/nms_plugin.cpp |
| 204 | + ) |
| 205 | + |
| 206 | + ament_auto_add_library(yolo SHARED |
| 207 | + lib/src/trt_yolo.cpp |
| 208 | + ) |
| 209 | + |
| 210 | + target_include_directories(yolo PUBLIC |
| 211 | + lib/include |
| 212 | + ) |
| 213 | + |
| 214 | + target_link_libraries(yolo |
| 215 | + ${NVINFER} |
| 216 | + ${NVONNXPARSER} |
| 217 | + ${NVINFER_PLUGIN} |
| 218 | + ${CUDA_LIBRARIES} |
| 219 | + ${CUBLAS_LIBRARIES} |
| 220 | + ${CUDA_curand_LIBRARY} |
| 221 | + ${CUDNN_LIBRARY} |
| 222 | + mish_plugin |
| 223 | + yolo_layer_plugin |
| 224 | + nms_plugin |
| 225 | + ) |
| 226 | + |
| 227 | + ament_auto_add_library(tensorrt_yolo_nodelet SHARED |
| 228 | + src/nodelet.cpp |
| 229 | + ) |
| 230 | + |
| 231 | + target_link_libraries(tensorrt_yolo_nodelet |
| 232 | + ${OpenCV_LIBS} |
| 233 | + yolo |
| 234 | + mish_plugin |
| 235 | + yolo_layer_plugin |
| 236 | + nms_plugin |
| 237 | + ) |
| 238 | + |
| 239 | + rclcpp_components_register_node(tensorrt_yolo_nodelet |
| 240 | + PLUGIN "object_recognition::TensorrtYoloNodelet" |
| 241 | + EXECUTABLE tensorrt_yolo_node |
| 242 | + ) |
| 243 | + |
| 244 | + if(BUILD_TESTING) |
| 245 | + find_package(ament_lint_auto REQUIRED) |
| 246 | + ament_lint_auto_find_test_dependencies() |
| 247 | + endif() |
| 248 | + |
| 249 | + ament_auto_package(INSTALL_TO_SHARE |
| 250 | + config |
| 251 | + data |
| 252 | + launch |
| 253 | + ) |
| 254 | + |
| 255 | + install( |
| 256 | + TARGETS |
| 257 | + mish_plugin |
| 258 | + yolo_layer_plugin |
| 259 | + nms_plugin |
| 260 | + LIBRARY DESTINATION lib |
| 261 | + ARCHIVE DESTINATION lib |
| 262 | + RUNTIME DESTINATION bin |
| 263 | + ) |
| 264 | + |
| 265 | +else() |
| 266 | + message("TensorrtYolo won't be built, CUDA and/or TensorRT were not found.") |
| 267 | + ament_auto_package(INSTALL_TO_SHARE |
| 268 | + config |
| 269 | + data |
| 270 | + launch |
| 271 | + ) |
| 272 | +endif() |
0 commit comments