├── CMakeLists.txt ├── LICENSE ├── README.md ├── bindings ├── CMakeLists.txt └── python │ ├── CMakeLists.txt │ └── pybind.cpp ├── docs ├── demo.png └── demo_robotsports.png ├── examples ├── CMakeLists.txt ├── batch │ ├── CMakeLists.txt │ ├── README.md │ ├── process_batch.cpp │ └── process_batch.py ├── builder │ ├── CMakeLists.txt │ ├── README.md │ ├── build_engine.cpp │ └── build_engine.py ├── coco.txt ├── image │ ├── CMakeLists.txt │ ├── README.md │ ├── process_image.cpp │ └── process_image.py └── live │ ├── CMakeLists.txt │ ├── README.md │ ├── process_live.cpp │ └── process_live.py ├── include ├── yolov5_builder.hpp ├── yolov5_common.hpp ├── yolov5_detection.hpp ├── yolov5_detector.hpp ├── yolov5_detector_internal.hpp └── yolov5_logging.hpp ├── pyproject.toml ├── setup.py ├── src ├── CMakeLists.txt ├── yolov5_builder.cpp ├── yolov5_common.cpp ├── yolov5_detection.cpp ├── yolov5_detector.cpp ├── yolov5_detector_internal.cpp └── yolov5_logging.cpp └── yolov5-tensorrt.pc.in /CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.1) 2 | project(yolov5-tensorrt VERSION 0.1 DESCRIPTION "Real-time object detection with YOLOv5 and TensorRT") 3 | 4 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O2 -g -Wall -Wextra -Wno-deprecated -fPIC") 5 | set(CMAKE_CXX_STANDARD 14) 6 | set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} -g -pthread") 7 | 8 | set(CMAKE_INCLUDE_CURRENT_DIR ON) 9 | 10 | set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}) 11 | if(NOT DEFINED CMAKE_LIBRARY_OUTPUT_DIRECTORY) 12 | set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}) 13 | endif() 14 | 15 | option(BUILD_PYTHON "Build the Python bindings" OFF) 16 | 17 | 18 | ## Pkg-Config support when installing 19 | include(GNUInstallDirs) 20 | configure_file(yolov5-tensorrt.pc.in yolov5-tensorrt.pc @ONLY) 21 | install(FILES ${CMAKE_BINARY_DIR}/yolov5-tensorrt.pc 22 | DESTINATION ${CMAKE_INSTALL_DATAROOTDIR}/pkgconfig) 23 | 24 | 25 | ## Libraries 26 | ##################### 27 | 28 | find_package(OpenCV REQUIRED) 29 | find_package(CUDA REQUIRED) 30 | 31 | 32 | if(BUILD_PYTHON) 33 | # use pybind11 to build the Python bindings 34 | find_package(pybind11 REQUIRED) 35 | endif() 36 | 37 | 38 | ## Sources 39 | ##################### 40 | 41 | include_directories(include) 42 | file(GLOB_RECURSE YOLOV5_INCLUDE_FILES "include/*.hpp*") 43 | 44 | add_subdirectory(src) 45 | 46 | add_subdirectory(bindings) 47 | 48 | ## Examples 49 | ##################### 50 | 51 | add_subdirectory(examples) 52 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021, Noah van der Meer 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to 7 | deal in the Software without restriction, including without limitation the 8 | rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 9 | sell copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in 13 | all copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 20 | FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 21 | IN THE SOFTWARE. -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # YOLOv5-TensorRT 2 | 3 | ![](./docs/demo.png) 4 | 5 | The goal of this library is to provide an accessible and robust method for performing efficient, real-time object detection with [YOLOv5](https://github.com/ultralytics/yolov5) using NVIDIA TensorRT. The library was developed with real-world deployment and robustness in mind. Moreover, the library is extensively documented and comes with various guided examples. 6 | 7 |
8 | 9 | [![platform](https://img.shields.io/badge/platform-linux%20%7C%20L4T-orange?style=for-the-badge)](https://github.com/noahmr/yolov5-tensorrt#install) 10 | [![api](https://img.shields.io/badge/api-c++%20%7C%20Python-blue?style=for-the-badge)](https://github.com/noahmr/yolov5-tensorrt#usage) 11 | [![license](https://img.shields.io/badge/license-MIT-green?style=for-the-badge)](LICENSE) 12 |
13 | 14 | ##
Features
15 | 16 | - C++ and Python API 17 | - FP32 and FP16 inference 18 | - Batch inference 19 | - Support for varying input dimensions 20 | - ONNX support 21 | - CUDA-accelerated pre-processing 22 | - Integration with OpenCV (with optionally also the OpenCV-CUDA module) 23 | - Modular logging and error reporting 24 | - Extensive documentation available on all classes, methods and functions 25 | 26 | 27 | ##
Platforms
28 | 29 |
30 | Platforms 31 | 32 | - Modern Linux distros 33 | - NVIDIA L4T (Jetson platform) 34 |
35 | 36 |
37 | Dependencies 38 | 39 | - TensorRT >=8 (including libnvonnxparsers-dev) 40 | - CUDA >= 10.2 41 | - OpenCV 42 | - Pybind11 (optional, for Python API) 43 | 44 |
45 | 46 | 47 | ##
Install
48 | 49 |
50 | Build/Install with Pip (for just Python) 51 | 52 | Ensure that at least the TensorRT, CUDA and OpenCV dependencies mentioned above are installed on your system, the rest can be handled by Pip. You can install the library using: 53 | ```bash 54 | pip3 install . 55 | ``` 56 | This will build and install the only the Python API, as well as the example scripts. 57 |
58 | 59 | 60 | 61 |
62 | Build/Install with CMake (for C++ and Python) 63 | 64 | Configure the build with CMake: 65 | 66 | ```bash 67 | mkdir build 68 | cd build 69 | cmake .. -DBUILD_PYTHON=OFF 70 | ``` 71 | If wish to also build and install the Python API, you can instead set ```-DBUILD_PYTHON=ON```. Next, build and install using 72 | ``` 73 | make 74 | sudo make install 75 | ``` 76 | This will build and install all of the example applications, as well as a shared library yolov5-tensorrt.so. 77 | 78 |
79 | 80 | 81 | 82 | ##
Usage
83 | 84 |
85 | Command-line Usage 86 | 87 | The library also comes with various tools/demos. If your YOLOv5 model is stored as yolov5.onnx, you can build a TensorRT engine using: 88 | 89 | 90 | ```bash 91 | build_engine --model yolov5.onnx --output yolov5.engine 92 | ``` 93 | The resulting engine will be stored to disk at yolov5.engine. See [build_engine](examples/builder) for more information. 94 | 95 | After the engine has been stored, you can load it and detect objects as following: 96 | ```bash 97 | process_image --engine yolov5.engine --input image.png --output result.png 98 | ``` 99 | A visualization of the result will be stored to disk at result.png. See [process_image](examples/image) for more information. 100 | 101 |
102 | 103 |
104 | C++ usage 105 | 106 | Include ```yolov5_builder.hpp``` in your code. If your YOLOv5 model is stored as yolov5.onnx, you can build the TensorRT engine using three lines of C++ code: 107 | 108 | ```cpp 109 | yolov5::Builder builder; 110 | builder.init(); 111 | builder.build("yolov5.onnx", "yolov5.engine"); 112 | ``` 113 | 114 | 115 | For detection, include ```yolov5_detector.hpp``` in your code. You can detect objects with the following code: 116 | 117 | ```cpp 118 | yolov5::Detector detector; 119 | detector.init(); 120 | detector.loadEngine("yolov5.engine"); 121 | 122 | cv::Mat image = cv::imread("image.png"); 123 | 124 | std::vector detections; 125 | detector.detect(image, &detections); 126 | ``` 127 | 128 |
129 | 130 |
131 | Python usage 132 | 133 | Import the ```yolov5tensorrt``` package in your code. If your YOLOv5 model is stored as yolov5.onnx, you can build the TensorRT engine using three lines of Python code: 134 | 135 | ```python 136 | builder = yolov5tensorrt.Builder() 137 | builder.init() 138 | builder.build("yolov5.onnx", "yolov5.engine") 139 | ``` 140 | 141 | Next, you can detect objects using with the following code: 142 | 143 | ```python 144 | detector = yolov5tensorrt.Detector() 145 | detector.init() 146 | detector.loadEngine("yolov5.engine") 147 | 148 | image = cv2.imread("image.png") 149 | 150 | r, detections = detector.detect(image) 151 | ``` 152 | 153 |
154 | 155 | 156 |
157 | Examples 158 | 159 | Various **documented** examples can be found in the [examples](examples) directory. 160 | 161 | In order to **build** a TensorRT engine based on an ONNX model, the following 162 | tool/example is available: 163 | - [build_engine](examples/builder) (C++/Python): build a TensorRT engine based on your ONNX model 164 | 165 | For **object detection**, the following tools/examples are available: 166 | - [process_image](examples/image) (C++/Python): detect objects in a single image 167 | - [process_live](examples/live) (C++/Python): detect objects live in a video stream (e.g. webcam) 168 | - [process_batch](examples/batch) (C++/Python): detect objects in multiple images (batch inference) 169 | 170 |
171 | 172 | 173 |
174 | Importing the library in your project: CMake 175 | 176 | After installing the library, you can include it in your CMake-based project through pkg-config using the following: 177 | ``` 178 | find_package(PkgConfig REQUIRED) 179 | pkg_check_modules(YOLOV5_TENSORRT yolov5-tensorrt) 180 | ``` 181 | This will provide the usual ```YOLOV5_TENSORRT_INCLUDE_DIRS```, ```YOLOV5_TENSORRT_LIBRARIES``` and ```YOLOV5_TENSORRT_VERSION``` variables in CMake. 182 |
183 | 184 | 185 | 186 |
187 | Importing the library in your project: pkg-config 188 | 189 | After installing the library, in order to use the library in your own project, you can include and link it in the usual manner through [pkg-config](https://www.freedesktop.org/wiki/Software/pkg-config/). To get the include directories of the library, use: 190 | 191 | ``` 192 | pkg-config --cflags yolov5-tensorrt 193 | ``` 194 | and similarly for linking: 195 | 196 | ``` 197 | pkg-config --libs yolov5-tensorrt 198 | ``` 199 |
200 | 201 |
202 | Additional Resources 203 | 204 | - [Use with Stereolabs ZED](https://github.com/noahmr/zed-yolov5) 205 | - [AI at RobotSports (Kaggle)](https://www.kaggle.com/charel/yolov5-1st-place-world-championships-robocup-2021) 206 |
207 | 208 | 209 | 210 | ##
About
211 | 212 | This library is developed at [VDL RobotSports](https://robotsports.nl), 213 | an industrial team based in the Netherlands participating in the RoboCup Middle 214 | Size League, and currently sees active use on the soccer robots. 215 | 216 |
217 | RobotSports Demo 218 | 219 | ![](./docs/demo_robotsports.png) 220 | 221 |
222 | 223 |
224 | Citing 225 | 226 | If you like this library and would like to cite it, please use the following (LateX): 227 | 228 | ```tex 229 | @misc{yolov5tensorrt, 230 | author = {van der Meer, Noah and van Hoof, Charel}, 231 | title = {{yolov5-tensorrt}: Real-time object detection with {YOLOv5} and {TensorRT}}, 232 | howpublished = {GitHub}, 233 | year = {2021}, 234 | note = {\url{https://github.com/noahmr/yolov5-tensorrt}} 235 | } 236 | ``` 237 |
238 | 239 | 240 | ##
License
241 | 242 | Copyright (c) 2021, Noah van der Meer 243 | 244 | This software is licenced under the MIT license, which can be found in [LICENCE.md](LICENCE.md). By using, distributing, or contributing to this project, you agree to the terms and conditions of this license. 245 | -------------------------------------------------------------------------------- /bindings/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | if(pybind11_FOUND) 2 | 3 | add_subdirectory(python) 4 | 5 | endif() -------------------------------------------------------------------------------- /bindings/python/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | pybind11_add_module(yolov5tensorrt pybind.cpp) 2 | target_include_directories(yolov5tensorrt PUBLIC 3 | ${OpenCV_INCLUDE_DIRS} 4 | ${CUDA_INCLUDE_DIRS} 5 | ) 6 | target_link_libraries(yolov5tensorrt PRIVATE 7 | yolov5-tensorrt 8 | nvinfer 9 | nvonnxparser 10 | ${CUDA_CUDART_LIBRARY} 11 | ${OpenCV_LIBRARIES} 12 | ) 13 | 14 | # Determine the appropriate location to install the library. In 15 | # newer versions of CMake this can be done through FindPython, but it is 16 | # not yet widely available. 17 | # 18 | # See https://cmake.org/cmake/help/latest/module/FindPython.html 19 | execute_process( 20 | COMMAND python3 -c "from distutils.sysconfig import get_python_lib; print(get_python_lib())" 21 | OUTPUT_VARIABLE YOLOV5_PYTHON_SITE_PATH 22 | OUTPUT_STRIP_TRAILING_WHITESPACE 23 | ) 24 | 25 | # install rule 26 | install(TARGETS yolov5tensorrt 27 | LIBRARY DESTINATION ${YOLOV5_PYTHON_SITE_PATH}) -------------------------------------------------------------------------------- /bindings/python/pybind.cpp: -------------------------------------------------------------------------------- 1 | /** 2 | * @file 3 | * 4 | * @author Noah van der Meer 5 | * @brief Python Bindings for the yolov5-tensorrt library 6 | * 7 | * 8 | * Copyright (c) 2021, Noah van der Meer 9 | * 10 | * Permission is hereby granted, free of charge, to any person obtaining a copy 11 | * of this software and associated documentation files (the "Software"), to 12 | * deal in the Software without restriction, including without limitation the 13 | * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 14 | * sell copies of the Software, and to permit persons to whom the Software is 15 | * furnished to do so, subject to the following conditions: 16 | * 17 | * The above copyright notice and this permission notice shall be included in 18 | * all copies or substantial portions of the Software. 19 | * 20 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 21 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 22 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 23 | * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 24 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 25 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 26 | * IN THE SOFTWARE. 27 | * 28 | */ 29 | 30 | #include 31 | #include 32 | #include 33 | 34 | #include "yolov5_builder.hpp" 35 | #include "yolov5_detector.hpp" 36 | 37 | using namespace yolov5; 38 | 39 | inline cv::Mat arrayToMat(const pybind11::array_t& img) 42 | { 43 | const auto rows = img.shape(0); 44 | const auto columns = img.shape(1); 45 | const auto channels = img.shape(2); 46 | return cv::Mat(rows, columns, CV_8UC(channels), 47 | (unsigned char*)img.data()); 48 | } 49 | 50 | PYBIND11_MODULE(yolov5tensorrt, m) 51 | { 52 | m.doc() = "yolov5-tensorrt python binding"; 53 | 54 | 55 | /* 56 | yolov5_common.hpp 57 | */ 58 | pybind11::enum_(m, "Result") 59 | .value("FAILURE_INVALID_INPUT", RESULT_FAILURE_INVALID_INPUT) 60 | .value("FAILURE_NOT_INITIALIZED", 61 | RESULT_FAILURE_NOT_INITIALIZED) 62 | .value("FAILURE_NOT_LOADED", RESULT_FAILURE_NOT_LOADED) 63 | .value("FAILURE_MODEL_ERROR", RESULT_FAILURE_MODEL_ERROR) 64 | .value("FAILURE_OPENCV_NO_CUDA", RESULT_FAILURE_OPENCV_NO_CUDA) 65 | .value("FAILURE_FILESYSTEM_ERROR", 66 | RESULT_FAILURE_FILESYSTEM_ERROR) 67 | .value("FAILURE_CUDA_ERROR", RESULT_FAILURE_CUDA_ERROR) 68 | .value("FAILURE_TENSORRT_ERROR", RESULT_FAILURE_TENSORRT_ERROR) 69 | .value("FAILURE_OPENCV_ERROR", RESULT_FAILURE_OPENCV_ERROR) 70 | .value("FAILURE_ALLOC", RESULT_FAILURE_ALLOC) 71 | .value("FAILURE_OTHER", RESULT_FAILURE_OTHER) 72 | .value("SUCCESS", RESULT_SUCCESS); 73 | m.def("result_to_string", 74 | pybind11::overload_cast(&result_to_string), 75 | "Get a textual description of a result code"); 76 | 77 | pybind11::enum_(m, "Precision") 78 | .value("FP32", PRECISION_FP32) 79 | .value("FP16", PRECISION_FP16); 80 | m.def("precision_to_string", 81 | pybind11::overload_cast(&precision_to_string), 82 | "Get a textual description of a precision code"); 83 | 84 | pybind11::enum_(m, "DetectorFlag") 85 | .value("INPUT_BGR", INPUT_BGR) 86 | .value("INPUT_RGB", INPUT_RGB) 87 | .value("PREPROCESSOR_CVCUDA", PREPROCESSOR_CVCUDA) 88 | .value("PREPROCESSOR_CVCPU", PREPROCESSOR_CVCPU); 89 | 90 | 91 | /* 92 | yolov5_builder.hpp 93 | */ 94 | pybind11::class_(m, "Builder") 95 | .def(pybind11::init<>()) 96 | .def("init", &Builder::init, "Initialize the Builder.") 97 | .def("buildEngine", 98 | [](const Builder& builder, 99 | const std::string& inputPath, const std::string& outputPath, 100 | Precision precision) 101 | { 102 | return builder.buildEngine(inputPath, outputPath, precision); 103 | }, 104 | pybind11::arg("inputPath"), pybind11::arg("outputPath"), 105 | pybind11::arg("precision") = PRECISION_FP32, 106 | "Build an engine from ONNX model input, save it to disk" 107 | ); 108 | 109 | 110 | /* 111 | yolov5_detection.hpp 112 | */ 113 | pybind11::class_(m, "Detection") 114 | .def(pybind11::init<>()) 115 | .def("classId", &Detection::classId, 116 | "Retrieve the class id of the detection") 117 | .def("boundingBox", 118 | [](const Detection& detection) 119 | { 120 | const cv::Rect& r = detection.boundingBox(); 121 | return pybind11::make_tuple(r.x, r.y, r.width, r.height); 122 | }, 123 | "Retrieve a bounding box of the detection" 124 | ) 125 | .def("score", &Detection::score, 126 | "Retrieve the score assigned to this detection") 127 | .def("className", &Detection::className, 128 | "Retrieve the name of the class of this detection, if known.") 129 | .def("setClassName", &Detection::setClassName, 130 | "Set the class name"); 131 | 132 | m.def("visualizeDetection", 133 | [](const Detection& detection, 134 | pybind11::array_t& img, 136 | const std::tuple& color, 137 | const double& fontScale) 138 | { 139 | const cv::Scalar colorScalar(std::get<0>(color), 140 | std::get<1>(color), std::get<2>(color)); 141 | cv::Mat mat = arrayToMat(img); 142 | return visualizeDetection(detection, &mat, colorScalar, fontScale); 143 | }, 144 | pybind11::arg("detection"), pybind11::arg("img"), 145 | pybind11::arg("color"), pybind11::arg("fontScale"), 146 | "Helper method for visualizing a Detection in an image" 147 | ); 148 | 149 | pybind11::class_(m, "Classes") 150 | .def(pybind11::init<>()) 151 | .def("loadFromFile", &Classes::loadFromFile, 152 | "Try loading the class names as a list from a file") 153 | .def("isLoaded", &Classes::isLoaded, 154 | "Query whether the classes have been loaded") 155 | .def("getName", 156 | [](const Classes& classes, const int& classId) 157 | { 158 | std::string name; 159 | const Result r = classes.getName(classId, &name); 160 | return pybind11::make_tuple(r, name); 161 | }, 162 | "Get the Class name corresponding to a ClassId" 163 | ); 164 | 165 | 166 | /* 167 | yolov5_detector.hpp 168 | */ 169 | pybind11::class_(m, "Detector") 170 | .def(pybind11::init<>()) 171 | .def("init", &Detector::init, 172 | pybind11::arg("flags") = 0, 173 | "Initialize the Detector.") 174 | .def("isInitialized", &Detector::isInitialized, 175 | "Query whether the Detector is initialized") 176 | .def("loadEngine", 177 | [](Detector& detector, const std::string& filepath) 178 | { 179 | return detector.loadEngine(filepath); 180 | }, 181 | "Load a TensorRT engine from a file" 182 | ) 183 | .def("isEngineLoaded", &Detector::isEngineLoaded, 184 | "Query whether an inference engine has been loaded already") 185 | .def("numClasses", &Detector::numClasses, 186 | "Retrieve the number of classes of the engine/network") 187 | .def("setClasses", &Detector::setClasses, 188 | "Set the classes of the network") 189 | .def("detect", 190 | [](Detector& detector, 191 | const pybind11::array_t& img, 193 | int flags = 0) 194 | { 195 | std::vector detections; 196 | const Result r = 197 | detector.detect(arrayToMat(img), &detections, flags); 198 | return pybind11::make_tuple(r, detections); 199 | }, 200 | pybind11::arg("img"), 201 | pybind11::arg("flags") = 0, 202 | "Detect objects in the specified image using the YoloV5 model" 203 | ) 204 | .def("detectBatch", 205 | [](Detector& detector, 206 | const std::vector< 207 | pybind11::array_t>& images, 209 | int flags = 0) 210 | { 211 | std::vector> detections; 212 | 213 | std::vector mats; 214 | for(unsigned int i = 0; i < images.size(); ++i) 215 | { 216 | mats.push_back(arrayToMat(images[i])); 217 | } 218 | const Result r = 219 | detector.detectBatch(mats, &detections, flags); 220 | return pybind11::make_tuple(r, detections); 221 | }, 222 | pybind11::arg("images"), 223 | pybind11::arg("flags") = 0, 224 | "Detect objects in the specified images using batch inference " 225 | "with the YoloV5 model" 226 | ) 227 | .def("scoreThreshold", &Detector::scoreThreshold, 228 | "Obtain the score threshold") 229 | .def("setScoreThreshold", &Detector::setScoreThreshold, 230 | "Set the Score threshold: used to filter objects by score") 231 | .def("nmsThreshold", &Detector::nmsThreshold, 232 | "Obtain the NMS threshold") 233 | .def("setNmsThreshold", &Detector::setNmsThreshold, 234 | "Set the NMS threshold") 235 | .def("batchSize", &Detector::batchSize, 236 | "Retrieve the batch size of the engine/network") 237 | .def("inferenceSize", 238 | [](Detector& detector) 239 | { 240 | const cv::Size s = detector.inferenceSize(); 241 | return pybind11::make_tuple(s.width, s.height); 242 | }, 243 | "Retrieve the input size for which the network was configured"); 244 | } 245 | -------------------------------------------------------------------------------- /docs/demo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/noahmr/yolov5-tensorrt/6af93611fa275b4a4ff09a4c1fdbe717b433e4a1/docs/demo.png -------------------------------------------------------------------------------- /docs/demo_robotsports.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/noahmr/yolov5-tensorrt/6af93611fa275b4a4ff09a4c1fdbe717b433e4a1/docs/demo_robotsports.png -------------------------------------------------------------------------------- /examples/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | add_subdirectory(builder) 2 | add_subdirectory(image) 3 | add_subdirectory(batch) 4 | add_subdirectory(live) 5 | -------------------------------------------------------------------------------- /examples/batch/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # 2 | # process_batch 3 | # 4 | 5 | add_executable(process_batch 6 | process_batch.cpp 7 | ) 8 | target_include_directories(process_batch PUBLIC 9 | ${OpenCV_INCLUDE_DIRS} 10 | ${CUDA_INCLUDE_DIRS} 11 | ) 12 | target_link_libraries(process_batch 13 | yolov5-tensorrt 14 | nvinfer 15 | nvonnxparser 16 | ${CUDA_CUDART_LIBRARY} 17 | ${OpenCV_LIBRARIES} 18 | ) 19 | 20 | install(TARGETS process_batch 21 | DESTINATION ${CMAKE_INSTALL_BINDIR}) -------------------------------------------------------------------------------- /examples/batch/README.md: -------------------------------------------------------------------------------- 1 | ## process_batch 2 | 3 | The ```process_batch``` tool loads a YOLOv5 TensorRT engine, and performs object detection on a batch of images at once, i.e. batch inference. Note that the YOLOv5 model should have been exported with an explicit batch size. 4 | 5 | Basic usage: 6 | ``` 7 | ./process_batch --engine ENGINE_FILE --inputs INPUT_DIRECTORY --outputs OUTPUT_DIRECTORY 8 | ``` 9 | 10 | Arguments: 11 | - ```--engine```: path to the YOLOv5 TensorRT engine 12 | - ```--inputs```: path to the directory in which the images are stored 13 | - ```--outputs```: path to the directory in which the results should be written (directory should exist already) 14 | - ```--classes```: (optional) path to a file containing the class names 15 | 16 | 17 | ### Class names 18 | 19 | By default, the ```process_batch``` program will attach numbers representing the class to all of the detections. For instance, when using the COCO dataset, this ranges from 0 to 79. By specifying the class names corresponding to each class id, human-readable names (e.g. "car", "truck") are displayed instead. 20 | 21 | For convenience, a file containing the class names for COCO dataset can be found [here](../coco.txt). 22 | 23 | ### Example Usage 24 | 25 | Given that your YOLOv5 TensorRT engine is yolov5s.engine and your input images are stored in directory input_images, you can detect objects using: 26 | ``` 27 | ./process_batch --engine yolov5s.engine --inputs input_images --outputs output_images 28 | ``` 29 | Visualizations will be stored to disk in the directory output_images. 30 | -------------------------------------------------------------------------------- /examples/batch/process_batch.cpp: -------------------------------------------------------------------------------- 1 | /** 2 | * @file 3 | * 4 | * @author Noah van der Meer 5 | * @brief YoloV5-TensorRT example: inference on a batch of images 6 | * 7 | * Copyright (c) 2021, Noah van der Meer 8 | * 9 | * Permission is hereby granted, free of charge, to any person obtaining a copy 10 | * of this software and associated documentation files (the "Software"), to 11 | * deal in the Software without restriction, including without limitation the 12 | * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 13 | * sell copies of the Software, and to permit persons to whom the Software is 14 | * furnished to do so, subject to the following conditions: 15 | * 16 | * The above copyright notice and this permission notice shall be included in 17 | * all copies or substantial portions of the Software. 18 | * 19 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 20 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 21 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 22 | * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 23 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 24 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 25 | * IN THE SOFTWARE. 26 | * 27 | */ 28 | 29 | #include "yolov5_detector.hpp" 30 | 31 | #include 32 | #include 33 | 34 | #include 35 | 36 | char* getCmdOption(char** begin, char** end, const std::string& option) 37 | { 38 | /* From https://stackoverflow.com/questions/865668/parsing- 39 | command-line-arguments-in-c */ 40 | char** itr = std::find(begin, end, option); 41 | if(itr != end && ++itr != end) 42 | { 43 | return *itr; 44 | } 45 | return 0; 46 | } 47 | 48 | bool cmdOptionExists(char** begin, char** end, const std::string& option, 49 | bool value = false) 50 | { 51 | /* From https://stackoverflow.com/questions/865668/parsing- 52 | command-line-arguments-in-c */ 53 | char** itr = std::find(begin, end, option); 54 | if(itr == end) 55 | { 56 | return false; 57 | } 58 | if(value && itr == end-1) 59 | { 60 | std::cout << "Warning: option '" << option << "'" 61 | << " requires a value" << std::endl; 62 | return false; 63 | } 64 | return true; 65 | } 66 | 67 | bool listFiles(const std::string& directory, std::vector* out) 68 | { 69 | DIR* dir = opendir(directory.c_str()); 70 | if(dir == nullptr) 71 | { 72 | std::cout << "listFiles() failed, could not open directory: " 73 | << std::strerror(errno) << std::endl; 74 | return false; 75 | } 76 | 77 | struct dirent* i = readdir(dir); 78 | while(i != nullptr) 79 | { 80 | const std::string name(i->d_name); 81 | if(name != "." && name != "..") 82 | { 83 | std::cout << "Input image: " << name << std::endl; 84 | out->push_back(name); 85 | } 86 | i = readdir(dir); 87 | } 88 | return true; 89 | } 90 | 91 | void printHelp() 92 | { 93 | std::cout << "Options:\n" 94 | "-h --help : show this help menu\n" 95 | "--engine : [mandatory] specify the engine file\n" 96 | "--inputs : [mandatory] specify the input directory\n" 97 | "--outputs : [mandatory] specify the output directory\n" 98 | "--classes : [optional] specify list of class names\n\n" 99 | "Example usage:\n" 100 | "process_batch --engine yolov5s.engine --inputs input_dir " 101 | "--outputs output_dir" << std::endl; 102 | } 103 | 104 | int main(int argc, char* argv[]) 105 | { 106 | /* 107 | Handle arguments 108 | */ 109 | if(cmdOptionExists(argv, argv+argc, "--help") || 110 | cmdOptionExists(argv, argv+argc, "-h")) 111 | { 112 | printHelp(); 113 | return 0; 114 | } 115 | 116 | if(!cmdOptionExists(argv, argv+argc, "--engine", true) || 117 | !cmdOptionExists(argv, argv+argc, "--inputs", true) || 118 | !cmdOptionExists(argv, argv+argc, "--outputs", true)) 119 | { 120 | std::cout << "Missing mandatory argument" << std::endl; 121 | printHelp(); 122 | return 1; 123 | } 124 | const std::string engineFile(getCmdOption(argv, argv+argc, "--engine")); 125 | const std::string inputDir(getCmdOption(argv, argv+argc, "--inputs")); 126 | const std::string outputDir(getCmdOption(argv, argv+argc, "--outputs")); 127 | 128 | std::string classesFile; 129 | if(cmdOptionExists(argv, argv+argc, "--classes", true)) 130 | { 131 | classesFile = getCmdOption(argv, argv+argc, "--classes"); 132 | } 133 | 134 | 135 | /* 136 | Create the YoloV5 Detector object. 137 | */ 138 | yolov5::Detector detector; 139 | 140 | 141 | /* 142 | Initialize the YoloV5 Detector. 143 | */ 144 | yolov5::Result r = detector.init(); 145 | if(r != yolov5::RESULT_SUCCESS) 146 | { 147 | std::cout << "init() failed: " << yolov5::result_to_string(r) 148 | << std::endl; 149 | return 1; 150 | } 151 | 152 | 153 | /* 154 | Load the engine from file. 155 | */ 156 | r = detector.loadEngine(engineFile); 157 | if(r != yolov5::RESULT_SUCCESS) 158 | { 159 | std::cout << "loadEngine() failed: " << yolov5::result_to_string(r) 160 | << std::endl; 161 | return 1; 162 | } 163 | 164 | 165 | /* 166 | Load the Class names from file, and pass these on to the Detector 167 | */ 168 | if(classesFile.length() > 0) 169 | { 170 | yolov5::Classes classes; 171 | classes.setLogger(detector.logger()); 172 | r = classes.loadFromFile(classesFile); 173 | if(r != yolov5::RESULT_SUCCESS) 174 | { 175 | std::cout << "classes.loadFromFile() failed: " 176 | << yolov5::result_to_string(r) << std::endl; 177 | return 1; 178 | } 179 | detector.setClasses(classes); 180 | } 181 | 182 | 183 | /* 184 | List all files in the specified directory 185 | */ 186 | std::vector filenames; 187 | if(!listFiles(inputDir, &filenames)) 188 | { 189 | return 1; 190 | } 191 | std::cout << "Found " << filenames.size() << " files in specified input " 192 | << "directory" << std::endl; 193 | 194 | /* 195 | Load the images from disk and store in CPU memory. 196 | */ 197 | std::vector images; 198 | for(unsigned int i = 0; i < filenames.size(); ++i) 199 | { 200 | cv::Mat image; 201 | try 202 | { 203 | image = cv::imread(inputDir + "/" + filenames[i]); 204 | } 205 | catch(const std::exception& e) 206 | { 207 | std::cout << "Could not load image " << filenames[i] 208 | << " . Exception: " << e.what() << std::endl; 209 | return 1; 210 | } 211 | 212 | if(!image.empty()) 213 | { 214 | images.push_back(image); 215 | } 216 | else 217 | { 218 | std::cout << "Could not load image: empty image " 219 | << filenames[i] << std::endl; 220 | return 1; 221 | } 222 | } 223 | 224 | 225 | /* 226 | The first one/two runs of the engine typically take significantly 227 | longer. To get an accurate timing for inference, first do two 228 | runs. These can of course also be performed on other representative 229 | images 230 | */ 231 | detector.detectBatch(images, nullptr); 232 | detector.detectBatch(images, nullptr); 233 | 234 | 235 | auto ts = std::chrono::high_resolution_clock::now(); /* timing */ 236 | 237 | /* 238 | Detect objects in the images using the detectBatch(...) method. The 239 | detections are inserted into the 'detections' vector. 240 | */ 241 | std::vector> detections; 242 | r = detector.detectBatch(images, &detections, yolov5::INPUT_BGR); 243 | if(r != yolov5::RESULT_SUCCESS) 244 | { 245 | std::cout << "detectBatch() failed: " << yolov5::result_to_string(r) 246 | << std::endl; 247 | return 1; 248 | } 249 | 250 | /* timing */ 251 | auto duration = std::chrono::duration_cast( 252 | std::chrono::high_resolution_clock::now() - ts); 253 | std::cout << "detectBatch() took: " << duration.count() << "ms" 254 | << std::endl; 255 | 256 | 257 | /* 258 | Visualize all of the detections & store to disk 259 | */ 260 | cv::Mat visualization; 261 | const cv::Scalar magenta(255, 51, 153); /* BGR */ 262 | for(unsigned int i = 0; i < detections.size(); ++i) 263 | { 264 | images[i].copyTo(visualization); 265 | 266 | const std::vector& lst = detections[i]; 267 | for(unsigned int j = 0; j < lst.size(); ++j) 268 | { 269 | yolov5::visualizeDetection(lst[j], &visualization, magenta, 1.0); 270 | } 271 | 272 | /* 273 | Store the visualization to disk again. 274 | */ 275 | const std::string outputName = outputDir + "/" + filenames[i]; 276 | try 277 | { 278 | cv::imwrite(outputName, visualization); 279 | } 280 | catch(const std::exception& e) 281 | { 282 | std::cout << "Warning: could not save image '" << outputName 283 | << "'. Exception: " << e.what() << std::endl; 284 | } 285 | } 286 | 287 | return 0; 288 | } -------------------------------------------------------------------------------- /examples/batch/process_batch.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python 2 | 3 | # 4 | # Author: Noah van der Meer 5 | # Description: YoloV5-TensorRT example: inference on a batch of images 6 | # 7 | # 8 | # Copyright (c) 2021, Noah van der Meer 9 | # 10 | # Permission is hereby granted, free of charge, to any person obtaining a copy 11 | # of this software and associated documentation files (the "Software"), to 12 | # deal in the Software without restriction, including without limitation the 13 | # rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 14 | # sell copies of the Software, and to permit persons to whom the Software is 15 | # furnished to do so, subject to the following conditions: 16 | # 17 | # The above copyright notice and this permission notice shall be included in 18 | # all copies or substantial portions of the Software. 19 | # 20 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 21 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 22 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 23 | # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 24 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 25 | # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 26 | # IN THE SOFTWARE. 27 | # 28 | # 29 | 30 | # note: import cv2 _before_ yolov5tensorrt; Otherwise it may lead to 31 | # issues see https://github.com/opencv/opencv/issues/14884 32 | import cv2 33 | import yolov5tensorrt 34 | 35 | import argparse 36 | import time 37 | import os 38 | 39 | def main(args): 40 | 41 | # 42 | # Create the YoloV5 Detector object 43 | # 44 | detector = yolov5tensorrt.Detector() 45 | 46 | 47 | # 48 | # Initialize the YoloV5 Detector. 49 | # 50 | r = detector.init() 51 | if r != yolov5tensorrt.Result.SUCCESS: 52 | print("init() failed:", yolov5tensorrt.result_to_string(r)) 53 | return 1 54 | 55 | 56 | # 57 | # Load the engine from file. 58 | # 59 | r = detector.loadEngine(args.engine) 60 | if r != yolov5tensorrt.Result.SUCCESS: 61 | print("loadEngine() failed:", yolov5tensorrt.result_to_string(r)) 62 | return 1 63 | 64 | 65 | # 66 | # Load the Class names from file, and pass these on to the Detector 67 | # 68 | if args.classes is not None: 69 | classes = yolov5tensorrt.Classes() 70 | r = classes.loadFromFile(args.classes) 71 | if r != yolov5tensorrt.Result.SUCCESS: 72 | print("classes.loadFromFile() failed:", 73 | yolov5tensorrt.result_to_string(r)) 74 | return 1 75 | detector.setClasses(classes) 76 | 77 | 78 | # 79 | # List all files in the specified directory 80 | # 81 | filenames = os.listdir(args.inputs) 82 | print("Found", len(filenames), "files in specified input directory") 83 | 84 | 85 | images = [] 86 | for f in filenames: 87 | image = cv2.imread(args.inputs + "/" + f) 88 | if image is not None: 89 | images.append(image) 90 | else: 91 | print("Could not load file ", f) 92 | return 1 93 | 94 | 95 | # 96 | # The first one/two runs of the engine typically take significantly 97 | # longer. To get an accurate timing for inference, first do two 98 | # runs. These can of course also be performed on other representative 99 | # images 100 | # 101 | detector.detect(image) 102 | detector.detect(image) 103 | 104 | 105 | ts = time.perf_counter() 106 | 107 | # 108 | # Detect objects in the images using the detectBatch(...) method. 109 | # 110 | r, detections = detector.detectBatch(images) 111 | if r != yolov5tensorrt.Result.SUCCESS: 112 | print("detectBatch() failed:", yolov5tensorrt.result_to_string(r)) 113 | return 1 114 | 115 | # timing 116 | duration = time.perf_counter() - ts 117 | print("detectBatch() took:", duration*1000, "milliseconds") 118 | 119 | 120 | # 121 | # Visualize all of the detections & store to disk 122 | # 123 | magenta = (255, 51, 153) # BGR 124 | visualization = [] 125 | for i in range(0, len(detections)): 126 | img = images[i] 127 | 128 | lst = detections[i] 129 | for d in lst: 130 | yolov5tensorrt.visualizeDetection(d, img, magenta, 1.0) 131 | 132 | # 133 | # Store the visualization to disk again. 134 | # 135 | outputName = args.outputs + "/" + filenames[i] 136 | cv2.imwrite(outputName, img) 137 | 138 | return 0 139 | 140 | 141 | if __name__ == '__main__': 142 | # 143 | # Handle arguments 144 | # 145 | parser = argparse.ArgumentParser(add_help=True) 146 | parser.add_argument('--engine', 147 | required = True, 148 | dest ='engine', 149 | type = str, 150 | help = '[mandatory] specify the engine file') 151 | parser.add_argument('--inputs', 152 | required = True, 153 | dest ='inputs', 154 | type = str, 155 | help = '[mandatory] specify the input directory') 156 | parser.add_argument('--outputs', 157 | required = True, 158 | dest ='outputs', 159 | type = str, 160 | help = '[mandatory] specify the output directory') 161 | parser.add_argument('--classes', 162 | dest ='classes', 163 | type = str, 164 | help = '[optional] specify list of class names') 165 | args = parser.parse_args() 166 | 167 | main(args) -------------------------------------------------------------------------------- /examples/builder/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # 2 | # build_engine 3 | # 4 | 5 | add_executable(build_engine 6 | build_engine.cpp 7 | ) 8 | target_include_directories(build_engine PUBLIC 9 | ${CUDA_INCLUDE_DIRS} 10 | ) 11 | target_link_libraries(build_engine 12 | yolov5-tensorrt 13 | nvinfer 14 | nvonnxparser 15 | ${CUDA_CUDART_LIBRARY} 16 | ${OpenCV_LIBRARIES} 17 | ) 18 | 19 | install(TARGETS build_engine 20 | DESTINATION ${CMAKE_INSTALL_BINDIR}) -------------------------------------------------------------------------------- /examples/builder/README.md: -------------------------------------------------------------------------------- 1 | ## build_engine 2 | 3 | The ```build_engine``` tool takes the YOLOv5 model expressed in ONNX form, and creates a TensorRT engine based on it. This TensorRT engine is optimized for the 4 | specific hardware being used, and can be used repeadetly for object detection. Note that you can export your trained YOLOv5 model to ONNX using the [official YOLOv5 guide](https://github.com/ultralytics/yolov5/issues/251). 5 | 6 | Basic usage: 7 | ``` 8 | ./build_engine --model ONNX_MODEL --output OUTPUT_ENGINE 9 | ``` 10 | 11 | Arguments: 12 | - ```--model```: path to the input ONNX model 13 | - ```--output```: path at which the output engine should be written 14 | - ```--precision```: (optional) the precision that should be used for the engine. The available options are "fp32" and "fp16". This argument is optional, and if it is not specified, the default "fp32" is used 15 | 16 | 17 | ### Example Usage 18 | 19 | Assuming that your YOLOv5 model is yolov5s.onnx (i.e. the S variant of YOLOv5), you can build an engine with FP16 inference as following: 20 | ``` 21 | ./build_engine --model yolov5s.onnx --output yolov5s.engine --precision fp16 22 | ``` 23 | The resulting engine will be stored to disk as yolov5s.engine. After this, you may perform inference using one of the other tools, such as [process_image](examples/image). 24 | -------------------------------------------------------------------------------- /examples/builder/build_engine.cpp: -------------------------------------------------------------------------------- 1 | /** 2 | * @file 3 | * 4 | * @author Noah van der Meer 5 | * @brief YoloV5-TensorRT example: build a TensorRT engine 6 | * 7 | * Copyright (c) 2021, Noah van der Meer 8 | * 9 | * Permission is hereby granted, free of charge, to any person obtaining a copy 10 | * of this software and associated documentation files (the "Software"), to 11 | * deal in the Software without restriction, including without limitation the 12 | * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 13 | * sell copies of the Software, and to permit persons to whom the Software is 14 | * furnished to do so, subject to the following conditions: 15 | * 16 | * The above copyright notice and this permission notice shall be included in 17 | * all copies or substantial portions of the Software. 18 | * 19 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 20 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 21 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 22 | * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 23 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 24 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 25 | * IN THE SOFTWARE. 26 | * 27 | */ 28 | 29 | #include "yolov5_builder.hpp" 30 | 31 | #include 32 | #include 33 | 34 | 35 | char* getCmdOption(char** begin, char** end, const std::string& option) 36 | { 37 | /* From https://stackoverflow.com/questions/865668/parsing- 38 | command-line-arguments-in-c */ 39 | char** itr = std::find(begin, end, option); 40 | if(itr != end && ++itr != end) 41 | { 42 | return *itr; 43 | } 44 | return 0; 45 | } 46 | 47 | bool cmdOptionExists(char** begin, char** end, const std::string& option, 48 | bool value = false) 49 | { 50 | /* From https://stackoverflow.com/questions/865668/parsing- 51 | command-line-arguments-in-c */ 52 | char** itr = std::find(begin, end, option); 53 | if(itr == end) 54 | { 55 | return false; 56 | } 57 | if(value && itr == end-1) 58 | { 59 | std::cout << "Warning: option '" << option << "'" 60 | << " requires a value" << std::endl; 61 | return false; 62 | } 63 | return true; 64 | } 65 | 66 | void printHelp() 67 | { 68 | std::cout << "Options:\n" 69 | "-h --help : show this help menu\n" 70 | "--model : [mandatory] specify the ONNX model file\n" 71 | "--output : [mandatory] specify the engine output " 72 | "file\n" 73 | "--precision : [optional] specify the precision. " 74 | "Options: fp32, fp16\n\n" 75 | "Example usage:\n" 76 | "build_engine --model yolov5.onnx --output yolov5.engine " 77 | "--precision fp32" << std::endl; 78 | } 79 | 80 | int main(int argc, char* argv[]) 81 | { 82 | /* 83 | Handle arguments 84 | */ 85 | if(cmdOptionExists(argv, argv+argc, "--help") || 86 | cmdOptionExists(argv, argv+argc, "-h")) 87 | { 88 | printHelp(); 89 | return 0; 90 | } 91 | 92 | if(!cmdOptionExists(argv, argv+argc, "--model", true) || 93 | !cmdOptionExists(argv, argv+argc, "--output", true)) 94 | { 95 | std::cout << "Missing mandatory argument" << std::endl; 96 | printHelp(); 97 | return 1; 98 | } 99 | 100 | const std::string modelFile(getCmdOption(argv, argv+argc, "--model")); 101 | const std::string outputFile(getCmdOption(argv, argv+argc, "--output")); 102 | 103 | yolov5::Precision precision = yolov5::PRECISION_FP32; /* default */ 104 | if(cmdOptionExists(argv, argv+argc, "--precision", true)) 105 | { 106 | const std::string s = getCmdOption(argv, argv+argc, "--precision"); 107 | if(s == "fp32") 108 | { 109 | } 110 | else if(s == "fp16") 111 | { 112 | precision = yolov5::PRECISION_FP16; 113 | } 114 | else 115 | { 116 | std::cout << "Invalid precision specified: " << s << std::endl; 117 | printHelp(); 118 | return 1; 119 | } 120 | } 121 | 122 | 123 | 124 | /* 125 | Create the YoloV5 Builder object. 126 | */ 127 | yolov5::Builder builder; 128 | 129 | 130 | /* 131 | Initialize the YoloV5 Builder. This should be done first, before 132 | building the engine. 133 | 134 | The init() method (like most of the methods) returns a result code, 135 | of the type yolov5::Result. If initialization was successfull, this 136 | will be RESULT_SUCCESS. If unsuccessfull, it will be set to one of the 137 | error codes, and you can get a description through the 138 | yolov5::result_to_string() function. 139 | 140 | Note that the Builder also performs extensive logging itself, 141 | so in case of failure, you will see a more detailed description of 142 | the problem in the console output 143 | */ 144 | yolov5::Result r = builder.init(); 145 | if(r != yolov5::RESULT_SUCCESS) 146 | { 147 | std::cout << "init() failed: " << yolov5::result_to_string(r) 148 | << std::endl; 149 | return 1; 150 | } 151 | 152 | 153 | /* 154 | Build the TensorRT engine 155 | */ 156 | r = builder.buildEngine(modelFile, outputFile, precision); 157 | if(r != yolov5::RESULT_SUCCESS) 158 | { 159 | std::cout << "buildEngine() failed: " << yolov5::result_to_string(r) 160 | << std::endl; 161 | return 1; 162 | } 163 | 164 | return 0; 165 | } -------------------------------------------------------------------------------- /examples/builder/build_engine.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python 2 | 3 | # 4 | # Author: Noah van der Meer 5 | # Description: YoloV5-TensorRT example: build a TensorRT engine 6 | # 7 | # 8 | # Copyright (c) 2021, Noah van der Meer 9 | # 10 | # Permission is hereby granted, free of charge, to any person obtaining a copy 11 | # of this software and associated documentation files (the "Software"), to 12 | # deal in the Software without restriction, including without limitation the 13 | # rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 14 | # sell copies of the Software, and to permit persons to whom the Software is 15 | # furnished to do so, subject to the following conditions: 16 | # 17 | # The above copyright notice and this permission notice shall be included in 18 | # all copies or substantial portions of the Software. 19 | # 20 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 21 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 22 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 23 | # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 24 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 25 | # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 26 | # IN THE SOFTWARE. 27 | # 28 | # 29 | 30 | import yolov5tensorrt 31 | import argparse 32 | 33 | def main(args): 34 | 35 | precision = yolov5tensorrt.Precision.FP32 36 | if args.precision is not None: 37 | if args.precision == "fp32": 38 | precision = yolov5tensorrt.Precision.FP32 39 | elif args.precision == 'fp16': 40 | precision = yolov5tensorrt.Precision.FP16 41 | else: 42 | print("Invalid precision specified:", args.precision) 43 | return 1 44 | 45 | 46 | # 47 | # Create the YoloV5 Builder object. 48 | # 49 | builder = yolov5tensorrt.Builder() 50 | 51 | 52 | # 53 | # Initialize the YoloV5 Builder. This should be done first, before 54 | # loading the engine. 55 | # 56 | # The init() method (like most of the methods) returns a result code, 57 | # of the type yolov5tensorrt.Result. If initialization was successfull, this 58 | # will be Result.SUCCESS. If unsuccessfull, it will be set to one of the 59 | # error codes, and you can get a description through the 60 | # yolov5tensorrt.result_to_string() function. 61 | # 62 | # Note that the Builder also performs extensive logging itself, 63 | # so in case of failure, you will see a more detailed description of 64 | # the problem in the console output 65 | # 66 | r = builder.init() 67 | if r != yolov5tensorrt.Result.SUCCESS: 68 | print("init() failed:", yolov5tensorrt.result_to_string(r)) 69 | return 1 70 | 71 | 72 | # 73 | # Build the TensorRT engine 74 | # 75 | r = builder.buildEngine(args.model, args.output, precision) 76 | if r != yolov5tensorrt.Result.SUCCESS: 77 | print("buildEngine() failed:", yolov5tensorrt.result_to_string(r)) 78 | return 1 79 | 80 | 81 | return 0 82 | 83 | if __name__ == '__main__': 84 | # 85 | # Handle arguments 86 | # 87 | parser = argparse.ArgumentParser(add_help=True) 88 | parser.add_argument('--model', 89 | required = True, 90 | dest ='model', 91 | type = str, 92 | help = '[mandatory] specify the ONNX model file') 93 | parser.add_argument('--output', 94 | required = True, 95 | dest ='output', 96 | type = str, 97 | help = '[mandatory] specify the engine output file') 98 | parser.add_argument('--precision', 99 | dest ='precision', 100 | type = str, 101 | help = '[optional] specify the precision') 102 | args = parser.parse_args() 103 | 104 | main(args) -------------------------------------------------------------------------------- /examples/coco.txt: -------------------------------------------------------------------------------- 1 | person 2 | bicycle 3 | car 4 | motorbike 5 | aeroplane 6 | bus 7 | train 8 | truck 9 | boat 10 | traffic light 11 | fire hydrant 12 | stop sign 13 | parking meter 14 | bench 15 | bird 16 | cat 17 | dog 18 | horse 19 | sheep 20 | cow 21 | elephant 22 | bear 23 | zebra 24 | giraffe 25 | backpack 26 | umbrella 27 | handbag 28 | tie 29 | suitcase 30 | frisbee 31 | skis 32 | snowboard 33 | sports ball 34 | kite 35 | baseball bat 36 | baseball glove 37 | skateboard 38 | surfboard 39 | tennis racket 40 | bottle 41 | wine glass 42 | cup 43 | fork 44 | knife 45 | spoon 46 | bowl 47 | banana 48 | apple 49 | sandwich 50 | orange 51 | broccoli 52 | carrot 53 | hot dog 54 | pizza 55 | donut 56 | cake 57 | chair 58 | sofa 59 | pottedplant 60 | bed 61 | diningtable 62 | toilet 63 | tvmonitor 64 | laptop 65 | mouse 66 | remote 67 | keyboard 68 | cell phone 69 | microwave 70 | oven 71 | toaster 72 | sink 73 | refrigerator 74 | book 75 | clock 76 | vase 77 | scissors 78 | teddy bear 79 | hair drier 80 | toothbrush -------------------------------------------------------------------------------- /examples/image/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # 2 | # process_image 3 | # 4 | 5 | add_executable(process_image 6 | process_image.cpp 7 | ) 8 | target_include_directories(process_image PUBLIC 9 | ${OpenCV_INCLUDE_DIRS} 10 | ${CUDA_INCLUDE_DIRS} 11 | ) 12 | target_link_libraries(process_image 13 | yolov5-tensorrt 14 | nvinfer 15 | nvonnxparser 16 | ${CUDA_CUDART_LIBRARY} 17 | ${OpenCV_LIBRARIES} 18 | ) 19 | 20 | install(TARGETS process_image 21 | DESTINATION ${CMAKE_INSTALL_BINDIR}) -------------------------------------------------------------------------------- /examples/image/README.md: -------------------------------------------------------------------------------- 1 | ## process_image 2 | 3 | The ```process_image``` tool takes a YOLOv5 TensorRT engine, and performs object detection on on a single input image. The tool measures the time the overall detection process takes (i.e. pre-processing, inference, post-processing together) and prints the result. 4 | 5 | Basic usage: 6 | ``` 7 | ./process_image --engine ENGINE_FILE --input INPUT_IMAGE --output OUTPUT_IMAGE 8 | ``` 9 | 10 | Arguments: 11 | - ```--engine```: path to the YOLOv5 TensorRT engine 12 | - ```--input```: path to the input image 13 | - ```--output```: path at which the visualized output image should be written 14 | - ```--classes```: (optional) path to a file containing the class names 15 | 16 | 17 | ### Class names 18 | 19 | By default, the ```process_image``` will attach numbers representing the class to all of the detections. For instance, when using the COCO dataset, this ranges from 0 to 79. By specifying the class names corresponding to each class id, human-readable names (e.g. "car", "truck") are displayed instead. 20 | 21 | For convenience, a file containing the class names for COCO dataset can be found [here](../coco.txt). 22 | 23 | ### Example Usage 24 | 25 | Assuming that your YOLOv5 TensorRT engine is yolov5s.engine and your input image is image.png, you can detect objects using: 26 | ``` 27 | ./process_image --engine yolov5s.engine --input image.png --output result.png 28 | ``` 29 | A visualization will be stored to disk as result.png. 30 | -------------------------------------------------------------------------------- /examples/image/process_image.cpp: -------------------------------------------------------------------------------- 1 | /** 2 | * @file 3 | * 4 | * @author Noah van der Meer 5 | * @brief YoloV5-TensorRT example: inference on a single image 6 | * 7 | * Copyright (c) 2021, Noah van der Meer 8 | * 9 | * Permission is hereby granted, free of charge, to any person obtaining a copy 10 | * of this software and associated documentation files (the "Software"), to 11 | * deal in the Software without restriction, including without limitation the 12 | * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 13 | * sell copies of the Software, and to permit persons to whom the Software is 14 | * furnished to do so, subject to the following conditions: 15 | * 16 | * The above copyright notice and this permission notice shall be included in 17 | * all copies or substantial portions of the Software. 18 | * 19 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 20 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 21 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 22 | * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 23 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 24 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 25 | * IN THE SOFTWARE. 26 | * 27 | */ 28 | 29 | #include "yolov5_detector.hpp" 30 | 31 | #include 32 | #include 33 | 34 | char* getCmdOption(char** begin, char** end, const std::string& option) 35 | { 36 | /* From https://stackoverflow.com/questions/865668/parsing- 37 | command-line-arguments-in-c */ 38 | char** itr = std::find(begin, end, option); 39 | if(itr != end && ++itr != end) 40 | { 41 | return *itr; 42 | } 43 | return 0; 44 | } 45 | 46 | bool cmdOptionExists(char** begin, char** end, const std::string& option, 47 | bool value = false) 48 | { 49 | /* From https://stackoverflow.com/questions/865668/parsing- 50 | command-line-arguments-in-c */ 51 | char** itr = std::find(begin, end, option); 52 | if(itr == end) 53 | { 54 | return false; 55 | } 56 | if(value && itr == end-1) 57 | { 58 | std::cout << "Warning: option '" << option << "'" 59 | << " requires a value" << std::endl; 60 | return false; 61 | } 62 | return true; 63 | } 64 | 65 | void printHelp() 66 | { 67 | std::cout << "Options:\n" 68 | "-h --help : show this help menu\n" 69 | "--engine : [mandatory] specify the engine file\n" 70 | "--input : [mandatory] specify the input image file\n" 71 | "--output : [mandatory] specify the output image file\n" 72 | "--classes : [optional] specify list of class names\n\n" 73 | "Example usage:\n" 74 | "process_image --engine yolov5s.engine --input test_image.png " 75 | "--output result.png" << std::endl; 76 | } 77 | 78 | int main(int argc, char* argv[]) 79 | { 80 | /* 81 | Handle arguments 82 | */ 83 | if(cmdOptionExists(argv, argv+argc, "--help") || 84 | cmdOptionExists(argv, argv+argc, "-h")) 85 | { 86 | printHelp(); 87 | return 0; 88 | } 89 | 90 | if(!cmdOptionExists(argv, argv+argc, "--engine", true) || 91 | !cmdOptionExists(argv, argv+argc, "--input", true) || 92 | !cmdOptionExists(argv, argv+argc, "--output", true)) 93 | { 94 | std::cout << "Missing mandatory argument" << std::endl; 95 | printHelp(); 96 | return 1; 97 | } 98 | const std::string engineFile(getCmdOption(argv, argv+argc, "--engine")); 99 | const std::string inputFile(getCmdOption(argv, argv+argc, "--input")); 100 | const std::string outputFile(getCmdOption(argv, argv+argc, "--output")); 101 | 102 | std::string classesFile; 103 | if(cmdOptionExists(argv, argv+argc, "--classes", true)) 104 | { 105 | classesFile = getCmdOption(argv, argv+argc, "--classes"); 106 | } 107 | 108 | 109 | /* 110 | Create the YoloV5 Detector object. 111 | */ 112 | yolov5::Detector detector; 113 | 114 | 115 | /* 116 | Initialize the YoloV5 Detector. This should be done first, before 117 | loading the engine. 118 | 119 | The init() method (like most of the methods) returns a result code, 120 | of the type yolov5::Result. If initialization was successfull, this 121 | will be RESULT_SUCCESS. If unsuccessfull, it will be set to one of the 122 | error codes, and you can get a description through the 123 | yolov5_result_to_string() function. 124 | 125 | Note that the Detector also performs extensive logging itself, 126 | so in case of failure, you will see a more detailed description of 127 | the problem in the console output 128 | */ 129 | yolov5::Result r = detector.init(); 130 | if(r != yolov5::RESULT_SUCCESS) 131 | { 132 | std::cout << "init() failed: " << yolov5::result_to_string(r) 133 | << std::endl; 134 | return 1; 135 | } 136 | 137 | 138 | /* 139 | Load the engine from file. 140 | */ 141 | r = detector.loadEngine(engineFile); 142 | if(r != yolov5::RESULT_SUCCESS) 143 | { 144 | std::cout << "loadEngine() failed: " << yolov5::result_to_string(r) 145 | << std::endl; 146 | return 1; 147 | } 148 | 149 | 150 | /* 151 | Load the Class names from file, and pass these on to the Detector 152 | */ 153 | if(classesFile.length() > 0) 154 | { 155 | yolov5::Classes classes; 156 | classes.setLogger(detector.logger()); 157 | r = classes.loadFromFile(classesFile); 158 | if(r != yolov5::RESULT_SUCCESS) 159 | { 160 | std::cout << "classes.loadFromFile() failed: " 161 | << yolov5::result_to_string(r) << std::endl; 162 | return 1; 163 | } 164 | detector.setClasses(classes); 165 | } 166 | 167 | 168 | /* 169 | Load an image from disk and store it in CPU memory. 170 | 171 | Note that by default, OpenCV will represent the image in BGR 172 | format. 173 | */ 174 | cv::Mat image; 175 | try 176 | { 177 | image = cv::imread(inputFile); 178 | } 179 | catch(const std::exception& e) 180 | { 181 | std::cout << "Failed to load input image: " << e.what() << std::endl; 182 | return 1; /* fatal */ 183 | } 184 | if(image.empty()) 185 | { 186 | std::cout << "Failed to load input image" << std::endl; 187 | return 1; 188 | } 189 | 190 | 191 | /* 192 | The first one/two runs of the engine typically take significantly 193 | longer. To get an accurate timing for inference, first do two 194 | runs. These can of course also be performed on other representative 195 | images 196 | */ 197 | detector.detect(image, nullptr); 198 | detector.detect(image, nullptr); 199 | 200 | 201 | auto ts = std::chrono::high_resolution_clock::now(); /* timing */ 202 | 203 | /* 204 | Detect objects in the image using the detect(...) method. The 205 | detections are inserted into the 'detections' vector. 206 | 207 | The detect(...) method can optionally also take flags. Through these 208 | flags, the type of input image can be specified specified, e.g. BGR 209 | or RGB. By default, the detect(...) method assumes that the input 210 | is stored as BGR. Thus while not necessary in this case, for clarity 211 | we specifically specify that the input is BGR here. 212 | 213 | Note that the detect() method might also fail in some 214 | cases, which can be checked through the returned result code. 215 | */ 216 | std::vector detections; 217 | r = detector.detect(image, &detections, yolov5::INPUT_BGR); 218 | if(r != yolov5::RESULT_SUCCESS) 219 | { 220 | std::cout << "detect() failed: " << yolov5::result_to_string(r) 221 | << std::endl; 222 | return 1; 223 | } 224 | 225 | /* timing */ 226 | auto duration = std::chrono::duration_cast( 227 | std::chrono::high_resolution_clock::now() - ts); 228 | std::cout << "detect() took: " << duration.count() << "ms" << std::endl; 229 | 230 | 231 | /* 232 | Visualize all of the detections 233 | 234 | The detections are provided in the form of yolov5::Detection 235 | objects. These contain information regarding the location in the image, 236 | confidence, and class. 237 | */ 238 | const cv::Scalar magenta(255, 51, 153); /* BGR */ 239 | for(unsigned int i = 0; i < detections.size(); ++i) 240 | { 241 | yolov5::visualizeDetection(detections[i], &image, magenta, 1.0); 242 | } 243 | 244 | 245 | /* 246 | Store the visualization to disk again. 247 | */ 248 | try 249 | { 250 | cv::imwrite(outputFile, image); 251 | } 252 | catch(const std::exception& e) 253 | { 254 | std::cout << "Failed to write output image: " << e.what() << std::endl; 255 | } 256 | 257 | return 0; 258 | } -------------------------------------------------------------------------------- /examples/image/process_image.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python 2 | 3 | # 4 | # Author: Noah van der Meer 5 | # Description: YoloV5-TensorRT example: inference on a single image 6 | # 7 | # 8 | # Copyright (c) 2021, Noah van der Meer 9 | # 10 | # Permission is hereby granted, free of charge, to any person obtaining a copy 11 | # of this software and associated documentation files (the "Software"), to 12 | # deal in the Software without restriction, including without limitation the 13 | # rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 14 | # sell copies of the Software, and to permit persons to whom the Software is 15 | # furnished to do so, subject to the following conditions: 16 | # 17 | # The above copyright notice and this permission notice shall be included in 18 | # all copies or substantial portions of the Software. 19 | # 20 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 21 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 22 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 23 | # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 24 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 25 | # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 26 | # IN THE SOFTWARE. 27 | # 28 | # 29 | 30 | # note: import cv2 _before_ yolov5tensorrt; Otherwise it may lead to 31 | # issues see https://github.com/opencv/opencv/issues/14884 32 | import cv2 33 | import yolov5tensorrt 34 | 35 | import argparse 36 | import time 37 | 38 | def main(args): 39 | 40 | # 41 | # Create the YoloV5 Detector object 42 | # 43 | detector = yolov5tensorrt.Detector() 44 | 45 | 46 | # 47 | # Initialize the YoloV5 Detector. This should be done first, before 48 | # loading the engine. 49 | # 50 | # The init() method (like most of the methods) returns a result code, 51 | # of the type yolov5tensorrt.Result. If initialization was successfull, this 52 | # will be Result.SUCCESS. If unsuccessfull, it will be set to one of the 53 | # error codes, and you can get a description through the 54 | # yolov5tensorrt.result_to_string() function. 55 | # 56 | # Note that the Detector also performs extensive logging itself, 57 | # so in case of failure, you will see a more detailed description of 58 | # the problem in the console output 59 | # 60 | r = detector.init() 61 | if r != yolov5tensorrt.Result.SUCCESS: 62 | print("init() failed:", yolov5tensorrt.result_to_string(r)) 63 | return 1 64 | 65 | 66 | # 67 | # Load the engine from file. 68 | # 69 | r = detector.loadEngine(args.engine) 70 | if r != yolov5tensorrt.Result.SUCCESS: 71 | print("loadEngine() failed:", yolov5tensorrt.result_to_string(r)) 72 | return 1 73 | 74 | 75 | # 76 | # Load the Class names from file, and pass these on to the Detector 77 | # 78 | if args.classes is not None: 79 | classes = yolov5tensorrt.Classes() 80 | r = classes.loadFromFile(args.classes) 81 | if r != yolov5tensorrt.Result.SUCCESS: 82 | print("classes.loadFromFile() failed:", 83 | yolov5tensorrt.result_to_string(r)) 84 | return 1 85 | detector.setClasses(classes) 86 | 87 | 88 | # 89 | # Load an image from disk and store it in CPU memory. 90 | # 91 | # Note that by default, OpenCV will represent the image in BGR format. 92 | # 93 | image = cv2.imread(args.input) 94 | if image is None: 95 | print("Failed to load input image") 96 | return 1 97 | 98 | # 99 | # The first one/two runs of the engine typically take significantly 100 | # longer. To get an accurate timing for inference, first do two 101 | # runs. These can of course also be performed on other representative 102 | # images 103 | # 104 | detector.detect(image) 105 | detector.detect(image) 106 | 107 | 108 | ts = time.perf_counter() 109 | 110 | # 111 | # Detect objects in the image using the detect(...) method. 112 | # 113 | # The detect(...) method can optionally also take flags. Through these 114 | # flags, the type of input image can be specified specified, e.g. BGR 115 | # or RGB. By default, the detect(...) method assumes that the input 116 | # is stored as BGR. Thus while not necessary in this case, for clarity 117 | # we specifically specify that the input is BGR here. 118 | # 119 | # Note that the detect() method might also fail in some 120 | # cases, which can be checked through the returned result code. 121 | # 122 | r, detections = detector.detect(image, flags = yolov5tensorrt.DetectorFlag.INPUT_BGR) 123 | if r != yolov5tensorrt.Result.SUCCESS: 124 | print("detect() failed:", yolov5tensorrt.result_to_string(r)) 125 | return 1 126 | 127 | # timing 128 | duration = time.perf_counter() - ts 129 | print("detect() took:", duration*1000, "milliseconds") 130 | 131 | 132 | # 133 | # Visualize all of the detections 134 | # 135 | # The detections are provided in the form of yolov5tensorrt.Detection 136 | # objects. These contain information regarding the location in the image, 137 | # confidence, and class. 138 | # 139 | magenta = (255, 51, 153) # BGR 140 | for d in detections: 141 | yolov5tensorrt.visualizeDetection(d, image, magenta, 1.0) 142 | 143 | 144 | # 145 | # Store the visualization to disk again. 146 | # 147 | cv2.imwrite(args.output, image) 148 | 149 | return 0 150 | 151 | 152 | if __name__ == '__main__': 153 | # 154 | # Handle arguments 155 | # 156 | parser = argparse.ArgumentParser(add_help=True) 157 | parser.add_argument('--engine', 158 | required = True, 159 | dest ='engine', 160 | type = str, 161 | help = '[mandatory] specify the engine file') 162 | parser.add_argument('--input', 163 | required = True, 164 | dest ='input', 165 | type = str, 166 | help = '[mandatory] specify the input image file') 167 | parser.add_argument('--output', 168 | required = True, 169 | dest ='output', 170 | type = str, 171 | help = '[mandatory] specify the output image file') 172 | parser.add_argument('--classes', 173 | dest ='classes', 174 | type = str, 175 | help = '[optional] specify list of class names') 176 | args = parser.parse_args() 177 | 178 | main(args) -------------------------------------------------------------------------------- /examples/live/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # 2 | # process_live 3 | # 4 | 5 | add_executable(process_live 6 | process_live.cpp 7 | ) 8 | target_include_directories(process_live PUBLIC 9 | ${OpenCV_INCLUDE_DIRS} 10 | ${CUDA_INCLUDE_DIRS} 11 | ) 12 | target_link_libraries(process_live 13 | yolov5-tensorrt 14 | nvinfer 15 | nvonnxparser 16 | ${CUDA_CUDART_LIBRARY} 17 | ${OpenCV_LIBRARIES} 18 | ) 19 | 20 | install(TARGETS process_live 21 | DESTINATION ${CMAKE_INSTALL_BINDIR}) -------------------------------------------------------------------------------- /examples/live/README.md: -------------------------------------------------------------------------------- 1 | ## process_live 2 | 3 | The ```process_live``` tool takes a YOLOv5 TensorRT engine, and performs object detection on a live source such as a camera. The result is visualized through a graphical user interface (highgui). 4 | 5 | Basic usage: 6 | ``` 7 | ./process_live --engine ENGINE_FILE 8 | ``` 9 | 10 | Arguments: 11 | - ```--engine```: path to the YOLOv5 TensorRT engine 12 | - ```--camera```: (optional) camera index. The default is 0 13 | - ```--classes```: (optional) path to a file containing the class names 14 | 15 | 16 | ### Class names 17 | 18 | By default, the ```process_live``` will attach numbers representing the class to all of the detections. For instance, when using the COCO dataset, this ranges from 0 to 79. By specifying the class names corresponding to each class id, human-readable names (e.g. "car", "truck") are displayed instead. 19 | 20 | For convenience, a file containing the class names for COCO dataset can be found [here](../coco.txt). 21 | 22 | ### Example Usage 23 | 24 | Assuming that your YOLOv5 TensorRT engine is yolov5s.engine, you can detect objects using: 25 | ``` 26 | ./process_live --engine yolov5s.engine 27 | ``` 28 | A graphical user interface will show up, displaying the detections live. 29 | -------------------------------------------------------------------------------- /examples/live/process_live.cpp: -------------------------------------------------------------------------------- 1 | /** 2 | * @file 3 | * 4 | * @author Noah van der Meer 5 | * @brief YoloV5-TensorRT example: inference on a live video source 6 | * 7 | * Copyright (c) 2021, Noah van der Meer 8 | * 9 | * Permission is hereby granted, free of charge, to any person obtaining a copy 10 | * of this software and associated documentation files (the "Software"), to 11 | * deal in the Software without restriction, including without limitation the 12 | * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 13 | * sell copies of the Software, and to permit persons to whom the Software is 14 | * furnished to do so, subject to the following conditions: 15 | * 16 | * The above copyright notice and this permission notice shall be included in 17 | * all copies or substantial portions of the Software. 18 | * 19 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 20 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 21 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 22 | * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 23 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 24 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 25 | * IN THE SOFTWARE. 26 | * 27 | */ 28 | 29 | #include "yolov5_detector.hpp" 30 | 31 | #include 32 | #include 33 | 34 | #include 35 | #include 36 | 37 | char* getCmdOption(char** begin, char** end, const std::string& option) 38 | { 39 | /* From https://stackoverflow.com/questions/865668/parsing- 40 | command-line-arguments-in-c */ 41 | char** itr = std::find(begin, end, option); 42 | if(itr != end && ++itr != end) 43 | { 44 | return *itr; 45 | } 46 | return 0; 47 | } 48 | 49 | bool cmdOptionExists(char** begin, char** end, const std::string& option, 50 | bool value = false) 51 | { 52 | /* From https://stackoverflow.com/questions/865668/parsing- 53 | command-line-arguments-in-c */ 54 | char** itr = std::find(begin, end, option); 55 | if(itr == end) 56 | { 57 | return false; 58 | } 59 | if(value && itr == end-1) 60 | { 61 | std::cout << "Warning: option '" << option << "'" 62 | << " requires a value" << std::endl; 63 | return false; 64 | } 65 | return true; 66 | } 67 | 68 | void printHelp() 69 | { 70 | std::cout << "Options:\n" 71 | "-h --help : show this help menu\n" 72 | "--engine : [mandatory] specify the engine file\n" 73 | "--camera : [optional] camera index\n" 74 | "--classes : [optional] specify list of class names\n\n" 75 | "Example usage:\n" 76 | "process_live --engine yolov5s.engine --camera 0" << std::endl; 77 | } 78 | 79 | int main(int argc, char* argv[]) 80 | { 81 | /* 82 | Handle arguments 83 | */ 84 | if(cmdOptionExists(argv, argv+argc, "--help") || 85 | cmdOptionExists(argv, argv+argc, "-h")) 86 | { 87 | printHelp(); 88 | return 0; 89 | } 90 | 91 | if(!cmdOptionExists(argv, argv+argc, "--engine", true)) 92 | { 93 | std::cout << "Missing mandatory argument" << std::endl; 94 | printHelp(); 95 | return 1; 96 | } 97 | const std::string engineFile(getCmdOption(argv, argv+argc, "--engine")); 98 | 99 | int cameraIndex = 0; 100 | if(cmdOptionExists(argv, argv+argc, "--camera", true)) 101 | { 102 | const std::string option = 103 | getCmdOption(argv, argv+argc, "--camera"); 104 | cameraIndex = std::atoi(option.c_str()); 105 | } 106 | 107 | std::string classesFile; 108 | if(cmdOptionExists(argv, argv+argc, "--classes", true)) 109 | { 110 | classesFile = getCmdOption(argv, argv+argc, "--classes"); 111 | } 112 | 113 | 114 | /* 115 | Create the YoloV5 Detector object. 116 | */ 117 | yolov5::Detector detector; 118 | 119 | 120 | /* 121 | Initialize the YoloV5 Detector. This should be done first, before 122 | loading the engine. 123 | */ 124 | yolov5::Result r = detector.init(); 125 | if(r != yolov5::RESULT_SUCCESS) 126 | { 127 | std::cout << "init() failed: " << yolov5::result_to_string(r) 128 | << std::endl; 129 | return 1; 130 | } 131 | 132 | 133 | /* 134 | Load the engine from file. 135 | */ 136 | r = detector.loadEngine(engineFile); 137 | if(r != yolov5::RESULT_SUCCESS) 138 | { 139 | std::cout << "loadEngine() failed: " << yolov5::result_to_string(r) 140 | << std::endl; 141 | return 1; 142 | } 143 | 144 | 145 | /* 146 | Load the Class names from file, and pass these on to the Detector 147 | */ 148 | if(classesFile.length() > 0) 149 | { 150 | yolov5::Classes classes; 151 | classes.setLogger(detector.logger()); 152 | r = classes.loadFromFile(classesFile); 153 | if(r != yolov5::RESULT_SUCCESS) 154 | { 155 | std::cout << "classes.loadFromFile() failed: " 156 | << yolov5::result_to_string(r) << std::endl; 157 | return 1; 158 | } 159 | detector.setClasses(classes); 160 | } 161 | 162 | 163 | /* 164 | Set up the GUI 165 | */ 166 | cv::namedWindow("live"); 167 | 168 | 169 | /* 170 | Set up the Camera 171 | */ 172 | cv::VideoCapture capture; 173 | if(!capture.open(cameraIndex, cv::CAP_ANY)) 174 | { 175 | std::cout << "failure: could not open capture device" << std::endl; 176 | return 1; 177 | } 178 | 179 | /* 180 | Start Inference 181 | */ 182 | cv::Mat image; 183 | std::vector detections; 184 | while(true) 185 | { 186 | if(!capture.read(image)) 187 | { 188 | std::cout << "failure: could not read new frames" << std::endl; 189 | break; 190 | } 191 | 192 | r = detector.detect(image, &detections, yolov5::INPUT_BGR); 193 | if(r != yolov5::RESULT_SUCCESS) 194 | { 195 | std::cout << "detect() failed: " << yolov5::result_to_string(r) 196 | << std::endl; 197 | return 1; 198 | } 199 | 200 | /* 201 | Visualize the detections 202 | */ 203 | for(unsigned int i = 0; i < detections.size(); ++i) 204 | { 205 | const cv::Scalar magenta(255, 51, 153); /* BGR */ 206 | yolov5::visualizeDetection(detections[i], &image, magenta, 1.0); 207 | } 208 | cv::imshow("live", image); 209 | 210 | cv::waitKey(1); 211 | } 212 | capture.release(); 213 | 214 | cv::destroyAllWindows(); 215 | 216 | return 0; 217 | } -------------------------------------------------------------------------------- /examples/live/process_live.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python 2 | 3 | # 4 | # Author: Noah van der Meer 5 | # Description: YoloV5-TensorRT example: inference on a live video source 6 | # 7 | # 8 | # Copyright (c) 2020, Noah van der Meer 9 | # 10 | # Permission is hereby granted, free of charge, to any person obtaining a copy 11 | # of this software and associated documentation files (the "Software"), to 12 | # deal in the Software without restriction, including without limitation the 13 | # rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 14 | # sell copies of the Software, and to permit persons to whom the Software is 15 | # furnished to do so, subject to the following conditions: 16 | # 17 | # The above copyright notice and this permission notice shall be included in 18 | # all copies or substantial portions of the Software. 19 | # 20 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 21 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 22 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 23 | # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 24 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 25 | # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 26 | # IN THE SOFTWARE. 27 | # 28 | # 29 | 30 | # note: import cv2 _before_ yolov5tensorrt; Otherwise it may lead to 31 | # issues see https://github.com/opencv/opencv/issues/14884 32 | import cv2 33 | import yolov5tensorrt 34 | 35 | import argparse 36 | import time 37 | 38 | def main(args): 39 | 40 | # 41 | # Create the YoloV5 Detector object 42 | # 43 | detector = yolov5tensorrt.Detector() 44 | 45 | 46 | # 47 | # Initialize the YoloV5 Detector. This should be done first, before 48 | # loading the engine. 49 | # 50 | r = detector.init() 51 | if r != yolov5tensorrt.Result.SUCCESS: 52 | print("init() failed:", yolov5tensorrt.result_to_string(r)) 53 | return 1 54 | 55 | 56 | # 57 | # Load the engine from file. 58 | # 59 | r = detector.loadEngine(args.engine) 60 | if r != yolov5tensorrt.Result.SUCCESS: 61 | print("loadEngine() failed:", yolov5tensorrt.result_to_string(r)) 62 | return 1 63 | 64 | 65 | # 66 | # Load the Class names from file, and pass these on to the Detector 67 | # 68 | if args.classes is not None: 69 | classes = yolov5tensorrt.Classes() 70 | r = classes.loadFromFile(args.classes) 71 | if r != yolov5tensorrt.Result.SUCCESS: 72 | print("classes.loadFromFile() failed:", 73 | yolov5tensorrt.result_to_string(r)) 74 | return 1 75 | detector.setClasses(classes) 76 | 77 | 78 | # 79 | # Set up the GUI 80 | # 81 | cv2.namedWindow("live") 82 | 83 | 84 | # 85 | # Set up the Camera 86 | # 87 | capture = cv2.VideoCapture() 88 | if not capture.open(args.camera, cv2.CAP_ANY): 89 | print("failure: could not open capture device") 90 | return 1 91 | 92 | 93 | # 94 | # Start Inference 95 | # 96 | while True: 97 | 98 | ret, image = capture.read() 99 | if not ret: 100 | print("failure: could not read new frames") 101 | break 102 | 103 | r, detections = detector.detect(image, flags = yolov5tensorrt.DetectorFlag.INPUT_BGR) 104 | if r != yolov5tensorrt.Result.SUCCESS: 105 | print("detect() failed:", yolov5tensorrt.result_to_string(r)) 106 | return 1 107 | 108 | # 109 | # Visualize the detections 110 | # 111 | magenta = (255, 51, 153) # BGR 112 | for d in detections: 113 | yolov5tensorrt.visualizeDetection(d, image, magenta, 1.0) 114 | cv2.imshow("live", image) 115 | 116 | cv2.waitKey(1) 117 | 118 | capture.release() 119 | 120 | cv2.destroyAllWindows() 121 | 122 | return 0 123 | 124 | 125 | if __name__ == '__main__': 126 | # 127 | # Handle arguments 128 | # 129 | parser = argparse.ArgumentParser(add_help=True) 130 | parser.add_argument('--engine', 131 | required = True, 132 | dest ='engine', 133 | type = str, 134 | help = '[mandatory] specify the engine file') 135 | parser.add_argument('--camera', 136 | dest ='camera', 137 | type = int, 138 | default = 0, 139 | help = '[optional] camera index') 140 | parser.add_argument('--classes', 141 | dest ='classes', 142 | type = str, 143 | help = '[optional] specify list of class names') 144 | args = parser.parse_args() 145 | 146 | main(args) -------------------------------------------------------------------------------- /include/yolov5_builder.hpp: -------------------------------------------------------------------------------- 1 | /** 2 | * @file 3 | * 4 | * @author Noah van der Meer 5 | * @brief YoloV5 inference through NVIDIA TensorRT (builder) 6 | * 7 | * 8 | * Copyright (c) 2021, Noah van der Meer 9 | * 10 | * Permission is hereby granted, free of charge, to any person obtaining a copy 11 | * of this software and associated documentation files (the "Software"), to 12 | * deal in the Software without restriction, including without limitation the 13 | * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 14 | * sell copies of the Software, and to permit persons to whom the Software is 15 | * furnished to do so, subject to the following conditions: 16 | * 17 | * The above copyright notice and this permission notice shall be included in 18 | * all copies or substantial portions of the Software. 19 | * 20 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 21 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 22 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 23 | * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 24 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 25 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 26 | * IN THE SOFTWARE. 27 | * 28 | */ 29 | 30 | /* include guard */ 31 | #ifndef _YOLOV5_BUILDER_HPP_ 32 | #define _YOLOV5_BUILDER_HPP_ 33 | 34 | #include 35 | #include 36 | #include 37 | 38 | #include "yolov5_logging.hpp" 39 | 40 | namespace yolov5 41 | { 42 | 43 | /** 44 | * Build the YoloV5 TensorRT engine, which can be used for detection 45 | * 46 | * Before building a TensorRT engine, you should first initialize the 47 | * builder by using the init() method. 48 | * 49 | * 50 | * ### Basic usage example 51 | * yolov5::Builder builder; 52 | * builder.init(); 53 | * builder.buildEngine("yolov5.onnx", "yolov5.engine"); 54 | */ 55 | class Builder 56 | { 57 | public: 58 | /// *** 59 | /// Constructor / Destructor 60 | /// *** 61 | 62 | Builder() noexcept; 63 | 64 | 65 | ~Builder(); 66 | 67 | public: 68 | /// *** 69 | /// Initialization 70 | /// *** 71 | 72 | /** 73 | * @brief Initialize the Builder. 74 | * 75 | * @return Result code 76 | */ 77 | Result init() noexcept; 78 | 79 | 80 | 81 | /// *** 82 | /// Building 83 | /// *** 84 | 85 | /** 86 | * @brief Build an engine from ONNX model input, save it 87 | * to disk 88 | * 89 | * The Builder should have been initialized already through 90 | * the init() method. 91 | * 92 | * 93 | * @param inputFilePath Path to ONNX model 94 | * @param outputFilePath Path where output (engine) should be written 95 | * 96 | * @param precision (optional) Desired precision 97 | * 98 | * @return Result code 99 | */ 100 | Result buildEngine(const std::string& inputFilePath, 101 | const std::string& outputFilePath, 102 | Precision precision = PRECISION_FP32) 103 | const noexcept; 104 | 105 | /** 106 | * @brief Build an engine from ONNX model input, store in 107 | * memory 108 | * 109 | * The Builder should have been initialized already through 110 | * the init() method. 111 | * 112 | * 113 | * @param inputFilePath Path to ONNX model 114 | * @param output Output data 115 | * 116 | * @param precision (optional) Desired precision 117 | * 118 | * @return Result Result code 119 | */ 120 | Result buildEngine(const std::string& inputFilePath, 121 | std::vector* output, 122 | Precision precision = PRECISION_FP32) 123 | const noexcept; 124 | 125 | 126 | /// *** 127 | /// Logging 128 | /// *** 129 | 130 | /** 131 | * @brief Set the logger to be used by the Builder 132 | * 133 | * Note that you can potentially use this method _before_ initializing 134 | * the Builder. 135 | * 136 | * @param logger New logger; Should NOT be a nullptr 137 | * 138 | * @return Result code 139 | */ 140 | Result setLogger(std::shared_ptr logger) noexcept; 141 | 142 | 143 | /** 144 | * @brief Retrieve the logger used by the Builder 145 | * 146 | * @return Logger. Could be a nullptr 147 | */ 148 | std::shared_ptr logger() const noexcept; 149 | 150 | 151 | private: 152 | Result _buildEngine(const std::string& inputFilePath, 153 | std::shared_ptr* output, 154 | Precision precision) const noexcept; 155 | 156 | private: 157 | bool _initialized; 158 | 159 | std::shared_ptr _logger; 160 | 161 | std::unique_ptr _trtLogger; 162 | }; 163 | 164 | } /* namespace yolov5 */ 165 | 166 | #endif /* include guard */ -------------------------------------------------------------------------------- /include/yolov5_common.hpp: -------------------------------------------------------------------------------- 1 | /** 2 | * @file 3 | * 4 | * @author Noah van der Meer 5 | * @brief YoloV5 inference through NVIDIA TensorRT (common utilities) 6 | * 7 | * 8 | * Copyright (c) 2021, Noah van der Meer 9 | * 10 | * Permission is hereby granted, free of charge, to any person obtaining a copy 11 | * of this software and associated documentation files (the "Software"), to 12 | * deal in the Software without restriction, including without limitation the 13 | * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 14 | * sell copies of the Software, and to permit persons to whom the Software is 15 | * furnished to do so, subject to the following conditions: 16 | * 17 | * The above copyright notice and this permission notice shall be included in 18 | * all copies or substantial portions of the Software. 19 | * 20 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 21 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 22 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 23 | * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 24 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 25 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 26 | * IN THE SOFTWARE. 27 | * 28 | */ 29 | 30 | /* include guard */ 31 | #ifndef _YOLOV5_COMMON_HPP_ 32 | #define _YOLOV5_COMMON_HPP_ 33 | 34 | /* C/C++ */ 35 | #include 36 | 37 | /* macro's for internal use */ 38 | #define YOLOV5_UNUSED(x) (void)x; 39 | 40 | namespace yolov5 41 | { 42 | 43 | enum Result 44 | { 45 | /* Invalid input specified. This typically indicates a programming 46 | error in your software (i.e. a bug in your software). 47 | */ 48 | RESULT_FAILURE_INVALID_INPUT = -100, 49 | 50 | 51 | /* Not initialized yet */ 52 | RESULT_FAILURE_NOT_INITIALIZED = -90, 53 | 54 | 55 | /* Not loaded yet (e.g. no engine loaded yet) */ 56 | RESULT_FAILURE_NOT_LOADED = -80, 57 | 58 | /* Issue with the loaded model (e.g. input binding is missing) */ 59 | RESULT_FAILURE_MODEL_ERROR = -70, 60 | 61 | /* Indicates that you are trying to use a function that OpenCV-CUDA, 62 | but your OpenCV has no support for this. This typically indicates a 63 | programming error in your software */ 64 | RESULT_FAILURE_OPENCV_NO_CUDA = -21, 65 | 66 | 67 | /* Error related to filesystem (e.g. could not open file) */ 68 | RESULT_FAILURE_FILESYSTEM_ERROR = -50, 69 | 70 | 71 | /* Internal cuda error (e.g. could not allocate memory) */ 72 | RESULT_FAILURE_CUDA_ERROR = -40, 73 | 74 | /* Internal TensorRT error (e.g. could not setup execution context */ 75 | RESULT_FAILURE_TENSORRT_ERROR = -30, 76 | 77 | /* Internal OpenCV error */ 78 | RESULT_FAILURE_OPENCV_ERROR = -20, 79 | 80 | 81 | /* Memory-related error */ 82 | RESULT_FAILURE_ALLOC = -11, 83 | 84 | /* Other error */ 85 | RESULT_FAILURE_OTHER = -10, 86 | 87 | 88 | /* Successfull execution */ 89 | RESULT_SUCCESS = 0 90 | }; 91 | 92 | 93 | /** 94 | * @brief Get a textual description of a result code 95 | * 96 | * If the specified value 'r' is not a valid yolov5 result code, an empty 97 | * string is returned. Note that the methods and functions of this library 98 | * always return valid result codes. 99 | * 100 | * Outputs: 101 | * - RESULT_FAILURE_INVALID_INPUT: "invalid input" 102 | * - RESULT_FAILURE_NOT_INITIALIZED: "not initialized" 103 | * - RESULT_FAILURE_NOT_LOADED: "not loaded" 104 | * - RESULT_FAILURE_MODEL_ERROR: "model error" 105 | * - RESULT_FAILURE_OPENCV_NO_CUDA: "opencv lacks cuda" 106 | * - RESULT_FAILURE_FILESYSTEM_ERROR: "filesystem error" 107 | * - RESULT_FAILURE_CUDA_ERROR: "cuda error" 108 | * - RESULT_FAILURE_TENSORRT_ERROR: "tensorrt error" 109 | * - RESULT_FAILURE_OPENCV_ERROR: "opencv error" 110 | * - RESULT_FAILURE_ALLOC: "memory error" 111 | * - RESULT_FAILURE_OTHER: "other error" 112 | * - RESULT_SUCCESS: "success" 113 | * - In case of an invalid result code: "" (empty string) 114 | * 115 | * 116 | * @param r Result code 117 | * 118 | * @return String 119 | */ 120 | const char* result_to_string(Result r) noexcept; 121 | 122 | /** 123 | * @brief Get a textual description of a result code 124 | * 125 | * See result_to_string(Result) for more information. 126 | * 127 | * If 'r' is not valid a valid result code, False is returned and 'out' is 128 | * left untouched. Note that the methods and functions f this library always 129 | * return valid result codes. 130 | * 131 | * @param r Result code 132 | * @param out Output. Can be nullptr 133 | * 134 | * @return True on success, False otherwise 135 | */ 136 | bool result_to_string(Result r, std::string* out) noexcept; 137 | 138 | 139 | 140 | enum Precision 141 | { 142 | PRECISION_FP32 = 0, /**< 32-bit floating point mode */ 143 | 144 | PRECISION_FP16 = 1, /**< 16-bit floating point mode */ 145 | }; 146 | 147 | /** 148 | * @brief Get a textual description of a precision code 149 | * 150 | * If the specified value 'p' is not a valid yolov5 precision, an empty 151 | * string is returned. 152 | * 153 | * Outputs: 154 | * - PRECISION_FP32: "fp32" 155 | * - PRECISION_FP16: "fp16" 156 | * - In case of an invalid input: "" (empty string) 157 | * 158 | * @param p Precision 159 | * @return String 160 | */ 161 | const char* precision_to_string(Precision p) noexcept; 162 | 163 | /** 164 | * @brief Get a textual description of a precision code 165 | * 166 | * See precision_to_string(Precision) for more information. 167 | * 168 | * If 'r' is not valid a valid precision code, False is returned and 'out' is 169 | * left untouched. 170 | * 171 | * @param p Precision 172 | * @param out Output. Can be nullptr 173 | * 174 | * @return True on success, False otherwise 175 | */ 176 | bool precision_to_string(Precision p, std::string* out) noexcept; 177 | 178 | 179 | /** 180 | * Additional flags that can be passed to the Detector 181 | */ 182 | enum DetectorFlag 183 | { 184 | INPUT_BGR = 1, 185 | /**< input image is in BGR colorspace(opencv default) */ 186 | 187 | INPUT_RGB = 2, 188 | /**< input image is in RGB colorspace */ 189 | 190 | PREPROCESSOR_CVCUDA = 4, 191 | /**< OpenCV-CUDA pre-processing should be used */ 192 | 193 | PREPROCESSOR_CVCPU = 8 194 | /**< OpenCV-CPU pre-processing should be used */ 195 | }; 196 | 197 | 198 | 199 | } /* namespace yolov5 */ 200 | 201 | #endif /* include guard */ -------------------------------------------------------------------------------- /include/yolov5_detection.hpp: -------------------------------------------------------------------------------- 1 | /** 2 | * @file 3 | * 4 | * @author Noah van der Meer 5 | * @brief YoloV5 inference through TensorRT 6 | * 7 | * 8 | * Copyright (c) 2021, Noah van der Meer 9 | * 10 | * Permission is hereby granted, free of charge, to any person obtaining a copy 11 | * of this software and associated documentation files (the "Software"), to 12 | * deal in the Software without restriction, including without limitation the 13 | * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 14 | * sell copies of the Software, and to permit persons to whom the Software is 15 | * furnished to do so, subject to the following conditions: 16 | * 17 | * The above copyright notice and this permission notice shall be included in 18 | * all copies or substantial portions of the Software. 19 | * 20 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 21 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 22 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 23 | * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 24 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 25 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 26 | * IN THE SOFTWARE. 27 | * 28 | */ 29 | 30 | /* include guard */ 31 | #ifndef _YOLOV5_DETECTION_HPP_ 32 | #define _YOLOV5_DETECTION_HPP_ 33 | 34 | #include 35 | 36 | #include 37 | 38 | #include 39 | 40 | 41 | namespace yolov5 42 | { 43 | 44 | /** 45 | * Represents an object detected in an image by the YoloV5 model 46 | */ 47 | class Detection 48 | { 49 | public: 50 | Detection() noexcept; 51 | 52 | 53 | Detection(const int& classId, const cv::Rect& boundingBox, 54 | const double& score) noexcept; 55 | 56 | 57 | ~Detection() noexcept; 58 | 59 | public: 60 | /** 61 | * @brief Retrieve the class id of the detection 62 | * 63 | * If invalid (i.e. class is not set yet), this method 64 | * returns -1. Otherwise, if your network has NUM_CLASSES classes, any 65 | * Detection object generated by this library is _guaranteed_ to have a 66 | * classId value within [0, NUM_CLASSES - 1]. 67 | */ 68 | const int32_t& classId() const noexcept; 69 | 70 | 71 | /** 72 | * @brief Retrieve a bounding box of the detection 73 | * 74 | * Any Detection object generated by this library is _guaranteed_ to have 75 | * a bounding box completely within the original image provided to the 76 | * library for inference. 77 | */ 78 | const cv::Rect& boundingBox() const noexcept; 79 | 80 | 81 | /** 82 | * @brief Retrieve the score assigned to this detection 83 | * 84 | * Any Detection object generated by this library is _guaranteed_ to have 85 | * a score within the interval [0, 1]. 86 | * 87 | * A high value indicates a high certainty, while a low value indicates 88 | * uncertainty. 89 | */ 90 | const double& score() const noexcept; 91 | 92 | 93 | /** 94 | * @brief Retrieve the name of the class of this detection, 95 | * if known. 96 | * 97 | * YoloV5 inference only provides class numbers, not the corresponding 98 | * names of the classes. If you provide the list of class names to 99 | * the yolov5::Detector, it will automatically set them for the Detection 100 | * objects it outputs. 101 | * 102 | * If the class name is unknown, the result is an empty string. 103 | */ 104 | const std::string& className() const noexcept; 105 | 106 | 107 | /** 108 | * @brief Set the class name 109 | * 110 | * @param name New name. Anything is allowed, even an empty string 111 | * @return True on success, False otherwise 112 | */ 113 | bool setClassName(const std::string& name) noexcept; 114 | 115 | private: 116 | int32_t _classId; 117 | std::string _className; 118 | 119 | cv::Rect _boundingBox; 120 | double _score; 121 | }; 122 | 123 | 124 | /** 125 | * @brief Helper method for visualizing a Detection in 126 | * an image 127 | * 128 | * Draws the bounding box of the detection in the specified color, as well as 129 | * a small label indicating the class(name) and the confidence. These texts 130 | * are drawn in white. 131 | * 132 | * @param detection Detection 133 | * @param image Output image. Can be nullptr, in which case this 134 | * function has no effect. 135 | * 136 | * @param color Color of the bounding box 137 | * @param fontScale Scaling for the label. E.g. 1.0 138 | * 139 | * Possible result codes: 140 | * - RESULT_SUCCESS : if successful 141 | * - RESULT_FAILURE_OPENCV_ERROR : in case an error was encountered when 142 | * visualizing using OpenCV 143 | * 144 | * @return Result code 145 | */ 146 | Result visualizeDetection(const Detection& detection, cv::Mat* image, 147 | const cv::Scalar& color, 148 | const double& fontScale) noexcept; 149 | 150 | 151 | /** 152 | * Represents the classes of your model 153 | * 154 | * This can be used to map classIds to actual understandable names, 155 | * such as "human" or "suitcase". 156 | */ 157 | class Classes 158 | { 159 | public: 160 | Classes() noexcept; 161 | 162 | 163 | ~Classes() noexcept; 164 | 165 | public: 166 | 167 | /** 168 | * @brief Set the class names that should be used 169 | * 170 | * ClassId 0 will correspond to names[0], ClassId 1 to names[1] etc. 171 | * 172 | * 173 | * If any code other than RESULT_SUCCESS is returned, this method has 174 | * no effect, and loading the classes may be attempted again at a later 175 | * time. 176 | * 177 | * 178 | * @param names List of class names 179 | * 180 | * @return Result code 181 | */ 182 | Result load(const std::vector& names) noexcept; 183 | 184 | 185 | /** 186 | * @brief Try loading the class names as a list from a file 187 | 188 | * The expected file format is as following: 189 | * human 190 | * suitcase 191 | * mailbox 192 | * bike 193 | * car 194 | * ... 195 | * 196 | * Each line contains one class name; this is also the format used by 197 | * the Darknet framework. Using the above example, human will correspond 198 | * to classId 0, suitcase to classId 1, etc... 199 | * 200 | * 201 | * If any code other than RESULT_SUCCESS is returned, this method has 202 | * no effect, and loading the classes may be attempted again at a later 203 | * time. 204 | * 205 | * 206 | * @param filepath Path to file 207 | * 208 | * @return Result code 209 | */ 210 | Result loadFromFile(const std::string& filepath) noexcept; 211 | 212 | 213 | /** 214 | * @brief Query whether the classes have been loaded 215 | * 216 | * @return True if loaded, False otherwise 217 | */ 218 | bool isLoaded() const noexcept; 219 | 220 | 221 | /** 222 | * @brief Get the Class name corresponding to a ClassId 223 | * 224 | * 225 | * @param classId Class id 226 | * @param out Output. Can be nullptr 227 | * 228 | * @return Result code 229 | */ 230 | Result getName(const int& classId, std::string* out) const noexcept; 231 | 232 | 233 | /** 234 | * @brief Set the Logger to be used 235 | * 236 | * Note that you normally do not have to worry about using this 237 | * method. 238 | * 239 | * @param logger Logger 240 | */ 241 | void setLogger(std::shared_ptr logger) noexcept; 242 | 243 | private: 244 | std::shared_ptr _logger; 245 | 246 | std::vector _names; 247 | }; 248 | 249 | } /* namespace yolov5 */ 250 | 251 | #endif /* include guard */ -------------------------------------------------------------------------------- /include/yolov5_detector.hpp: -------------------------------------------------------------------------------- 1 | /** 2 | * @file 3 | * 4 | * @author Noah van der Meer 5 | * @brief YoloV5 inference through TensorRT (detector) 6 | * 7 | * 8 | * Copyright (c) 2021, Noah van der Meer 9 | * 10 | * Permission is hereby granted, free of charge, to any person obtaining a copy 11 | * of this software and associated documentation files (the "Software"), to 12 | * deal in the Software without restriction, including without limitation the 13 | * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 14 | * sell copies of the Software, and to permit persons to whom the Software is 15 | * furnished to do so, subject to the following conditions: 16 | * 17 | * The above copyright notice and this permission notice shall be included in 18 | * all copies or substantial portions of the Software. 19 | * 20 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 21 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 22 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 23 | * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 24 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 25 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 26 | * IN THE SOFTWARE. 27 | * 28 | */ 29 | 30 | /* include guard */ 31 | #ifndef _YOLOV5_DETECTOR_HPP_ 32 | #define _YOLOV5_DETECTOR_HPP_ 33 | 34 | #include 35 | 36 | namespace yolov5 37 | { 38 | 39 | 40 | /** 41 | * The main class for YoloV5 detection using TensorRT. 42 | * 43 | * Before loading a TensorRT engine or performing inference, you should first 44 | * initialize the detector by using the init() method. 45 | * 46 | * 47 | * Basic usage example 48 | * 49 | * yolov5::Detector detector; 50 | * detector.init(); 51 | * detector.loadEngine("yolov5.engine"); 52 | * 53 | * cv::Mat image = cv::imread("image.png"); 54 | * 55 | * std::vector detections; 56 | * detector.detect(image, &detections); 57 | * 58 | */ 59 | class Detector 60 | { 61 | public: 62 | /// *** 63 | /// Constructor / Destructor 64 | /// *** 65 | 66 | /** 67 | * @brief Construct a new Detector object, using default 68 | * options for everything 69 | */ 70 | Detector() noexcept; 71 | 72 | 73 | /** 74 | * @brief Destroy Detector object. Frees up any resources 75 | */ 76 | ~Detector() noexcept; 77 | 78 | private: 79 | 80 | Detector(const Detector& src) noexcept; 81 | 82 | public: 83 | /// *** 84 | /// Initialization 85 | /// *** 86 | 87 | /** 88 | * @brief Initialize the Detector. 89 | * 90 | * The initialization consists of multiple steps. If a particular step 91 | * fails, the method returns False and appropriate error messages are 92 | * logged, and later steps are not performed. In this case, the method 93 | * might be called again at a later time to complete the initialization. 94 | * 95 | * If no logger has been set before, this method will create the 96 | * default logger provided by this library, which simply prints messages 97 | * to stdout. 98 | * 99 | * 100 | * This method will also set up the pre-processor that will be used for 101 | * object detection. By default, the OpenCV-CUDA pre-processor is picked if 102 | * it is available. If not (meaning either OpenCV was built without CUDA 103 | * support, or no CUDA devices are currently available), a CPU based 104 | * pre-processor is used. To change this behaviour, use the appropriate 105 | * flags. 106 | * 107 | * 108 | * Supported flags: 109 | * - PREPROCESSOR_CVCUDA : specify that the OpenCV-CUDA pre-processor 110 | * should be used. If it is not available, this method fails and 111 | * the RESULT_FAILURE_OPENCV_NO_CUDA code is returned. 112 | * 113 | * - PREPROCESSOR_CVCPU : specify that the OpenCV-CPU pre-processor 114 | * should be used. This pre-processor is always available. 115 | * 116 | * Any unsupported flags are ignored. 117 | * 118 | * 119 | * On success, RESULT_SUCCESS is returned and no messages are logged. 120 | * 121 | * 122 | * @param flags (Optional) Additional flags for initialization 123 | * 124 | * @return Result code 125 | */ 126 | Result init(int flags = 0) noexcept; 127 | 128 | 129 | /** 130 | * @brief Query whether the Detector is initialized 131 | * 132 | * @return True if initialized, False otherwise 133 | */ 134 | bool isInitialized() const noexcept; 135 | 136 | 137 | 138 | /// *** 139 | /// Engine 140 | /// *** 141 | 142 | /** 143 | * @brief Load a TensorRT engine from a file 144 | * 145 | * The initialization should have been completed (i.e. through 146 | * the init() method). 147 | * 148 | * If any code other than RESULT_SUCCESS is returned, this method has 149 | * no effect, and loading an engine may be attempted again at a later time, 150 | * for instance after freeing up memory on either the CUDA device or host. 151 | * 152 | * If an engine is already loaded, this method will first fully load the 153 | * new engine, and only if this is successfull, the old engine is replaced. 154 | * 155 | * 156 | * @param filepath Path to engine file in filesystem 157 | * 158 | * @return Result code 159 | */ 160 | Result loadEngine(const std::string& filepath) noexcept; 161 | 162 | 163 | /** 164 | * @brief Load a TensorRT engine from the provided data 165 | * 166 | * The initialization should have been completed (i.e. through 167 | * the init() method). 168 | * 169 | * If any code other than RESULT_SUCCESS is returned, this method has 170 | * no effect, and loading an engine may be attempted again at a later time, 171 | * for instance after freeing up memory on either the CUDA device or host. 172 | * 173 | * If an engine is already loaded, this method will first fully load the 174 | * new engine, and only if this is successfull, the old engine is replaced. 175 | * 176 | * 177 | * @param data Engine data 178 | * 179 | * @return Result code 180 | */ 181 | Result loadEngine(const std::vector& data) noexcept; 182 | 183 | 184 | /** 185 | * @brief Query whether an inference engine has been 186 | * loaded already 187 | * 188 | * @return True if loaded, False otherwise 189 | */ 190 | bool isEngineLoaded() const noexcept; 191 | 192 | 193 | 194 | /// *** 195 | /// Classes 196 | /// *** 197 | 198 | /** 199 | * @brief Retrieve the number of classes of 200 | * the engine/network 201 | * 202 | * An engine should have been loaded already. If not, an error message is 203 | * logged and 0 is returned. 204 | * 205 | * @return Number of classes 206 | */ 207 | int numClasses() const noexcept; 208 | 209 | 210 | /** 211 | * @brief Set the classes of the network 212 | * 213 | * See the 'Classes' class for more information. 214 | * 215 | * Note that it is NOT mandatory to set the Classes object. This is only 216 | * useful if you want ClassIds to be automatically mapped to 217 | * Class names (e.g. "human", "bike") in the detections that are given by 218 | * the Detector. 219 | * 220 | * 221 | * This method may be used at any point in time, e.g. before/after 222 | * initialization, before/after loading an engine etc. 223 | * 224 | * 225 | * @param classes Classes (e.g. names) 226 | * 227 | * @return Result code 228 | */ 229 | Result setClasses(const Classes& classes) noexcept; 230 | 231 | 232 | 233 | /// *** 234 | /// Detection 235 | /// *** 236 | 237 | /** 238 | * @brief Detect objects in the specified image using 239 | * the YoloV5 model 240 | * 241 | * An engine should have been loaded already. 242 | * 243 | * This method accepts input of any size, but providing an input 244 | * of the exact size for which the network was configured will result in 245 | * a lower detection time, since no pre-processing is required. The network 246 | * input size can be retrieved using the inferenceSize() method. 247 | * 248 | * 249 | * By default, this method assumes that your input is in BGR format. If 250 | * this is not the case, this can be specified by setting the appropriate 251 | * flags. 252 | * 253 | * 254 | * Supported flags: 255 | * - INPUT_BGR : specify that the input is in BGR format (opencv default) 256 | * - INPUT_RGB :specify that the input is in RGB format 257 | 258 | * Any unsupported flags are ignored. 259 | * 260 | * 261 | * If any code other than RESULT_SUCCESS is returned, the output 'out' is 262 | * left untouched. 263 | * 264 | * 265 | * @param img Input images 266 | * @param out Output; Can be nullptr 267 | * @param flags (Optional) Additional flags for detection 268 | * 269 | * @return True on success, False otherwise 270 | */ 271 | Result detect(const cv::Mat& img, 272 | std::vector* out, 273 | int flags = 0) noexcept; 274 | 275 | /** 276 | * @brief Detect objects in the specified image (in CUDA 277 | * memory) using the YoloV5 model 278 | * 279 | * See the documentation on detect(const cv::Mat&, std::vector*, 280 | * int) for more information. 281 | */ 282 | Result detect(const cv::cuda::GpuMat& img, 283 | std::vector* out, 284 | int flags = 0) noexcept; 285 | 286 | /** 287 | * @brief Detect objects in the specified images using 288 | * batch inference with the YoloV5 model 289 | * 290 | * An engine should have been loaded already. 291 | * 292 | * This method accepts inputs of any size, but providing inputs 293 | * of the exact size for which the network was configured will result in 294 | * a lower detection time, since no pre-processing is required. The network 295 | * input size can be retrieved using the inferenceSize() method. 296 | * 297 | * Note that the inputs can potentially all have different sizes if this is 298 | * desired. 299 | * 300 | * 301 | * By default, this method assumes that your inputs are in BGR format. If 302 | * this is not the case, this can be specified by setting the appropriate 303 | * flags. 304 | * 305 | * 306 | * Supported flags: 307 | * - INPUT_BGR : specify that the inputs are all in BGR 308 | * format (opencv default) 309 | * - INPUT_RGB : specify that the inputs are all in RGB format 310 | 311 | * Any unsupported flags are ignored. 312 | * 313 | * 314 | * If any code other than RESULT_SUCCESS is returned, the output 'out' is 315 | * left untouched. 316 | * 317 | * 318 | * @param images Input images 319 | * @param out Outputs for each image. 320 | * @param flags (Optional) Additional flags for detection 321 | * 322 | * @return True on success, False otherwise 323 | */ 324 | Result detectBatch(const std::vector& images, 325 | std::vector>* out, 326 | int flags = 0) noexcept; 327 | 328 | /** 329 | * @brief Detect objects in the specified images (in CUDA 330 | * memory) using batch inference with YoloV5 331 | 332 | * See the documentation on detectBatch(const std::vector&, 333 | * std::vector*, int) for more information. 334 | */ 335 | Result detectBatch(const std::vector& images, 336 | std::vector>* out, 337 | int flags = 0) noexcept; 338 | 339 | 340 | /// *** 341 | /// Detection Parameters 342 | /// *** 343 | 344 | /** 345 | * @brief Obtain the score threshold 346 | */ 347 | double scoreThreshold() const noexcept; 348 | 349 | 350 | /** 351 | * @brief Set the Score threshold: used to filter 352 | * objects by score 353 | * 354 | * @param v Score threshold. Should be in [0, 1] 355 | * @return Result code 356 | */ 357 | Result setScoreThreshold(const double& v) noexcept; 358 | 359 | 360 | /** 361 | * @brief Obtain the NMS threshold 362 | */ 363 | double nmsThreshold() const noexcept; 364 | 365 | 366 | /** 367 | * @brief Set the NMS threshold 368 | * 369 | * @param v NMS threshold. Should be in [0, 1] 370 | * @return Result code 371 | */ 372 | Result setNmsThreshold(const double& v) noexcept; 373 | 374 | 375 | 376 | /// *** 377 | /// Engine/Network properties 378 | /// *** 379 | 380 | /** 381 | * @brief Retrieve the batch size of the engine/network 382 | * 383 | * An engine should have been loaded already. If not, an error message is 384 | * logged and 0 is returned. 385 | * 386 | * @return Batch size 387 | */ 388 | int batchSize() const noexcept; 389 | 390 | /** 391 | * @brief Input size for which the network was configured 392 | * 393 | * An engine should have been loaded already. If not, an error message is 394 | * logged and Size(0, 0) is returned. 395 | * 396 | * @return Size 397 | */ 398 | cv::Size inferenceSize() const noexcept; 399 | 400 | 401 | /// *** 402 | /// Logging 403 | /// *** 404 | 405 | /** 406 | * @brief Set a custom logger to be used by the Detector 407 | * 408 | * This method can be called at any time, either before or after 409 | * initialization, and will take effect immediately. 410 | * 411 | * 412 | * @param logger New logger; Should NOT be a nullptr 413 | * 414 | * @return Result code 415 | */ 416 | Result setLogger(std::shared_ptr logger) noexcept; 417 | 418 | 419 | /** 420 | * @brief Retrieve the logger used by the Detector 421 | * 422 | * @return Logger. Can potentially be a nullptr if 423 | * the Detector has not been initialized yet 424 | */ 425 | std::shared_ptr logger() const noexcept; 426 | 427 | 428 | private: 429 | /** 430 | * @brief Not implemented 431 | */ 432 | Detector& operator=(const Detector& rhs); 433 | 434 | /** 435 | * @brief Load the engine from data 436 | * 437 | * @return Result code 438 | */ 439 | Result _loadEngine(const std::vector& data) noexcept; 440 | 441 | void _printBindings(const std::unique_ptr& engine) 442 | const noexcept; 443 | 444 | int _batchSize() const noexcept; 445 | 446 | int _numClasses() const noexcept; 447 | 448 | 449 | Result _detect(std::vector* out); 450 | 451 | Result _detectBatch(const int& nrImages, 452 | std::vector>* out); 453 | 454 | /** 455 | * @brief Run the TensorRT engine on the network inputs, 456 | * copy the output to host memory 457 | */ 458 | Result _inference(const char* logid); 459 | 460 | /** 461 | * @brief Decode network output, convert to 462 | * proper Detection objects 463 | */ 464 | Result _decodeOutput(const char* logid, const int& index, 465 | std::vector* out); 466 | 467 | private: 468 | bool _initialized; 469 | 470 | std::shared_ptr _logger; 471 | 472 | Classes _classes; 473 | double _scoreThreshold; 474 | double _nmsThreshold; 475 | 476 | 477 | /* TensorRT */ 478 | std::unique_ptr _trtLogger; 479 | std::unique_ptr _trtRuntime; 480 | 481 | /* note: execution context depends on the engine, and should be destroyed 482 | _before_ the engine is destroyed */ 483 | std::unique_ptr _trtEngine; 484 | std::unique_ptr _trtExecutionContext; 485 | 486 | 487 | /* I/O */ 488 | internal::EngineBinding _inputBinding; 489 | internal::EngineBinding _outputBinding; 490 | 491 | std::unique_ptr _preprocessor; 492 | 493 | internal::DeviceMemory _deviceMemory; 494 | 495 | std::vector _outputHostMemory; 496 | }; 497 | 498 | } /* namespace yolov5 */ 499 | 500 | #endif /* include guard */ -------------------------------------------------------------------------------- /include/yolov5_detector_internal.hpp: -------------------------------------------------------------------------------- 1 | /** 2 | * @file 3 | * 4 | * @author Noah van der Meer 5 | * @brief YoloV5 inference through TensorRT (detector internals) 6 | * 7 | * 8 | * Copyright (c) 2021, Noah van der Meer 9 | * 10 | * Permission is hereby granted, free of charge, to any person obtaining a copy 11 | * of this software and associated documentation files (the "Software"), to 12 | * deal in the Software without restriction, including without limitation the 13 | * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 14 | * sell copies of the Software, and to permit persons to whom the Software is 15 | * furnished to do so, subject to the following conditions: 16 | * 17 | * The above copyright notice and this permission notice shall be included in 18 | * all copies or substantial portions of the Software. 19 | * 20 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 21 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 22 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 23 | * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 24 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 25 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 26 | * IN THE SOFTWARE. 27 | * 28 | */ 29 | 30 | 31 | #include "yolov5_detection.hpp" 32 | #include "yolov5_logging.hpp" 33 | 34 | #include 35 | 36 | #include 37 | #include 38 | 39 | #include 40 | 41 | namespace yolov5 42 | { 43 | 44 | namespace internal 45 | { 46 | 47 | int32_t dimsVolume(const nvinfer1::Dims& dims) noexcept; 48 | 49 | 50 | bool dimsToString(const nvinfer1::Dims& dims, std::string* out) noexcept; 51 | 52 | 53 | /** 54 | * Used to store all (relevant) properties of an engine binding. This does 55 | * not include memory or any i/o. 56 | */ 57 | class EngineBinding 58 | { 59 | public: 60 | EngineBinding() noexcept; 61 | 62 | ~EngineBinding() noexcept; 63 | 64 | public: 65 | 66 | void swap(EngineBinding& other) noexcept; 67 | 68 | 69 | const int& index() const noexcept; 70 | 71 | const std::string& name() const noexcept; 72 | 73 | const nvinfer1::Dims& dims() const noexcept; 74 | const int& volume() const noexcept; 75 | 76 | bool isDynamic() const noexcept; 77 | 78 | const bool& isInput() const noexcept; 79 | 80 | 81 | void toString(std::string* out) const noexcept; 82 | 83 | 84 | static bool setup(const std::unique_ptr& engine, 85 | const std::string& name, EngineBinding* binding) noexcept; 86 | 87 | static bool setup(const std::unique_ptr& engine, 88 | const int& index, EngineBinding* binding) noexcept; 89 | private: 90 | int _index; 91 | 92 | std::string _name; 93 | 94 | nvinfer1::Dims _dims; 95 | int _volume; /* note: calculated based on dims */ 96 | 97 | bool _isInput; 98 | }; 99 | 100 | 101 | /** 102 | * Used to manage memory on the CUDA device, corresponding to the engine 103 | * bindings. 104 | */ 105 | class DeviceMemory 106 | { 107 | public: 108 | DeviceMemory() noexcept; 109 | 110 | ~DeviceMemory() noexcept; 111 | 112 | private: 113 | DeviceMemory(const DeviceMemory&); 114 | 115 | public: 116 | 117 | void swap(DeviceMemory& other) noexcept; 118 | 119 | /** 120 | * @brief Get the beginning of the data. This can be passed onto 121 | * the TensorRT engine 122 | */ 123 | void** begin() const noexcept; 124 | 125 | /** 126 | * @brief Obtain a pointer to the device memory corresponding to 127 | * the specified binding 128 | * 129 | * @param index Index of the engine binding 130 | */ 131 | void* at(const int& index) const noexcept; 132 | 133 | /** 134 | * @brief Try setting up the Device Memory based on the TensorRT 135 | * engine 136 | * 137 | * 138 | * @param logger Logger to be used 139 | * 140 | * @param engine TensorRT engine 141 | * @param output Output 142 | * 143 | * @return Result Result code 144 | */ 145 | static Result setup(const std::shared_ptr& logger, 146 | std::unique_ptr& engine, 147 | DeviceMemory* output) noexcept; 148 | 149 | private: 150 | std::vector _memory; 151 | }; 152 | 153 | 154 | /** 155 | * @brief Check whether OpenCV-CUDA is supported 156 | */ 157 | bool opencvHasCuda() noexcept; 158 | 159 | 160 | /** 161 | * Used to store the Letterbox parameters used for a particular image. These 162 | * can be used to transform the bounding boxes returned by the engine to use 163 | * coordinates in the original input image. 164 | */ 165 | class PreprocessorTransform 166 | { 167 | public: 168 | PreprocessorTransform() noexcept; 169 | 170 | PreprocessorTransform(const cv::Size& inputSize, 171 | const double& f, const int& leftWidth, 172 | const int& topHeight) noexcept; 173 | 174 | ~PreprocessorTransform() noexcept; 175 | 176 | private: 177 | 178 | public: 179 | 180 | /** 181 | * @brief Transform bounding box from network space to input 182 | * space 183 | */ 184 | cv::Rect transformBbox(const cv::Rect& input) const noexcept; 185 | 186 | private: 187 | cv::Size _inputSize; 188 | 189 | double _f; 190 | int _leftWidth; 191 | int _topHeight; 192 | }; 193 | 194 | 195 | /** 196 | * Used to perform pre-processing task, and to store intermediate buffers to 197 | * speed up repeated computations. 198 | * 199 | * Note that this base class does not actually do any processing. 200 | */ 201 | class Preprocessor 202 | { 203 | public: 204 | Preprocessor() noexcept; 205 | 206 | virtual ~Preprocessor() noexcept; 207 | 208 | enum InputType 209 | { 210 | INPUTTYPE_BGR = 0, 211 | INPUTTYPE_RGB 212 | }; 213 | 214 | public: 215 | 216 | void setLogger(std::shared_ptr logger) noexcept; 217 | 218 | /** 219 | * @brief Set up the Preprocessor 220 | * 221 | * The implementation of the base class only manages the transforms. 222 | * 223 | * @param inputDims Engine input dimensions 224 | * @param flags Additional flags 225 | * @param batchSize Number of images that will be processed (i.e. in 226 | * batch mode) 227 | * @param inputMemory Start of input on the CUDA device 228 | * 229 | * @return True on success, False otherwise 230 | */ 231 | virtual bool setup(const nvinfer1::Dims& inputDims, 232 | const int& flags, const int& batchSize, 233 | float* inputMemory) noexcept = 0; 234 | 235 | virtual void reset() noexcept = 0; 236 | 237 | /** 238 | * @brief Process the input 239 | * 240 | * @param index Index in the input batch 241 | * @param input Input image 242 | * @param last Boolean indicating whether this is the last image 243 | * in the batch 244 | * 245 | * @return True on success, False otherwise 246 | */ 247 | virtual bool process(const int& index, const cv::Mat& input, 248 | const bool& last) noexcept; 249 | 250 | virtual bool process(const int& index, 251 | const cv::cuda::GpuMat& input, const bool& last) noexcept; 252 | 253 | virtual cudaStream_t cudaStream() const noexcept = 0; 254 | 255 | virtual bool synchronizeCudaStream() noexcept = 0; 256 | 257 | 258 | /** 259 | * @brief Transform bounding box from network space to input 260 | * space, for a particular image in the batch 261 | * 262 | * @param index Index in the input batch 263 | * @param input Input bounding box 264 | */ 265 | cv::Rect transformBbox(const int& index, 266 | const cv::Rect& bbox) const noexcept; 267 | 268 | protected: 269 | std::shared_ptr _logger; 270 | 271 | std::vector _transforms; 272 | }; 273 | 274 | 275 | /** 276 | * Preprocessing based on letterboxing with OpenCV CPU operations 277 | */ 278 | class CvCpuPreprocessor : public Preprocessor 279 | { 280 | public: 281 | CvCpuPreprocessor() noexcept; 282 | 283 | virtual ~CvCpuPreprocessor() noexcept; 284 | 285 | public: 286 | 287 | virtual bool setup(const nvinfer1::Dims& inputDims, 288 | const int& flags, const int& batchSize, 289 | float* inputMemory) 290 | noexcept override; 291 | 292 | virtual void reset() noexcept override; 293 | 294 | virtual bool process(const int& index, 295 | const cv::Mat& input, const bool& last) noexcept override; 296 | 297 | virtual bool process(const int& index, 298 | const cv::cuda::GpuMat& input, const bool& last) 299 | noexcept override; 300 | 301 | virtual cudaStream_t cudaStream() const noexcept override; 302 | 303 | virtual bool synchronizeCudaStream() noexcept override; 304 | 305 | private: 306 | cudaStream_t _cudaStream; 307 | 308 | InputType _lastType; 309 | int _lastBatchSize; 310 | 311 | int _networkCols; 312 | int _networkRows; 313 | 314 | cv::Mat _buffer1; 315 | cv::Mat _buffer2; 316 | cv::Mat _buffer3; 317 | 318 | std::vector> _inputChannels; 319 | 320 | std::vector _hostInputMemory; 321 | float* _deviceInputMemory; 322 | }; 323 | 324 | 325 | /** 326 | * Preprocessing based on letterboxing with OpenCV-CUDA operations. Note 327 | * that OpenCV-CUDA must be available for this. If not, this class will 328 | * not perform any actual operations. 329 | */ 330 | class CvCudaPreprocessor : public Preprocessor 331 | { 332 | public: 333 | 334 | CvCudaPreprocessor() noexcept; 335 | 336 | virtual ~CvCudaPreprocessor() noexcept; 337 | 338 | public: 339 | 340 | virtual bool setup(const nvinfer1::Dims& inputDims, 341 | const int& flags, const int& batchSize, 342 | float* inputMemory) 343 | noexcept override; 344 | 345 | virtual void reset() noexcept override; 346 | 347 | virtual bool process(const int& index, 348 | const cv::Mat& input, const bool& last) noexcept override; 349 | 350 | virtual bool process(const int& index, 351 | const cv::cuda::GpuMat& input, 352 | const bool& last) noexcept override; 353 | 354 | virtual cudaStream_t cudaStream() const noexcept override; 355 | 356 | virtual bool synchronizeCudaStream() noexcept override; 357 | 358 | private: 359 | cv::cuda::Stream _cudaStream; 360 | 361 | InputType _lastType; 362 | int _lastBatchSize; 363 | 364 | int _networkCols; 365 | int _networkRows; 366 | 367 | cv::cuda::GpuMat _buffer0; 368 | cv::cuda::GpuMat _buffer1; 369 | cv::cuda::GpuMat _buffer2; 370 | cv::cuda::GpuMat _buffer3; 371 | 372 | std::vector> _inputChannels; 373 | }; 374 | 375 | 376 | } /* namespace internal */ 377 | 378 | } /* namespace yolov5 */ -------------------------------------------------------------------------------- /include/yolov5_logging.hpp: -------------------------------------------------------------------------------- 1 | /** 2 | * @file 3 | * 4 | * @author Noah van der Meer 5 | * @brief YoloV5 inference through TensorRT (logging) 6 | * 7 | * 8 | * Copyright (c) 2021, Noah van der Meer 9 | * 10 | * Permission is hereby granted, free of charge, to any person obtaining a copy 11 | * of this software and associated documentation files (the "Software"), to 12 | * deal in the Software without restriction, including without limitation the 13 | * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 14 | * sell copies of the Software, and to permit persons to whom the Software is 15 | * furnished to do so, subject to the following conditions: 16 | * 17 | * The above copyright notice and this permission notice shall be included in 18 | * all copies or substantial portions of the Software. 19 | * 20 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 21 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 22 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 23 | * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 24 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 25 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 26 | * IN THE SOFTWARE. 27 | * 28 | */ 29 | 30 | /* include guard */ 31 | #ifndef _YOLOV5_LOGGING_HPP_ 32 | #define _YOLOV5_LOGGING_HPP_ 33 | 34 | #include "yolov5_common.hpp" 35 | 36 | #include 37 | 38 | #include 39 | 40 | 41 | 42 | namespace yolov5 43 | { 44 | 45 | enum LogLevel 46 | { 47 | LOGGING_DEBUG = 0, /**< verbose, low-level details */ 48 | 49 | LOGGING_INFO = 1, /**< informational messages */ 50 | 51 | LOGGING_WARNING = 2, /**< warning messages */ 52 | 53 | LOGGING_ERROR = 3 /**< error messages */ 54 | }; 55 | 56 | /** 57 | * @brief Convert the LogLevel to string 58 | * 59 | * If the specified value 'l' is not a valid YoloV5 loglevel, an empty 60 | * string is returned. Note that the methods and functions of this library 61 | * always use proper loglevels when logging. 62 | * 63 | * Outputs: 64 | * - LOGGING_DEBUG: "debug" 65 | * - LOGGING_INFO: "info" 66 | * - LOGGING_WARNING: "warning" 67 | * - LOGGING_ERROR: "error" 68 | * 69 | * @param l Log level 70 | * 71 | * @return Log level string 72 | */ 73 | const char* loglevel_to_string(const LogLevel& l) noexcept; 74 | 75 | 76 | /** 77 | * @brief Convert the LogLevel to string 78 | * 79 | * See yolov5_loglevel_to_string(YoloV5_LogLevel) for more information. 80 | * 81 | * If 'r' is not a valid loglevel, False is returned and 'out' is left 82 | * untouched. Note that the methods and functions of this library 83 | * always use proper loglevels when logging. 84 | * 85 | * 86 | * @param l Log level 87 | * @param out Output. Can be nullptr 88 | * 89 | * @return True on success, False otherwise 90 | */ 91 | bool loglevel_to_string(const LogLevel& l, std::string* out) noexcept; 92 | 93 | 94 | /** 95 | * The main logger used in the yolov5-tensorrt library 96 | * 97 | * You can use this class to integrate the yolov5-tensorrt logging into 98 | * your own preferred logging facilities. To do so, create your own 99 | * class that inherits from this class, and override the print(...) 100 | * method. 101 | */ 102 | class Logger 103 | { 104 | public: 105 | Logger() noexcept; 106 | 107 | 108 | virtual ~Logger(); 109 | 110 | public: 111 | 112 | /** 113 | * @brief Print/Log a message. Override this method to 114 | * integrate logging with your own preferred 115 | * logging mechanism. 116 | * 117 | * The default implementation prints all messages to stdout, and appends 118 | * a newline at the end of all messages. 119 | * 120 | * 121 | * This method is not marked as 'noexcept' intentionally, in case a user 122 | * does not (want to)deal with exceptions properly in a derived class. 123 | * 124 | * @param level Logging level 125 | * @param msg Message to be printed. 126 | */ 127 | virtual void print(const LogLevel& level, const char* msg); 128 | 129 | 130 | /** 131 | * @brief Log a message 132 | * 133 | * Internally, this method will forward the message to print(). 134 | * 135 | * @param level Logging level 136 | * @param msg Message to be printed 137 | */ 138 | void log(const LogLevel& level, const char* msg) noexcept; 139 | 140 | 141 | /** 142 | * @brief Log a formatted message 143 | * 144 | * Internally, this method will forward the message to print(). 145 | * 146 | * @param level Logging level 147 | * @param fmt Format string 148 | */ 149 | void logf(const LogLevel& level, const char* fmt, ...) noexcept 150 | __attribute__ ((format (printf, 3, 4))); 151 | 152 | private: 153 | }; 154 | 155 | 156 | /** 157 | * Logger used to integrate TensorRT and yolov5-tensorrt logging 158 | * 159 | * This logger forwards all messages from the TensorRT logger 160 | * to a yolov5::Logger. 161 | * 162 | * Normally, it is not necessary for a user of the library to worry about 163 | * this class, unless you are using TensorRT in other places as well and wish 164 | * to integrate logging further. 165 | */ 166 | class TensorRT_Logger : public nvinfer1::ILogger 167 | { 168 | public: 169 | TensorRT_Logger() noexcept; 170 | 171 | 172 | /** 173 | * @brief Construct a new TensorRT_Logger object 174 | * 175 | * @param logger Pointer to Logger. Can be nullptr 176 | */ 177 | TensorRT_Logger(std::shared_ptr logger) noexcept; 178 | 179 | 180 | ~TensorRT_Logger(); 181 | public: 182 | 183 | /** 184 | * @brief Set the YoloV5-TensorRT logger 185 | * 186 | * @param logger Pointer to Logger. Can be nullptr 187 | */ 188 | void setLogger(std::shared_ptr logger) noexcept; 189 | 190 | 191 | virtual void log(nvinfer1::ILogger::Severity severity, 192 | const char* msg) noexcept override; 193 | 194 | private: 195 | std::shared_ptr _logger; 196 | }; 197 | 198 | } /* namespace yolov5 */ 199 | 200 | #endif /* include guard */ -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = [ 3 | "setuptools>=42", 4 | "wheel", 5 | "pybind11", 6 | "pybind11-global", 7 | "cmake>=3.1" 8 | ] 9 | build-backend = "setuptools.build_meta" -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import os 2 | import subprocess 3 | 4 | import setuptools 5 | from setuptools.command.build_ext import build_ext 6 | 7 | class CMakeExtension(setuptools.Extension): 8 | 9 | def __init__(self, name, sourcedir=''): 10 | setuptools.Extension.__init__(self, name, sources=[]) 11 | self.sourcedir = os.path.abspath(sourcedir) 12 | 13 | 14 | class CMakeBuildExt(build_ext): 15 | 16 | def build_extension(self, ext) -> None: 17 | 18 | if not os.path.exists(self.build_temp): 19 | os.makedirs(self.build_temp) 20 | 21 | cmakedir = os.path.abspath(os.path.dirname(self.get_ext_fullpath(ext.name))) 22 | if not cmakedir.endswith(os.path.sep): 23 | cmakedir += os.path.sep 24 | 25 | subprocess.check_call(['cmake', ext.sourcedir, '-DBUILD_PYTHON=ON', 26 | '-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=' + cmakedir], 27 | cwd=self.build_temp) 28 | subprocess.check_call(['cmake', '--build', '.'], cwd=self.build_temp) 29 | 30 | 31 | setuptools.setup( 32 | name = 'yolov5tensorrt', 33 | version = '0.1', 34 | author = 'Noah van der Meer', 35 | description = 'Real-time object detection with YOLOv5 and TensorRT', 36 | long_description = 'file: README.md', 37 | long_description_content_type = 'text/markdown', 38 | url = 'https://github.com/noahmr/yolov5-tensorrt', 39 | keywords = ['yolov5', 'tensorrt', 'object detection'], 40 | license = "MIT", 41 | classifiers = [ 42 | 'Programming Language :: Python :: 3', 43 | 'License :: OSI Approved :: MIT License', 44 | 'Operating System :: Linux' 45 | ], 46 | install_requires = ['numpy'], 47 | ext_modules = [CMakeExtension('yolov5tensorrt')], 48 | cmdclass = {'build_ext': CMakeBuildExt}, 49 | scripts = [ 50 | 'examples/builder/build_engine.py', 51 | 'examples/batch/process_batch.py', 52 | 'examples/image/process_image.py' 53 | ], 54 | python_requires = '>=3.6' 55 | ) -------------------------------------------------------------------------------- /src/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | add_library(yolov5-tensorrt SHARED 2 | yolov5_detection.cpp 3 | yolov5_detector.cpp 4 | yolov5_detector_internal.cpp 5 | yolov5_builder.cpp 6 | yolov5_common.cpp 7 | yolov5_logging.cpp 8 | ) 9 | target_include_directories(yolov5-tensorrt PUBLIC 10 | ${OpenCV_INCLUDE_DIRS} 11 | ${CUDA_INCLUDE_DIRS} 12 | ) 13 | set_target_properties(yolov5-tensorrt PROPERTIES PUBLIC_HEADER "${YOLOV5_INCLUDE_FILES}") 14 | target_link_libraries(yolov5-tensorrt) 15 | 16 | ## install rules 17 | ## 18 | ## Install the library to the default lib destination, 19 | ## but the header files in a dedicated subdirectory called 'yolov5-tensorrt' 20 | install(TARGETS yolov5-tensorrt 21 | LIBRARY DESTINATION lib 22 | PUBLIC_HEADER DESTINATION include/yolov5-tensorrt) -------------------------------------------------------------------------------- /src/yolov5_builder.cpp: -------------------------------------------------------------------------------- 1 | /** 2 | * @file 3 | * 4 | * @author Noah van der Meer 5 | * @brief YoloV5 inference through NVIDIA TensorRT (builder) 6 | * 7 | * 8 | * Copyright (c) 2021, Noah van der Meer 9 | * 10 | * Permission is hereby granted, free of charge, to any person obtaining a copy 11 | * of this software and associated documentation files (the "Software"), to 12 | * deal in the Software without restriction, including without limitation the 13 | * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 14 | * sell copies of the Software, and to permit persons to whom the Software is 15 | * furnished to do so, subject to the following conditions: 16 | * 17 | * The above copyright notice and this permission notice shall be included in 18 | * all copies or substantial portions of the Software. 19 | * 20 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 21 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 22 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 23 | * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 24 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 25 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 26 | * IN THE SOFTWARE. 27 | * 28 | */ 29 | 30 | #include "yolov5_builder.hpp" 31 | 32 | #include 33 | 34 | #include 35 | 36 | namespace yolov5 37 | { 38 | 39 | Builder::Builder() noexcept : _initialized(false) 40 | { 41 | } 42 | 43 | Builder::~Builder() 44 | { 45 | } 46 | 47 | Result Builder::init() noexcept 48 | { 49 | /* Initialize Logger */ 50 | if(!_logger) 51 | { 52 | try 53 | { 54 | _logger = std::make_shared(); 55 | } 56 | catch(const std::exception& e) 57 | { 58 | /* logging not available */ 59 | return RESULT_FAILURE_ALLOC; 60 | } 61 | } 62 | 63 | /* Initialize TensorRT logger */ 64 | if(!_trtLogger) 65 | { 66 | try 67 | { 68 | _trtLogger = std::make_unique(_logger); 69 | } 70 | catch(const std::exception& e) 71 | { 72 | _logger->logf(LOGGING_ERROR, "[Builder] init() failure: could not " 73 | "create TensorRT logger: %s", e.what()); 74 | return RESULT_FAILURE_ALLOC; 75 | } 76 | } 77 | 78 | _initialized = true; 79 | return RESULT_SUCCESS; 80 | } 81 | 82 | Result Builder::buildEngine(const std::string& inputFilePath, 83 | const std::string& outputFilePath, Precision precision) 84 | const noexcept 85 | { 86 | if(!_initialized) 87 | { 88 | if(_logger) 89 | { 90 | _logger->log(LOGGING_ERROR, "[Builder] buildEngine() failure: " 91 | "builder is not initialized yet"); 92 | } 93 | return RESULT_FAILURE_NOT_INITIALIZED; 94 | } 95 | 96 | std::shared_ptr engineOutput; 97 | Result r = _buildEngine(inputFilePath, &engineOutput, precision); 98 | if(r != RESULT_SUCCESS) 99 | { 100 | return r; 101 | } 102 | 103 | /* Write to disk */ 104 | _logger->logf(LOGGING_INFO, "[Builder] buildEngine(): writing serialized " 105 | "engine to file: %s", outputFilePath.c_str()); 106 | 107 | std::ofstream outputFile; 108 | outputFile.open(outputFilePath, std::ios::out | std::ios::binary); 109 | outputFile.write((char*)engineOutput->data(), engineOutput->size()); 110 | if(!outputFile.good()) 111 | { 112 | _logger->log(LOGGING_ERROR, "[Builder] buildEngine() failure: " 113 | "error encountered writing to output file"); 114 | return RESULT_FAILURE_FILESYSTEM_ERROR; 115 | } 116 | outputFile.close(); 117 | 118 | return RESULT_SUCCESS; 119 | } 120 | 121 | Result Builder::buildEngine(const std::string& inputFilePath, 122 | std::vector* output, Precision precision) 123 | const noexcept 124 | { 125 | if(!_initialized) 126 | { 127 | if(_logger) 128 | { 129 | _logger->log(LOGGING_ERROR, "[Builder] buildEngine() failure: " 130 | "builder is not initialized yet"); 131 | } 132 | return RESULT_FAILURE_NOT_INITIALIZED; 133 | } 134 | 135 | std::shared_ptr engineOutput; 136 | Result r = _buildEngine(inputFilePath, &engineOutput, precision); 137 | if(r != RESULT_SUCCESS) 138 | { 139 | return r; 140 | } 141 | 142 | if(!output) 143 | { 144 | return RESULT_SUCCESS; 145 | } 146 | 147 | try 148 | { 149 | output->resize(engineOutput->size()); 150 | } 151 | catch(const std::exception& e) 152 | { 153 | _logger->logf(LOGGING_ERROR, "[Builder] buildEngine() failure: could " 154 | "not set up output memory: %s", e.what()); 155 | return RESULT_FAILURE_ALLOC; 156 | } 157 | std::memcpy(output->data(), (char*)engineOutput->data(), 158 | engineOutput->size()); 159 | return RESULT_SUCCESS; 160 | } 161 | 162 | Result Builder::setLogger(std::shared_ptr logger) noexcept 163 | { 164 | if(!logger) 165 | { 166 | if(_logger) 167 | { 168 | _logger->log(LOGGING_ERROR, "[Builder] setLogger() failure: " 169 | "provided logger is nullptr"); 170 | } 171 | return RESULT_FAILURE_INVALID_INPUT; 172 | } 173 | _logger = logger; 174 | 175 | return RESULT_SUCCESS; 176 | } 177 | 178 | std::shared_ptr Builder::logger() const noexcept 179 | { 180 | return _logger; 181 | } 182 | 183 | Result Builder::_buildEngine(const std::string& inputFilePath, 184 | std::shared_ptr* output, 185 | Precision precision) 186 | const noexcept 187 | { 188 | const char* precisionStr = precision_to_string(precision); 189 | if(std::strlen(precisionStr) == 0) 190 | { 191 | _logger->log(LOGGING_ERROR, "[Builder] buildEngine() failure: " 192 | "invalid precision specified"); 193 | return RESULT_FAILURE_INVALID_INPUT; 194 | } 195 | 196 | try 197 | { 198 | std::unique_ptr builder( 199 | nvinfer1::createInferBuilder(*_trtLogger)); 200 | 201 | const auto explicitBatch = 1U << static_cast( 202 | nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); 203 | std::unique_ptr network( 204 | builder->createNetworkV2(explicitBatch)); 205 | 206 | builder->setMaxBatchSize(1); 207 | 208 | std::unique_ptr parser( 209 | nvonnxparser::createParser(*network, *_trtLogger)); 210 | if(!parser->parseFromFile(inputFilePath.c_str(), 211 | (int)nvinfer1::ILogger::Severity::kWARNING)) 212 | { 213 | _logger->log(LOGGING_ERROR, "[Builder] buildEngine() failure: " 214 | "could not parse ONNX model from file"); 215 | return RESULT_FAILURE_MODEL_ERROR; 216 | } 217 | 218 | std::unique_ptr config( 219 | builder->createBuilderConfig()); 220 | config->setMaxWorkspaceSize(1 << 20); 221 | 222 | if(precision == PRECISION_FP32) 223 | { 224 | /* this is the default */ 225 | } 226 | else if(precision == PRECISION_FP16) 227 | { 228 | if(!builder->platformHasFastFp16()) 229 | { 230 | _logger->log(LOGGING_ERROR, "[Builder] buildEngine() failure: " 231 | "fp16 precision specified, but not supported by " 232 | "current platform"); 233 | return RESULT_FAILURE_INVALID_INPUT; 234 | } 235 | config->setFlag(nvinfer1::BuilderFlag::kFP16); 236 | } 237 | 238 | _logger->logf(LOGGING_INFO, "[Builder] buildEngine(): building and " 239 | "serializing engine at %s precision. This may take a while", 240 | precisionStr); 241 | 242 | std::shared_ptr serialized( 243 | builder->buildSerializedNetwork(*network, *config)); 244 | if(!serialized) 245 | { 246 | _logger->log(LOGGING_ERROR, "[Builder] buildEngine() failure: " 247 | "could not build serialized engine"); 248 | return RESULT_FAILURE_TENSORRT_ERROR; 249 | } 250 | *output = serialized; 251 | } 252 | catch(const std::exception& e) 253 | { 254 | _logger->logf(LOGGING_ERROR, "[Builder] buildEngine() failure: got " 255 | "exception: %s", e.what()); 256 | return RESULT_FAILURE_OTHER; 257 | } 258 | return RESULT_SUCCESS; 259 | } 260 | 261 | } /* namespace yolov5 */ -------------------------------------------------------------------------------- /src/yolov5_common.cpp: -------------------------------------------------------------------------------- 1 | /** 2 | * @file 3 | * 4 | * @author Noah van der Meer 5 | * @brief YoloV5 inference through NVIDIA TensorRT (common utilities) 6 | * 7 | * 8 | * Copyright (c) 2021, Noah van der Meer 9 | * 10 | * Permission is hereby granted, free of charge, to any person obtaining a copy 11 | * of this software and associated documentation files (the "Software"), to 12 | * deal in the Software without restriction, including without limitation the 13 | * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 14 | * sell copies of the Software, and to permit persons to whom the Software is 15 | * furnished to do so, subject to the following conditions: 16 | * 17 | * The above copyright notice and this permission notice shall be included in 18 | * all copies or substantial portions of the Software. 19 | * 20 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 21 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 22 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 23 | * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 24 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 25 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 26 | * IN THE SOFTWARE. 27 | * 28 | */ 29 | 30 | #include "yolov5_common.hpp" 31 | 32 | #include 33 | 34 | namespace yolov5 35 | { 36 | 37 | const char* result_to_string(yolov5::Result r) noexcept 38 | { 39 | if(r == RESULT_FAILURE_INVALID_INPUT) 40 | { 41 | return "invalid input"; 42 | } 43 | else if(r == RESULT_FAILURE_NOT_INITIALIZED) 44 | { 45 | return "not initialized"; 46 | } 47 | else if(r == RESULT_FAILURE_NOT_LOADED) 48 | { 49 | return "not loaded"; 50 | } 51 | else if(r == RESULT_FAILURE_MODEL_ERROR) 52 | { 53 | return "model error"; 54 | } 55 | else if(r == RESULT_FAILURE_OPENCV_NO_CUDA) 56 | { 57 | return "opencv lacks cuda"; 58 | } 59 | else if(r == RESULT_FAILURE_FILESYSTEM_ERROR) 60 | { 61 | return "filesystem error"; 62 | } 63 | else if(r == RESULT_FAILURE_CUDA_ERROR) 64 | { 65 | return "cuda error"; 66 | } 67 | else if(r == RESULT_FAILURE_TENSORRT_ERROR) 68 | { 69 | return "tensorrt error"; 70 | } 71 | else if(r == RESULT_FAILURE_OPENCV_ERROR) 72 | { 73 | return "opencv error"; 74 | } 75 | else if(r == RESULT_FAILURE_ALLOC) 76 | { 77 | return "alloc error"; 78 | } 79 | else if(r == RESULT_FAILURE_OTHER) 80 | { 81 | return "other error"; 82 | } 83 | else if(r == RESULT_SUCCESS) 84 | { 85 | return "success"; 86 | } 87 | else 88 | { 89 | return ""; 90 | } 91 | } 92 | 93 | bool result_to_string(Result r, std::string* out) noexcept 94 | { 95 | const char* str = result_to_string(r); 96 | if(std::strlen(str) == 0) 97 | { 98 | return false; 99 | } 100 | 101 | if(out != nullptr) 102 | { 103 | try 104 | { 105 | *out = str; 106 | } 107 | catch(const std::exception& e) 108 | { 109 | } 110 | } 111 | return true; 112 | } 113 | 114 | const char* precision_to_string(Precision p) noexcept 115 | { 116 | if(p == PRECISION_FP32) 117 | { 118 | return "fp32"; 119 | } 120 | else if(p == PRECISION_FP16) 121 | { 122 | return "fp16"; 123 | } 124 | else 125 | { 126 | return ""; 127 | } 128 | } 129 | 130 | bool precision_to_string(Precision p, std::string* out) noexcept 131 | { 132 | const char* str = precision_to_string(p); 133 | if(std::strlen(str) == 0) 134 | { 135 | return false; 136 | } 137 | 138 | if(out != nullptr) 139 | { 140 | try 141 | { 142 | *out = str; 143 | } 144 | catch(const std::exception& e) 145 | { 146 | } 147 | } 148 | return true; 149 | } 150 | 151 | } /* namespace yolov5 */ -------------------------------------------------------------------------------- /src/yolov5_detection.cpp: -------------------------------------------------------------------------------- 1 | /** 2 | * @file 3 | * 4 | * @author Noah van der Meer 5 | * @brief YoloV5 inference through TensorRT 6 | * 7 | * 8 | * Copyright (c) 2021, Noah van der Meer 9 | * 10 | * Permission is hereby granted, free of charge, to any person obtaining a copy 11 | * of this software and associated documentation files (the "Software"), to 12 | * deal in the Software without restriction, including without limitation the 13 | * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 14 | * sell copies of the Software, and to permit persons to whom the Software is 15 | * furnished to do so, subject to the following conditions: 16 | * 17 | * The above copyright notice and this permission notice shall be included in 18 | * all copies or substantial portions of the Software. 19 | * 20 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 21 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 22 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 23 | * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 24 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 25 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 26 | * IN THE SOFTWARE. 27 | * 28 | */ 29 | 30 | #include "yolov5_detection.hpp" 31 | 32 | #include 33 | 34 | namespace yolov5 35 | { 36 | 37 | Detection::Detection() noexcept : _classId(-1), _score(0) 38 | { 39 | } 40 | 41 | Detection::Detection(const int& classId, 42 | const cv::Rect& boundingBox, const double& score) noexcept 43 | : _classId(classId), _boundingBox(boundingBox), _score(score) 44 | { 45 | } 46 | 47 | Detection::~Detection() noexcept 48 | { 49 | } 50 | 51 | const int32_t& Detection::classId() const noexcept 52 | { 53 | return _classId; 54 | } 55 | 56 | const cv::Rect& Detection::boundingBox() const noexcept 57 | { 58 | return _boundingBox; 59 | } 60 | 61 | const double& Detection::score() const noexcept 62 | { 63 | return _score; 64 | } 65 | 66 | const std::string& Detection::className() const noexcept 67 | { 68 | return _className; 69 | } 70 | 71 | bool Detection::setClassName(const std::string& name) noexcept 72 | { 73 | try 74 | { 75 | _className = name; 76 | } 77 | catch(const std::exception& e) 78 | { 79 | return false; 80 | } 81 | return true; 82 | } 83 | 84 | Result visualizeDetection(const Detection& detection, cv::Mat* image, 85 | const cv::Scalar& color, 86 | const double& fontScale) noexcept 87 | { 88 | if(image == nullptr) 89 | { 90 | return RESULT_SUCCESS; 91 | } 92 | 93 | try 94 | { 95 | /* Draw bounding box around the detection */ 96 | const int bboxThickness = 2; 97 | const cv::Rect& bbox = detection.boundingBox(); 98 | cv::rectangle(*image, bbox, color, bboxThickness); 99 | 100 | /* "className: score" 101 | 102 | or alternatively "classId: score" 103 | if no class name is known 104 | 105 | display the score with 2 decimal places 106 | */ 107 | std::string className = detection.className(); 108 | if(className.length() == 0) 109 | { 110 | className = std::to_string(detection.classId()); 111 | } 112 | std::stringstream ss; 113 | ss << std::fixed << std::setprecision(2) << detection.score(); 114 | const std::string label = className + ": " + ss.str(); 115 | 116 | /* Draw a rectangle above the bounding box, in which the 117 | label will be written */ 118 | const int textThickness = 1; 119 | 120 | int baseline = 0; 121 | const cv::Size textSize = cv::getTextSize(label, 122 | cv::FONT_HERSHEY_PLAIN, 123 | fontScale, textThickness, &baseline); 124 | const cv::Point tl(bbox.x - bboxThickness/2.0, bbox.y-textSize.height); 125 | const cv::Rect labelRect(tl, textSize); 126 | cv::rectangle(*image, labelRect, color, -1); /* filled rectangle */ 127 | 128 | /* white text on top of the previously drawn rectangle */ 129 | const cv::Point bl(tl.x, bbox.y - bboxThickness/2.0); 130 | cv::putText(*image, label, bl, cv::FONT_HERSHEY_PLAIN, 131 | fontScale, cv::Scalar(255, 255, 255), textThickness); 132 | } 133 | catch(const std::exception& e) 134 | { 135 | return RESULT_FAILURE_OPENCV_ERROR; 136 | } 137 | return RESULT_SUCCESS; 138 | } 139 | 140 | Classes::Classes() noexcept 141 | { 142 | } 143 | 144 | Classes::~Classes() noexcept 145 | { 146 | } 147 | 148 | Result Classes::load(const std::vector& names) noexcept 149 | { 150 | if(names.size() == 0 && _logger) 151 | { 152 | if(_logger) 153 | { 154 | _logger->log(LOGGING_ERROR, "[Classes] load() warning: specified " 155 | "list of class names is empty!"); 156 | } 157 | return RESULT_FAILURE_INVALID_INPUT; 158 | } 159 | 160 | try 161 | { 162 | _names = names; 163 | } 164 | catch(const std::exception& e) 165 | { 166 | if(_logger) 167 | { 168 | _logger->logf(LOGGING_ERROR, "[Classes] load() failure: got " 169 | "exception trying to copy names: %s", e.what()); 170 | } 171 | return RESULT_FAILURE_ALLOC; 172 | } 173 | 174 | if(_logger) 175 | { 176 | _logger->logf(LOGGING_INFO, "[Classes] Loaded %d classes", 177 | (unsigned int)names.size()); 178 | } 179 | return RESULT_SUCCESS; 180 | } 181 | 182 | Result Classes::loadFromFile(const std::string& filepath) noexcept 183 | { 184 | std::vector names; 185 | 186 | std::ifstream file(filepath, std::ios::binary); 187 | if(!file.good()) 188 | { 189 | if(_logger) 190 | { 191 | _logger->logf(LOGGING_ERROR, "[Classes] loadFromFile() failure: " 192 | "could not open file '%s'", filepath.c_str()); 193 | } 194 | return RESULT_FAILURE_FILESYSTEM_ERROR; 195 | } 196 | 197 | try 198 | { 199 | std::string line; 200 | while(std::getline(file, line)) 201 | { 202 | if(line.length() > 0) 203 | { 204 | names.push_back(line); 205 | } 206 | } 207 | } 208 | catch(const std::exception& e) 209 | { 210 | if(_logger) 211 | { 212 | _logger->logf(LOGGING_ERROR, "[Classes] loadFromFile() failure: " 213 | "got exception while reading classes from file: %s", 214 | e.what()); 215 | } 216 | return RESULT_FAILURE_ALLOC; 217 | } 218 | file.close(); 219 | 220 | if(names.size() == 0) 221 | { 222 | if(_logger) 223 | { 224 | _logger->log(LOGGING_ERROR, "[Classes] loadFromFile() failure: " 225 | "could not load any classes"); 226 | } 227 | return RESULT_FAILURE_OTHER; 228 | } 229 | 230 | names.swap(_names); 231 | if(_logger) 232 | { 233 | _logger->logf(LOGGING_INFO, "[Classes] Loaded %d classes", 234 | (unsigned int)_names.size()); 235 | } 236 | return RESULT_SUCCESS; 237 | } 238 | 239 | bool Classes::isLoaded() const noexcept 240 | { 241 | return (_names.size() > 0); 242 | } 243 | 244 | Result Classes::getName(const int& classId, std::string* out) const noexcept 245 | { 246 | if((unsigned int)classId >= _names.size() || classId < 0) 247 | { 248 | if(_logger) 249 | { 250 | _logger->logf(LOGGING_ERROR, "[Classes] getName() failure: no " 251 | "info about specified classId '%i'", classId); 252 | } 253 | return RESULT_FAILURE_INVALID_INPUT; 254 | } 255 | 256 | if(out != nullptr) 257 | { 258 | try 259 | { 260 | *out = _names[classId]; 261 | } 262 | catch(const std::exception& e) 263 | { 264 | if(_logger) 265 | { 266 | _logger->logf(LOGGING_ERROR, "[Classes] getName() failure: got" 267 | " exception when setting output: %s", e.what()); 268 | } 269 | return RESULT_FAILURE_ALLOC; 270 | } 271 | } 272 | return RESULT_SUCCESS; 273 | } 274 | 275 | void Classes::setLogger(std::shared_ptr logger) noexcept 276 | { 277 | _logger = logger; 278 | } 279 | 280 | } /* namespace yolov5 */ -------------------------------------------------------------------------------- /src/yolov5_detector.cpp: -------------------------------------------------------------------------------- 1 | /** 2 | * @file 3 | * 4 | * @author Noah van der Meer 5 | * @brief YoloV5 inference through TensorRT (detector) 6 | * 7 | * 8 | * Copyright (c) 2021, Noah van der Meer 9 | * 10 | * Permission is hereby granted, free of charge, to any person obtaining a copy 11 | * of this software and associated documentation files (the "Software"), to 12 | * deal in the Software without restriction, including without limitation the 13 | * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 14 | * sell copies of the Software, and to permit persons to whom the Software is 15 | * furnished to do so, subject to the following conditions: 16 | * 17 | * The above copyright notice and this permission notice shall be included in 18 | * all copies or substantial portions of the Software. 19 | * 20 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 21 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 22 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 23 | * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 24 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 25 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 26 | * IN THE SOFTWARE. 27 | * 28 | */ 29 | 30 | #include "yolov5_detector.hpp" 31 | 32 | #include 33 | #include 34 | #include 35 | 36 | #include 37 | #include 38 | 39 | /* CUDA */ 40 | #include 41 | 42 | namespace yolov5 43 | { 44 | 45 | Detector::Detector() noexcept 46 | : _initialized(false), _scoreThreshold(0.4), _nmsThreshold(0.4) 47 | { 48 | } 49 | 50 | Detector::~Detector() noexcept 51 | { 52 | } 53 | 54 | Result Detector::init(int flags) noexcept 55 | { 56 | /* Initialize Logger */ 57 | if(!_logger) 58 | { 59 | try 60 | { 61 | _logger = std::make_shared(); 62 | } 63 | catch(const std::exception& e) 64 | { 65 | /* logging not available */ 66 | return RESULT_FAILURE_ALLOC; 67 | } 68 | } 69 | 70 | /* Initialize TensorRT logger */ 71 | if(!_trtLogger) 72 | { 73 | try 74 | { 75 | _trtLogger = std::make_unique(_logger); 76 | } 77 | catch(const std::exception& e) 78 | { 79 | _logger->logf(LOGGING_ERROR, "[Detector] init() failure: could not" 80 | " create TensorRT logger: %s", e.what()); 81 | return RESULT_FAILURE_ALLOC; 82 | } 83 | } 84 | 85 | /* Set up Preprocessor */ 86 | if(!_preprocessor) 87 | { 88 | try 89 | { 90 | const bool cvCudaAvailable = internal::opencvHasCuda(); 91 | 92 | if((flags & PREPROCESSOR_CVCUDA) && (flags & PREPROCESSOR_CVCPU)) 93 | { 94 | _logger->log(LOGGING_ERROR, "[Detector] init() failure: " 95 | "both PREPROCESSOR_CVCUDA and PREPROCESSOR_CVCPU " 96 | "flags specified"); 97 | return RESULT_FAILURE_INVALID_INPUT; 98 | } 99 | 100 | /* If the CVCUDA flag was specified, OpenCV-CUDA has to be 101 | available or fail */ 102 | if(flags & PREPROCESSOR_CVCUDA && !cvCudaAvailable) 103 | { 104 | _logger->log(LOGGING_ERROR, "[Detector] init() failure: " 105 | "PREPROCESSOR_CVCUDA flag specified, but " 106 | "OpenCV-CUDA pre-processor is not available."); 107 | return RESULT_FAILURE_OPENCV_NO_CUDA; 108 | } 109 | 110 | bool useCudaPreprocessor = cvCudaAvailable; 111 | if(flags & PREPROCESSOR_CVCPU) 112 | { 113 | useCudaPreprocessor = false; 114 | } 115 | 116 | if(useCudaPreprocessor) 117 | { 118 | _logger->log(LOGGING_INFO, "[Detector] Using OpenCV-CUDA " 119 | "pre-processor"); 120 | _preprocessor = 121 | std::make_unique(); 122 | } 123 | else 124 | { 125 | _logger->log(LOGGING_INFO, "[Detector] Using OpenCV-CPU " 126 | "pre-processor"); 127 | _preprocessor = 128 | std::make_unique(); 129 | } 130 | } 131 | catch(const std::exception& e) 132 | { 133 | _logger->logf(LOGGING_ERROR, "[Detector] init() failure: " 134 | "could not set up preprocessor: %s", e.what()); 135 | return RESULT_FAILURE_ALLOC; 136 | } 137 | _preprocessor->setLogger(_logger); 138 | } 139 | 140 | /* Initialize TensorRT runtime */ 141 | if(!_trtRuntime) 142 | { 143 | nvinfer1::IRuntime* trtRuntime 144 | = nvinfer1::createInferRuntime(*_trtLogger); 145 | if(trtRuntime == nullptr) 146 | { 147 | _logger->log(LOGGING_ERROR, "[Detector] init() failure: could not" 148 | "create TensorRT runtime"); 149 | return RESULT_FAILURE_TENSORRT_ERROR; 150 | } 151 | std::unique_ptr ptr(trtRuntime); 152 | ptr.swap(_trtRuntime); 153 | } 154 | 155 | 156 | _initialized = true; 157 | return RESULT_SUCCESS; 158 | } 159 | 160 | bool Detector::isInitialized() const noexcept 161 | { 162 | return _initialized; 163 | } 164 | 165 | Result Detector::loadEngine(const std::string& filepath) noexcept 166 | { 167 | if(!_initialized) 168 | { 169 | if(_logger) 170 | { 171 | _logger->log(LOGGING_ERROR, "[Detector] loadEngine() failure: " 172 | "detector is not initialized yet"); 173 | } 174 | return RESULT_FAILURE_NOT_INITIALIZED; 175 | } 176 | _logger->logf(LOGGING_INFO, "[Detector] Loading TensorRT engine " 177 | "from '%s'", filepath.c_str()); 178 | 179 | std::ifstream file(filepath, std::ios::binary); 180 | if(!file.good()) 181 | { 182 | _logger->log(LOGGING_ERROR, "[Detector] loadEngine() failure: could " 183 | "not open specified file"); 184 | return RESULT_FAILURE_FILESYSTEM_ERROR; 185 | } 186 | 187 | std::vector data; 188 | try 189 | { 190 | file.seekg(0, file.end); 191 | const auto size = file.tellg(); 192 | file.seekg(0, file.beg); 193 | 194 | /* read entire file into vector */ 195 | data.resize(size); 196 | file.read(data.data(), size); 197 | } 198 | catch(const std::exception& e) 199 | { 200 | _logger->logf(LOGGING_ERROR, "[Detector] loadEngine() failure: could " 201 | "not load file into memory: %s", e.what()); 202 | return RESULT_FAILURE_ALLOC; 203 | } 204 | file.close(); 205 | return _loadEngine(data); 206 | } 207 | 208 | Result Detector::loadEngine(const std::vector& data) noexcept 209 | { 210 | if(!_initialized) 211 | { 212 | if(_logger) 213 | { 214 | _logger->log(LOGGING_ERROR, "[Detector] loadEngine() failure: " 215 | "detector is not initialized yet"); 216 | } 217 | return RESULT_FAILURE_NOT_INITIALIZED; 218 | } 219 | return _loadEngine(data); 220 | } 221 | 222 | bool Detector::isEngineLoaded() const noexcept 223 | { 224 | return (bool)_trtEngine; 225 | } 226 | 227 | int Detector::numClasses() const noexcept 228 | { 229 | if(!isEngineLoaded()) 230 | { 231 | if(_logger) 232 | { 233 | _logger->log(LOGGING_ERROR, "[Detector] numClasses() failure: no " 234 | "engine loaded"); 235 | } 236 | return 0; 237 | } 238 | return _numClasses(); 239 | } 240 | 241 | Result Detector::setClasses(const Classes& classes) noexcept 242 | { 243 | if(!classes.isLoaded()) 244 | { 245 | if(_logger) 246 | { 247 | _logger->log(LOGGING_ERROR, "[Detector] setClasses() failure: " 248 | "invalid input specified: classes not yet loaded"); 249 | } 250 | return RESULT_FAILURE_INVALID_INPUT; 251 | } 252 | 253 | try 254 | { 255 | _classes = classes; 256 | } 257 | catch(const std::exception& e) 258 | { 259 | if(_logger) 260 | { 261 | _logger->logf(LOGGING_ERROR, "[Detector] setClasses() failure: " 262 | "could not set up classes: %s", e.what()); 263 | } 264 | return RESULT_FAILURE_ALLOC; 265 | } 266 | return RESULT_SUCCESS; 267 | } 268 | 269 | Result Detector::detect(const cv::Mat& img, 270 | std::vector* out, 271 | int flags) noexcept 272 | { 273 | if(!isEngineLoaded()) 274 | { 275 | if(_logger) 276 | { 277 | _logger->log(LOGGING_ERROR, "[Detector] detect() failure: no " 278 | "engine loaded"); 279 | } 280 | return RESULT_FAILURE_NOT_LOADED; 281 | } 282 | 283 | /** Pre-processing **/ 284 | if(!_preprocessor->setup( _inputBinding.dims(), flags, _batchSize(), 285 | (float*)_deviceMemory.at(_inputBinding.index()) )) 286 | { 287 | _logger->log(LOGGING_ERROR, "[Detector] detect() failure: could not " 288 | "set up pre-processor"); 289 | return RESULT_FAILURE_OTHER; 290 | } 291 | if(!_preprocessor->process(0, img, true)) 292 | { 293 | _logger->log(LOGGING_ERROR, "[Detector] detect() failure: could not " 294 | "pre-process input"); 295 | return RESULT_FAILURE_OTHER; 296 | } 297 | 298 | return _detect(out); 299 | } 300 | 301 | Result Detector::detect(const cv::cuda::GpuMat& img, 302 | std::vector* out, 303 | int flags) noexcept 304 | { 305 | if(!isEngineLoaded()) 306 | { 307 | if(_logger) 308 | { 309 | _logger->log(LOGGING_ERROR, "[Detector] detect() failure: no " 310 | "engine loaded"); 311 | } 312 | return RESULT_FAILURE_NOT_LOADED; 313 | } 314 | 315 | if(!internal::opencvHasCuda()) 316 | { 317 | _logger->log(LOGGING_ERROR, "[Detector] detect() failure: this method " 318 | "requires OpenCV-CUDA support, which is not " 319 | "available. Use detect(const cv::Mat&, ...) instead"); 320 | return RESULT_FAILURE_OPENCV_NO_CUDA; 321 | } 322 | 323 | /** Pre-processing **/ 324 | if(!_preprocessor->setup(_inputBinding.dims(), flags, _batchSize(), 325 | (float*)_deviceMemory.at(_inputBinding.index()) )) 326 | { 327 | _logger->log(LOGGING_ERROR, "[Detector] detect() failure: could not " 328 | "set up pre-processor"); 329 | return RESULT_FAILURE_OTHER; 330 | } 331 | if(!_preprocessor->process(0, img, true)) 332 | { 333 | _logger->log(LOGGING_ERROR, "[Detector] detect() failure: could not " 334 | "pre-process input"); 335 | return RESULT_FAILURE_OTHER; 336 | } 337 | 338 | return _detect(out); 339 | } 340 | 341 | Result Detector::detectBatch(const std::vector& images, 342 | std::vector>* out, 343 | int flags) noexcept 344 | { 345 | if(!isEngineLoaded()) 346 | { 347 | if(_logger) 348 | { 349 | _logger->log(LOGGING_ERROR, "[Detector] detectBatch() failure: no " 350 | "engine loaded"); 351 | } 352 | return RESULT_FAILURE_NOT_LOADED; 353 | } 354 | if(images.size() == 0) 355 | { 356 | _logger->log(LOGGING_ERROR, "[Detector] detectBatch() failure: list " 357 | "of inputs is empty"); 358 | return RESULT_FAILURE_INVALID_INPUT; 359 | } 360 | if((int)images.size() > _batchSize()) 361 | { 362 | _logger->logf(LOGGING_ERROR, "[Detector] detectBatch() failure: " 363 | "specified %d images, but batch size is %i", 364 | (unsigned int)images.size(), _batchSize()); 365 | return RESULT_FAILURE_INVALID_INPUT; 366 | } 367 | const int numProcessed = MIN((int)images.size(), _batchSize()); 368 | 369 | /** Pre-processing **/ 370 | if(!_preprocessor->setup(_inputBinding.dims(), flags, _batchSize(), 371 | (float*)_deviceMemory.at(_inputBinding.index()) )) 372 | { 373 | _logger->log(LOGGING_ERROR, "[Detector] detectBatch() failure: could " 374 | "not set up pre-processor"); 375 | return RESULT_FAILURE_OTHER; 376 | } 377 | 378 | for(int i = 0; i < numProcessed; ++i) 379 | { 380 | if(!_preprocessor->process(i, images[i], i == numProcessed - 1)) 381 | { 382 | _logger->logf(LOGGING_ERROR, "[Detector] detectBatch() " 383 | "failure: preprocessing for image %i failed", i); 384 | return RESULT_FAILURE_OTHER; 385 | } 386 | } 387 | 388 | return _detectBatch(numProcessed, out); 389 | } 390 | 391 | Result Detector::detectBatch(const std::vector& images, 392 | std::vector>* out, 393 | int flags) noexcept 394 | { 395 | if(!isEngineLoaded()) 396 | { 397 | if(_logger) 398 | { 399 | _logger->log(LOGGING_ERROR, "[Detector] detectBatch() failure: no " 400 | "engine loaded"); 401 | } 402 | return RESULT_FAILURE_NOT_LOADED; 403 | } 404 | 405 | if(!internal::opencvHasCuda()) 406 | { 407 | _logger->log(LOGGING_ERROR, "[Detector] detectBatch() failure: this " 408 | "method requires OpenCV-CUDA support, which is not " 409 | "available. Use detectBatch(const " 410 | "std::vector&, ...) instead"); 411 | return RESULT_FAILURE_OPENCV_NO_CUDA; 412 | } 413 | if(images.size() == 0) 414 | { 415 | _logger->log(LOGGING_ERROR, "[Detector] detectBatch() failure: list " 416 | "of inputs is empty"); 417 | return RESULT_FAILURE_INVALID_INPUT; 418 | } 419 | if((int)images.size() > _batchSize()) 420 | { 421 | _logger->logf(LOGGING_ERROR, "[Detector] detectBatch() failure: " 422 | "specified %d images, but batch size is %i", 423 | (unsigned int)images.size(), _batchSize()); 424 | return RESULT_FAILURE_INVALID_INPUT; 425 | } 426 | const int numProcessed = MIN((int)images.size(), _batchSize()); 427 | 428 | /** Pre-processing **/ 429 | if(!_preprocessor->setup(_inputBinding.dims(), flags, _batchSize(), 430 | (float*)_deviceMemory.at(_inputBinding.index()) )) 431 | { 432 | _logger->log(LOGGING_ERROR, "[Detector] detectBatch() failure: could " 433 | "not set up pre-processor"); 434 | return RESULT_FAILURE_OTHER; 435 | } 436 | 437 | for(int i = 0; i < numProcessed; ++i) 438 | { 439 | if(!_preprocessor->process(i, images[i], i == numProcessed - 1)) 440 | { 441 | _logger->logf(LOGGING_ERROR, "[Detector] detectBatch() " 442 | "failure: preprocessing for image %i failed", i); 443 | return RESULT_FAILURE_OTHER; 444 | } 445 | } 446 | 447 | return _detectBatch(numProcessed, out); 448 | } 449 | 450 | double Detector::scoreThreshold() const noexcept 451 | { 452 | return _scoreThreshold; 453 | } 454 | 455 | Result Detector::setScoreThreshold(const double& v) noexcept 456 | { 457 | if(v < 0 || v > 1) 458 | { 459 | if(_logger) 460 | { 461 | _logger->log(LOGGING_ERROR, "[Detector] setScoreThreshold() " 462 | "failure: invalid value specified"); 463 | } 464 | return RESULT_FAILURE_INVALID_INPUT; 465 | } 466 | _scoreThreshold = v; 467 | return RESULT_SUCCESS; 468 | } 469 | 470 | double Detector::nmsThreshold() const noexcept 471 | { 472 | return _nmsThreshold; 473 | } 474 | 475 | Result Detector::setNmsThreshold(const double& v) noexcept 476 | { 477 | if(v < 0 || v > 1) 478 | { 479 | if(_logger) 480 | { 481 | _logger->log(LOGGING_ERROR, "[Detector] setNmsThreshold() " 482 | "failure: invalid value specified"); 483 | } 484 | return RESULT_FAILURE_INVALID_INPUT; 485 | } 486 | _nmsThreshold = v; 487 | return RESULT_SUCCESS; 488 | } 489 | 490 | int Detector::batchSize() const noexcept 491 | { 492 | if(!isEngineLoaded()) 493 | { 494 | if(_logger) 495 | { 496 | _logger->log(LOGGING_ERROR, "[Detector] batchSize() failure: no " 497 | "engine loaded"); 498 | } 499 | return 0; 500 | } 501 | return _batchSize(); 502 | } 503 | 504 | cv::Size Detector::inferenceSize() const noexcept 505 | { 506 | if(!isEngineLoaded()) 507 | { 508 | if(_logger) 509 | { 510 | _logger->log(LOGGING_ERROR, "[Detector] inferenceSize() failure: " 511 | "no engine loaded"); 512 | } 513 | return cv::Size(0, 0); 514 | } 515 | const auto& inputDims = _inputBinding.dims(); 516 | const int rows = inputDims.d[2]; 517 | const int cols = inputDims.d[3]; 518 | return cv::Size(cols, rows); 519 | } 520 | 521 | Result Detector::setLogger(std::shared_ptr logger) noexcept 522 | { 523 | if(!logger) 524 | { 525 | if(_logger) 526 | { 527 | _logger->log(LOGGING_ERROR, "[Detector] setLogger() failure: " 528 | "provided logger is nullptr"); 529 | } 530 | return RESULT_FAILURE_INVALID_INPUT; 531 | } 532 | _logger = logger; /* note: operator= of shared_ptr is marked noexcept */ 533 | 534 | if(_preprocessor) 535 | { 536 | _preprocessor->setLogger(_logger); 537 | } 538 | return RESULT_SUCCESS; 539 | } 540 | 541 | std::shared_ptr Detector::logger() const noexcept 542 | { 543 | return _logger; 544 | } 545 | 546 | Result Detector::_loadEngine(const std::vector& data) noexcept 547 | { 548 | /* Try to deserialize engine */ 549 | _logger->log(LOGGING_INFO, "[Detector] Deserializing inference engine. " 550 | "This may take a while..."); 551 | std::unique_ptr engine( 552 | _trtRuntime->deserializeCudaEngine(data.data(), data.size())); 553 | if(!engine) 554 | { 555 | _logger->log(LOGGING_ERROR, "[Detector] loadEngine() failure: could " 556 | "not deserialize engine"); 557 | return RESULT_FAILURE_TENSORRT_ERROR; 558 | } 559 | 560 | /* Create execution context */ 561 | std::unique_ptr executionContext( 562 | engine->createExecutionContext()); 563 | if(!executionContext) 564 | { 565 | _logger->log(LOGGING_ERROR, "[Detector] loadEngine() failure: could " 566 | "not create execution context"); 567 | return RESULT_FAILURE_TENSORRT_ERROR; 568 | } 569 | 570 | _printBindings(engine); 571 | 572 | /* Determine input bindings & verify that it matches what is expected */ 573 | internal::EngineBinding input; 574 | if(!internal::EngineBinding::setup(engine, "images", &input)) 575 | { 576 | _logger->log(LOGGING_ERROR, "[Detector] loadEngine() failure: could " 577 | "not set up input binding"); 578 | return RESULT_FAILURE_MODEL_ERROR; 579 | } 580 | if(input.dims().nbDims != 4) 581 | { 582 | std::string str; 583 | internal::dimsToString(input.dims(), &str); 584 | _logger->logf(LOGGING_ERROR, "[Detector] loadEngine() failure: " 585 | "unexpected input dimensions: %s", str.c_str()); 586 | return RESULT_FAILURE_MODEL_ERROR; 587 | } 588 | if(input.isDynamic()) 589 | { 590 | _logger->log(LOGGING_ERROR, "[Detector] loadEngine() failure: " 591 | "input binding has dynamic dimensions. This is not " 592 | "supported at this time!"); 593 | return RESULT_FAILURE_MODEL_ERROR; 594 | } 595 | 596 | 597 | /* Determine output binding & verify that it matches what is expected */ 598 | internal::EngineBinding output; 599 | if(!internal::EngineBinding::setup(engine, "output", &output)) 600 | { 601 | _logger->log(LOGGING_ERROR, "[Detector] loadEngine() failure: could " 602 | "not set up output binding"); 603 | return RESULT_FAILURE_MODEL_ERROR; 604 | } 605 | if(output.dims().nbDims != 3) 606 | { 607 | std::string str; 608 | internal::dimsToString(output.dims(), &str); 609 | _logger->logf(LOGGING_ERROR, "[Detector] loadEngine() failure: " 610 | "unexpected output dimensions: %s", str.c_str()); 611 | return RESULT_FAILURE_MODEL_ERROR; 612 | } 613 | if(output.isDynamic()) 614 | { 615 | _logger->log(LOGGING_ERROR, "[Detector] loadEngine() failure: " 616 | "output binding has dynamic dimensions. This is not " 617 | "supported at this time!"); 618 | return RESULT_FAILURE_MODEL_ERROR; 619 | } 620 | 621 | /* Set up Device memory for input & output */ 622 | internal::DeviceMemory memory; 623 | const Result r = internal::DeviceMemory::setup(_logger, engine, &memory); 624 | if(r != RESULT_SUCCESS) 625 | { 626 | _logger->log(LOGGING_ERROR, "[Detector] loadEngine() failure: " 627 | "could not set up device memory"); 628 | return r; 629 | } 630 | 631 | /* Set up memory on host for post-processing */ 632 | std::vector outputHostMemory; 633 | try 634 | { 635 | outputHostMemory.resize(output.volume()); 636 | } 637 | catch(const std::exception& e) 638 | { 639 | _logger->logf(LOGGING_ERROR, "[Detector] loadEngine() failure: " 640 | "could not set up output host memory: %s", e.what()); 641 | return RESULT_FAILURE_ALLOC; 642 | } 643 | 644 | 645 | /* commit to the new engine; at this point, there is nothing that 646 | will fail anymore */ 647 | if(isEngineLoaded()) 648 | { 649 | _logger->log(LOGGING_INFO, "[Detector] loadEngine() info: an engine " 650 | "is already loaded; Replacing it"); 651 | } 652 | 653 | engine.swap(_trtEngine); 654 | executionContext.swap(_trtExecutionContext); 655 | 656 | memory.swap(_deviceMemory); 657 | outputHostMemory.swap(_outputHostMemory); 658 | 659 | input.swap(_inputBinding); 660 | output.swap(_outputBinding); 661 | 662 | /* Note: this is the PreProcessor::reset() method, not the reset() 663 | method of unique_ptr (!) */ 664 | _preprocessor->reset(); 665 | 666 | _logger->log(LOGGING_INFO, "[Detector] Successfully loaded inference " 667 | "engine"); 668 | return RESULT_SUCCESS; 669 | } 670 | 671 | void Detector::_printBindings( 672 | const std::unique_ptr& engine) 673 | const noexcept 674 | { 675 | const int32_t nbBindings = engine->getNbBindings(); 676 | 677 | for(int i = 0; i < nbBindings; ++i) 678 | { 679 | internal::EngineBinding binding; 680 | internal::EngineBinding::setup(engine, i, &binding); 681 | 682 | std::string str; 683 | binding.toString(&str); 684 | 685 | _logger->logf(LOGGING_DEBUG, "[Detector] loadEngine() info: Binding " 686 | "%i - %s", i, str.c_str()); 687 | } 688 | } 689 | 690 | int Detector::_batchSize() const noexcept 691 | { 692 | return _inputBinding.dims().d[0]; 693 | } 694 | 695 | int Detector::_numClasses() const noexcept 696 | { 697 | return _outputBinding.dims().d[2] - 5; 698 | } 699 | 700 | Result Detector::_detect(std::vector* out) 701 | { 702 | /** Inference **/ 703 | Result r = _inference("detect()"); 704 | if(r != RESULT_SUCCESS) 705 | { 706 | return r; 707 | } 708 | 709 | 710 | /** Post-processing **/ 711 | std::vector lst; 712 | r = _decodeOutput("detect()", 0, &lst); 713 | if(r != RESULT_SUCCESS) 714 | { 715 | return r; 716 | } 717 | 718 | if(out != nullptr) 719 | { 720 | std::swap(lst, *out); 721 | } 722 | return RESULT_SUCCESS; 723 | } 724 | 725 | Result Detector::_detectBatch(const int& nrImages, 726 | std::vector>* out) 727 | { 728 | /** Inference **/ 729 | Result r = _inference("detectBatch()"); 730 | if(r != RESULT_SUCCESS) 731 | { 732 | return r; 733 | } 734 | 735 | 736 | /** Post-processing **/ 737 | std::vector> lst; 738 | try 739 | { 740 | lst.resize(nrImages); 741 | } 742 | catch(const std::exception& e) 743 | { 744 | _logger->logf(LOGGING_ERROR, "[Detector] detectBatch() failure: could " 745 | "not allocate output space: %s", e.what()); 746 | return RESULT_FAILURE_ALLOC; 747 | } 748 | 749 | for(int i = 0; i < nrImages; ++i) 750 | { 751 | r = _decodeOutput("detectBatch()", i, &lst[i]); 752 | if(r != RESULT_SUCCESS) 753 | { 754 | return r; 755 | } 756 | } 757 | 758 | if(out != nullptr) 759 | { 760 | std::swap(lst, *out); 761 | } 762 | return RESULT_SUCCESS; 763 | } 764 | 765 | Result Detector::_inference(const char* logid) 766 | { 767 | /* Enqueue for inference */ 768 | if(!_trtExecutionContext->enqueueV2(_deviceMemory.begin(), 769 | _preprocessor->cudaStream(), nullptr)) 770 | { 771 | _logger->logf(LOGGING_ERROR, "[Detector] %s failure: could not enqueue " 772 | "data for inference", logid); 773 | return RESULT_FAILURE_TENSORRT_ERROR; 774 | } 775 | 776 | /* Copy output back from device memory to host memory */ 777 | auto r = cudaMemcpyAsync(_outputHostMemory.data(), 778 | _deviceMemory.at(_outputBinding.index()), 779 | (int)(_outputBinding.volume() * sizeof(float)), 780 | cudaMemcpyDeviceToHost, _preprocessor->cudaStream()); 781 | if(r != 0) 782 | { 783 | _logger->logf(LOGGING_ERROR, "[Detector] %s failure: could not set up " 784 | "device-to-host transfer for output: %s", 785 | logid, cudaGetErrorString(r)); 786 | return RESULT_FAILURE_CUDA_ERROR; 787 | } 788 | 789 | /* Synchronize */ 790 | if(!_preprocessor->synchronizeCudaStream()) 791 | { 792 | return RESULT_FAILURE_CUDA_ERROR; 793 | } 794 | return RESULT_SUCCESS; 795 | } 796 | 797 | Result Detector::_decodeOutput(const char* logid, const int& index, 798 | std::vector* out) 799 | { 800 | std::vector boxes; 801 | std::vector scores; 802 | std::vector classes; 803 | 804 | const int nrClasses = numClasses(); 805 | 806 | /* Decode YoloV5 output */ 807 | const int numGridBoxes = _outputBinding.dims().d[1]; 808 | const int rowSize = _outputBinding.dims().d[2]; 809 | 810 | float* begin = _outputHostMemory.data() + index * numGridBoxes * rowSize; 811 | 812 | for(int i = 0; i < numGridBoxes; ++i) 813 | { 814 | float* ptr = begin + i * rowSize; 815 | 816 | const float objectness = ptr[4]; 817 | if(objectness < _scoreThreshold) 818 | { 819 | continue; 820 | } 821 | 822 | /* Get the class with the highest score attached to it */ 823 | double maxClassScore = 0.0; 824 | int maxScoreIndex = 0; 825 | for(int i = 0; i < nrClasses; ++i) 826 | { 827 | const float& v = ptr[5 + i]; 828 | if(v > maxClassScore) 829 | { 830 | maxClassScore = v; 831 | maxScoreIndex = i; 832 | } 833 | } 834 | const double score = objectness * maxClassScore; 835 | if(score < _scoreThreshold) 836 | { 837 | continue; 838 | } 839 | 840 | const float w = ptr[2]; 841 | const float h = ptr[3]; 842 | const float x = ptr[0] - w / 2.0; 843 | const float y = ptr[1] - h / 2.0; 844 | 845 | try 846 | { 847 | boxes.push_back(cv::Rect(x, y, w, h)); 848 | scores.push_back(score); 849 | classes.push_back(maxScoreIndex); 850 | } 851 | catch(const std::exception& e) 852 | { 853 | _logger->logf(LOGGING_ERROR, "[Detector] %s failure: got " 854 | "exception setting up model detection: %s", 855 | logid, e.what()); 856 | return RESULT_FAILURE_ALLOC; 857 | } 858 | } 859 | 860 | 861 | /* Apply non-max-suppression */ 862 | std::vector indices; 863 | try 864 | { 865 | cv::dnn::NMSBoxes(boxes, scores, _scoreThreshold, _nmsThreshold, 866 | indices); 867 | } 868 | catch(const std::exception& e) 869 | { 870 | _logger->logf(LOGGING_ERROR, "[Detector] %s failure: got exception " 871 | "applying OpenCV non-max-suppression: %s", 872 | logid, e.what()); 873 | return RESULT_FAILURE_OPENCV_ERROR; 874 | } 875 | 876 | /* Convert to Detection objects */ 877 | for(unsigned int i = 0; i < indices.size(); ++i) 878 | { 879 | const int& j = indices[i]; 880 | /* transform bounding box from network space to input space */ 881 | const cv::Rect bbox = _preprocessor->transformBbox(index, boxes[j]); 882 | const double score = MAX(0.0, MIN(1.0, scores[j])); 883 | try 884 | { 885 | out->push_back(Detection(classes[j], bbox, score)); 886 | } 887 | catch(const std::exception& e) 888 | { 889 | _logger->logf(LOGGING_ERROR, "[Detector] %s failure: got " 890 | "exception setting up Detection output: %s", 891 | logid, e.what()); 892 | return RESULT_FAILURE_ALLOC; 893 | } 894 | 895 | if(_classes.isLoaded()) 896 | { 897 | Detection& det = out->back(); 898 | 899 | std::string className; 900 | _classes.getName(det.classId(), &className); 901 | det.setClassName(className); 902 | } 903 | } 904 | return RESULT_SUCCESS; 905 | } 906 | 907 | 908 | } /* namespace yolov5 */ -------------------------------------------------------------------------------- /src/yolov5_detector_internal.cpp: -------------------------------------------------------------------------------- 1 | /** 2 | * @file 3 | * 4 | * @author Noah van der Meer 5 | * @brief YoloV5 inference through TensorRT (detector internals) 6 | * 7 | * 8 | * Copyright (c) 2021, Noah van der Meer 9 | * 10 | * Permission is hereby granted, free of charge, to any person obtaining a copy 11 | * of this software and associated documentation files (the "Software"), to 12 | * deal in the Software without restriction, including without limitation the 13 | * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 14 | * sell copies of the Software, and to permit persons to whom the Software is 15 | * furnished to do so, subject to the following conditions: 16 | * 17 | * The above copyright notice and this permission notice shall be included in 18 | * all copies or substantial portions of the Software. 19 | * 20 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 21 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 22 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 23 | * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 24 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 25 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 26 | * IN THE SOFTWARE. 27 | * 28 | */ 29 | 30 | #include "yolov5_detector_internal.hpp" 31 | 32 | #if __has_include() && __has_include() 33 | 34 | #define YOLOV5_OPENCV_HAS_CUDA 1 35 | 36 | #include /* cv::cuda::split(...) */ 37 | #include /* cv::cuda::resize(...) */ 38 | #include 39 | #endif 40 | 41 | #include 42 | 43 | namespace yolov5 44 | { 45 | 46 | namespace internal 47 | { 48 | 49 | int32_t dimsVolume(const nvinfer1::Dims& dims) noexcept 50 | { 51 | int32_t r = 0; 52 | if(dims.nbDims > 0) 53 | { 54 | r = 1; 55 | } 56 | 57 | for(int32_t i = 0; i < dims.nbDims; ++i) 58 | { 59 | r = r * dims.d[i]; 60 | } 61 | return r; 62 | } 63 | 64 | 65 | bool dimsToString(const nvinfer1::Dims& dims, std::string* out) noexcept 66 | { 67 | try 68 | { 69 | *out = "("; 70 | for(int32_t i = 0; i < dims.nbDims; ++i) 71 | { 72 | *out += std::to_string(dims.d[i]); 73 | if(i < dims.nbDims - 1) 74 | { 75 | out->push_back(','); 76 | } 77 | } 78 | out->push_back(')'); 79 | } 80 | catch(const std::exception& e) 81 | { 82 | return false; 83 | } 84 | return true; 85 | } 86 | 87 | EngineBinding::EngineBinding() noexcept 88 | { 89 | } 90 | 91 | EngineBinding::~EngineBinding() noexcept 92 | { 93 | } 94 | 95 | void EngineBinding::swap(EngineBinding& other) noexcept 96 | { 97 | std::swap(_index, other._index); 98 | std::swap(_name, other._name); 99 | std::swap(_volume, other._volume); 100 | 101 | nvinfer1::Dims tmp; 102 | std::memcpy(&tmp, &_dims, sizeof(nvinfer1::Dims)); 103 | std::memcpy(&_dims, &other._dims, sizeof(nvinfer1::Dims)); 104 | std::memcpy(&other._dims, &tmp, sizeof(nvinfer1::Dims)); 105 | 106 | std::swap(_isInput, other._isInput); 107 | } 108 | 109 | const int& EngineBinding::index() const noexcept 110 | { 111 | return _index; 112 | } 113 | 114 | const std::string& EngineBinding::name() const noexcept 115 | { 116 | return _name; 117 | } 118 | 119 | const nvinfer1::Dims& EngineBinding::dims() const noexcept 120 | { 121 | return _dims; 122 | } 123 | 124 | const int& EngineBinding::volume() const noexcept 125 | { 126 | return _volume; 127 | } 128 | 129 | bool EngineBinding::isDynamic() const noexcept 130 | { 131 | for(int i = 0; i < _dims.nbDims; ++i) 132 | { 133 | if(_dims.d[i] == -1) 134 | { 135 | return true; 136 | } 137 | } 138 | return false; 139 | } 140 | 141 | const bool& EngineBinding::isInput() const noexcept 142 | { 143 | return _isInput; 144 | } 145 | 146 | void EngineBinding::toString(std::string* out) const noexcept 147 | { 148 | std::string dimsStr; 149 | if(!dimsToString(_dims, &dimsStr)) 150 | { 151 | return; 152 | } 153 | 154 | try 155 | { 156 | *out = "name: '" + _name + "'" 157 | + " ; dims: " + dimsStr 158 | + " ; isInput: " + (_isInput ? "true" : "false") 159 | + " ; dynamic: " + (isDynamic() ? "true" : "false"); 160 | } 161 | catch(const std::exception& e) 162 | { 163 | } 164 | } 165 | 166 | bool EngineBinding::setup(const std::unique_ptr& engine, 167 | const std::string& name, EngineBinding* binding) noexcept 168 | { 169 | try 170 | { 171 | binding->_name = name; 172 | } 173 | catch(const std::exception& e) 174 | { 175 | return false; 176 | } 177 | 178 | binding->_index = engine->getBindingIndex(name.c_str()); 179 | if(binding->_index == -1) 180 | { 181 | return false; 182 | } 183 | 184 | binding->_dims = engine->getBindingDimensions(binding->_index); 185 | binding->_volume = dimsVolume(binding->_dims); 186 | 187 | binding->_isInput = engine->bindingIsInput(binding->_index); 188 | 189 | return true; 190 | } 191 | 192 | bool EngineBinding::setup(const std::unique_ptr& engine, 193 | const int& index, EngineBinding* binding) noexcept 194 | { 195 | binding->_index = index; 196 | const char* name = engine->getBindingName(index); 197 | if(name == nullptr) 198 | { 199 | return false; 200 | } 201 | 202 | try 203 | { 204 | binding->_name = std::string(name); 205 | } 206 | catch(const std::exception& e) 207 | { 208 | return false; 209 | } 210 | 211 | binding->_dims = engine->getBindingDimensions(binding->_index); 212 | binding->_volume = dimsVolume(binding->_dims); 213 | 214 | binding->_isInput = engine->bindingIsInput(binding->_index); 215 | 216 | return true; 217 | } 218 | 219 | DeviceMemory::DeviceMemory() noexcept 220 | { 221 | } 222 | 223 | DeviceMemory::~DeviceMemory() noexcept 224 | { 225 | for(unsigned int i = 0; i < _memory.size(); ++i) 226 | { 227 | if(_memory[i]) 228 | { 229 | cudaFree(_memory[i]); 230 | } 231 | } 232 | _memory.clear(); 233 | } 234 | 235 | void DeviceMemory::swap(DeviceMemory& other) noexcept 236 | { 237 | std::swap(_memory, other._memory); 238 | } 239 | 240 | void** DeviceMemory::begin() const noexcept 241 | { 242 | return (void**)_memory.data(); 243 | } 244 | 245 | void* DeviceMemory::at(const int& index) const noexcept 246 | { 247 | return _memory[index]; 248 | } 249 | 250 | Result DeviceMemory::setup(const std::shared_ptr& logger, 251 | std::unique_ptr& engine, 252 | DeviceMemory* output) noexcept 253 | { 254 | const int32_t nbBindings = engine->getNbBindings(); 255 | for(int i = 0; i < nbBindings; ++i) 256 | { 257 | const nvinfer1::Dims dims = engine->getBindingDimensions(i); 258 | const int volume = dimsVolume(dims); 259 | 260 | try 261 | { 262 | output->_memory.push_back(nullptr); 263 | } 264 | catch(const std::exception& e) 265 | { 266 | logger->logf(LOGGING_ERROR, "[DeviceMemory] setup() failure: " 267 | "exception: %s", e.what()); 268 | return RESULT_FAILURE_ALLOC; 269 | } 270 | void** ptr = &output->_memory.back(); 271 | 272 | auto r = cudaMalloc(ptr, volume * sizeof(float)); 273 | if(r != 0 || *ptr == nullptr) 274 | { 275 | logger->logf(LOGGING_ERROR, "[DeviceMemory] setup() failure: " 276 | "could not allocate device memory: %s", 277 | cudaGetErrorString(r)); 278 | return RESULT_FAILURE_CUDA_ERROR; 279 | } 280 | 281 | } 282 | return RESULT_SUCCESS; 283 | } 284 | 285 | bool opencvHasCuda() noexcept 286 | { 287 | int r = 0; 288 | try 289 | { 290 | r = cv::cuda::getCudaEnabledDeviceCount(); 291 | } 292 | catch(const std::exception& e) 293 | { 294 | return 0; 295 | } 296 | return (r > 0); 297 | } 298 | 299 | PreprocessorTransform::PreprocessorTransform() noexcept 300 | : _inputSize(0, 0), _f(1), _leftWidth(0), _topHeight(0) 301 | { 302 | } 303 | 304 | PreprocessorTransform::PreprocessorTransform(const cv::Size& inputSize, 305 | const double& f, const int& leftWidth, 306 | const int& topHeight) noexcept 307 | : _inputSize(inputSize), _f(f), 308 | _leftWidth(leftWidth), _topHeight(topHeight) 309 | { 310 | } 311 | 312 | PreprocessorTransform::~PreprocessorTransform() noexcept 313 | { 314 | } 315 | 316 | cv::Rect PreprocessorTransform::transformBbox( 317 | const cv::Rect& input) const noexcept 318 | { 319 | cv::Rect r; 320 | r.x = (input.x - _leftWidth) / _f; 321 | r.x = MAX(0, MIN(r.x, _inputSize.width-1)); 322 | 323 | r.y = (input.y - _topHeight) / _f; 324 | r.y = MAX(0, MIN(r.y, _inputSize.height-1)); 325 | 326 | r.width = input.width / _f; 327 | if(r.x + r.width > _inputSize.width) 328 | { 329 | r.width = _inputSize.width - r.x; 330 | } 331 | r.height = input.height / _f; 332 | if(r.y + r.height > _inputSize.height) 333 | { 334 | r.height = _inputSize.height - r.y; 335 | } 336 | return r; 337 | } 338 | 339 | 340 | Preprocessor::Preprocessor() noexcept 341 | { 342 | } 343 | 344 | Preprocessor::~Preprocessor() noexcept 345 | { 346 | } 347 | 348 | void Preprocessor::setLogger(std::shared_ptr logger) noexcept 349 | { 350 | _logger = logger; 351 | } 352 | 353 | bool Preprocessor::process(const int& index, const cv::Mat& input, 354 | const bool& last) noexcept 355 | { 356 | YOLOV5_UNUSED(input); 357 | YOLOV5_UNUSED(last); 358 | 359 | if(_transforms.size() < (unsigned int)index+1) 360 | { 361 | try 362 | { 363 | _transforms.resize(index + 1); 364 | } 365 | catch(const std::exception& e) 366 | { 367 | _logger->logf(LOGGING_ERROR, "[Preprocessor] process() " 368 | "failure: got exception setting up transforms: " 369 | "%s", e.what()); 370 | return false; 371 | } 372 | } 373 | return true; 374 | } 375 | 376 | bool Preprocessor::process(const int& index, 377 | const cv::cuda::GpuMat& input, const bool& last) noexcept 378 | { 379 | YOLOV5_UNUSED(input); 380 | YOLOV5_UNUSED(last); 381 | 382 | if(_transforms.size() < (unsigned int)index+1) 383 | { 384 | try 385 | { 386 | _transforms.resize(index + 1); 387 | } 388 | catch(const std::exception& e) 389 | { 390 | _logger->logf(LOGGING_ERROR, "[Preprocessor] process() " 391 | "failure: got exception setting up transforms: " 392 | "%s", e.what()); 393 | return false; 394 | } 395 | } 396 | return true; 397 | } 398 | 399 | cv::Rect Preprocessor::transformBbox(const int& index, 400 | const cv::Rect& bbox) const noexcept 401 | { 402 | return _transforms[index].transformBbox(bbox); 403 | } 404 | 405 | template 406 | static void setupChannels(const cv::Size& size, 407 | const Preprocessor::InputType& inputType, 408 | float* inputPtr, 409 | std::vector& channels) 410 | { 411 | const int channelSize = size.area(); 412 | if(inputType == Preprocessor::INPUTTYPE_BGR) /* INPUT_BGR */ 413 | { 414 | /* B channel will go here */ 415 | channels.push_back(T(size, CV_32FC1, inputPtr + 2 * channelSize)); 416 | /* G channel will go here */ 417 | channels.push_back(T(size, CV_32FC1, inputPtr + 1 * channelSize)); 418 | /* R channel will go here */ 419 | channels.push_back(T(size, CV_32FC1, inputPtr)); 420 | } 421 | else /* INPUTTYPE_RGB */ 422 | { 423 | /* R channel will go here */ 424 | channels.push_back(T(size, CV_32FC1, inputPtr)); 425 | /* G channel will go here */ 426 | channels.push_back(T(size, CV_32FC1, inputPtr + 1 * channelSize)); 427 | /* B channel will go here */ 428 | channels.push_back(T(size, CV_32FC1, inputPtr + 2 * channelSize)); 429 | } 430 | } 431 | 432 | CvCpuPreprocessor::CvCpuPreprocessor() noexcept 433 | : _cudaStream(nullptr), _lastType((InputType)-1), _lastBatchSize(-1), 434 | _networkCols(0), _networkRows(0) 435 | { 436 | } 437 | 438 | CvCpuPreprocessor::~CvCpuPreprocessor() noexcept 439 | { 440 | } 441 | 442 | bool CvCpuPreprocessor::setup(const nvinfer1::Dims& inputDims, 443 | const int& flags, const int& batchSize, 444 | float* inputMemory) noexcept 445 | { 446 | if(!_cudaStream) 447 | { 448 | auto r = cudaStreamCreate(&_cudaStream); 449 | if(r != 0) 450 | { 451 | _logger->logf(LOGGING_ERROR, "[CvCpuPreprocessor] setup() " 452 | "failure: could not create cuda stream: %s", 453 | cudaGetErrorString(r)); 454 | return false; 455 | } 456 | } 457 | 458 | if((flags & INPUT_RGB) && (flags & INPUT_BGR)) 459 | { 460 | _logger->log(LOGGING_ERROR, "[CvCpuPreprocessor] setup() " 461 | "failure: both INPUT_RGB and INPUT_BGR flags specified"); 462 | return false; 463 | } 464 | InputType inputType = INPUTTYPE_BGR; 465 | if(flags & INPUT_RGB) 466 | { 467 | inputType = INPUTTYPE_RGB; 468 | } 469 | 470 | if(_lastType == inputType && _lastBatchSize == batchSize) 471 | { 472 | return true; 473 | } 474 | _lastType = inputType; 475 | _lastBatchSize = batchSize; 476 | 477 | _networkRows = inputDims.d[2]; 478 | _networkCols = inputDims.d[3]; 479 | const cv::Size networkSize(_networkCols, _networkRows); 480 | 481 | _deviceInputMemory = inputMemory; 482 | 483 | _inputChannels.clear(); 484 | try 485 | { 486 | /* Set up host input memory */ 487 | _hostInputMemory.resize(dimsVolume(inputDims)); 488 | 489 | _inputChannels.resize(batchSize); 490 | for(unsigned int i = 0; i < _inputChannels.size(); ++i) 491 | { 492 | float* inputPtr = _hostInputMemory.data() 493 | + i * networkSize.area() * 3; 494 | 495 | std::vector& channels = _inputChannels[i]; 496 | setupChannels(networkSize, inputType, inputPtr, channels); 497 | } 498 | } 499 | catch(const std::exception& e) 500 | { 501 | _logger->logf(LOGGING_ERROR, "[CvCpuPreprocessor] setup() failure: " 502 | "got exception while trying to set up input channels: %s", 503 | e.what()); 504 | return false; 505 | } 506 | return true; 507 | } 508 | 509 | void CvCpuPreprocessor::reset() noexcept 510 | { 511 | /* this will trigger setup() to take effect next time */ 512 | _lastType = (InputType)-1; 513 | } 514 | 515 | bool CvCpuPreprocessor::process(const int& index, 516 | const cv::Mat& input, 517 | const bool& last) noexcept 518 | { 519 | if(!Preprocessor::process(index, input, last)) 520 | { 521 | return false; 522 | } 523 | 524 | PreprocessorTransform& transform = _transforms[index]; 525 | try 526 | { 527 | if(input.rows == _networkRows && input.cols == _networkCols) 528 | { 529 | input.convertTo(_buffer3, CV_32FC3, 1.0f / 255.0f); 530 | transform = PreprocessorTransform(input.size(), 1.0, 0, 0); 531 | } 532 | else 533 | { 534 | const double f = 535 | MIN((double)_networkRows / (double)input.rows, 536 | (double)_networkCols / (double)input.cols); 537 | const cv::Size boxSize = cv::Size(input.cols * f, 538 | input.rows * f); 539 | 540 | const int dr = _networkRows - boxSize.height; 541 | const int dc = _networkCols - boxSize.width; 542 | const int topHeight = std::floor(dr / 2.0); 543 | const int bottomHeight = std::ceil(dr / 2.0); 544 | const int leftWidth = std::floor(dc / 2.0); 545 | const int rightWidth = std::ceil(dc / 2.0); 546 | 547 | transform = PreprocessorTransform(input.size(), 548 | f, leftWidth, topHeight); 549 | 550 | cv::resize(input, _buffer1, boxSize, 0, 0, cv::INTER_LINEAR); 551 | cv::copyMakeBorder(_buffer1, _buffer2, 552 | topHeight, bottomHeight, leftWidth, rightWidth, 553 | cv::BORDER_CONSTANT, cv::Scalar(0, 0, 0)); 554 | _buffer2.convertTo(_buffer3, CV_32FC3, 1.0f / 255.0f); 555 | } 556 | cv::split(_buffer3, _inputChannels[index]); 557 | } 558 | catch(const std::exception& e) 559 | { 560 | _logger->logf(LOGGING_ERROR, "[CvCpuPreprocessor] process() " 561 | "failure: got exception setting up input: " 562 | "%s", e.what()); 563 | return false; 564 | } 565 | 566 | /* Copy from host to device */ 567 | if(last) 568 | { 569 | const int volume = _inputChannels.size() /* batch size */ 570 | * 3 * _buffer3.size().area(); /* channels * rows*cols */ 571 | 572 | auto r = cudaMemcpyAsync(_deviceInputMemory, 573 | (void*)_hostInputMemory.data(), 574 | (int)(volume * sizeof(float)), 575 | cudaMemcpyHostToDevice, _cudaStream); 576 | if(r != 0) 577 | { 578 | _logger->logf(LOGGING_ERROR, "[CvCpuPreprocessor] process() " 579 | "failure: could not set up host-to-device transfer " 580 | "for input: %s", cudaGetErrorString(r)); 581 | return false; 582 | } 583 | } 584 | return true; 585 | } 586 | 587 | bool CvCpuPreprocessor::process(const int& index, 588 | const cv::cuda::GpuMat& input, 589 | const bool& last) noexcept 590 | { 591 | YOLOV5_UNUSED(index); 592 | YOLOV5_UNUSED(input); 593 | YOLOV5_UNUSED(last); 594 | _logger->log(LOGGING_ERROR, "[CvCpuPreprocessor] process() failure: CUDA " 595 | "method is not supported. Use CPU variant instead"); 596 | /* not supported */ 597 | return false; 598 | } 599 | 600 | cudaStream_t CvCpuPreprocessor::cudaStream() const noexcept 601 | { 602 | return _cudaStream; 603 | } 604 | 605 | bool CvCpuPreprocessor::synchronizeCudaStream() noexcept 606 | { 607 | auto r = cudaStreamSynchronize(_cudaStream); 608 | if(r != 0) 609 | { 610 | _logger->logf(LOGGING_ERROR, "[CvCpuPreprocessor] " 611 | "synchronizeCudaStream() failure: %s", cudaGetErrorString(r)); 612 | return false; 613 | } 614 | return true; 615 | } 616 | 617 | CvCudaPreprocessor::CvCudaPreprocessor() noexcept 618 | : _lastType((InputType)-1), _lastBatchSize(-1), 619 | _networkCols(0), _networkRows(0) 620 | { 621 | } 622 | 623 | CvCudaPreprocessor::~CvCudaPreprocessor() noexcept 624 | { 625 | } 626 | 627 | bool CvCudaPreprocessor::setup(const nvinfer1::Dims& inputDims, 628 | const int& flags, const int& batchSize, 629 | float* inputMemory) noexcept 630 | { 631 | #ifdef YOLOV5_OPENCV_HAS_CUDA 632 | if((flags & INPUT_RGB) && (flags & INPUT_BGR)) 633 | { 634 | _logger->log(LOGGING_ERROR, "[CvCudaPreprocessor] setup() " 635 | "failure: both INPUT_RGB and INPUT_BGR flags specified"); 636 | return false; 637 | } 638 | InputType inputType = INPUTTYPE_BGR; 639 | if(flags & INPUT_RGB) 640 | { 641 | inputType = INPUTTYPE_RGB; 642 | } 643 | 644 | if(_lastType == inputType && _lastBatchSize == batchSize) 645 | { 646 | return true; 647 | } 648 | _lastType = inputType; 649 | _lastBatchSize = batchSize; 650 | 651 | _networkRows = inputDims.d[2]; 652 | _networkCols = inputDims.d[3]; 653 | const cv::Size networkSize(_networkCols, _networkRows); 654 | 655 | _inputChannels.clear(); 656 | try 657 | { 658 | _inputChannels.resize(batchSize); 659 | for(unsigned int i = 0; i < _inputChannels.size(); ++i) 660 | { 661 | float* inputPtr = inputMemory + i * networkSize.area() * 3; 662 | 663 | std::vector& channels = _inputChannels[i]; 664 | setupChannels(networkSize, inputType, inputPtr, 665 | channels); 666 | } 667 | } 668 | catch(const std::exception& e) 669 | { 670 | _logger->logf(LOGGING_ERROR, "[CvCudaPreprocessor] setup() failure: " 671 | "got exception while trying to set up input channels: %s", 672 | e.what()); 673 | return false; 674 | } 675 | return true; 676 | #else 677 | YOLOV5_UNUSED(inputDims); 678 | YOLOV5_UNUSED(flags); 679 | YOLOV5_UNUSED(batchSize); 680 | YOLOV5_UNUSED(inputMemory); 681 | _logger->log(LOGGING_ERROR, "[CvCudaPreprocessor] setup() failure: " 682 | "OpenCV without CUDA support"); 683 | return false; 684 | #endif 685 | } 686 | 687 | void CvCudaPreprocessor::reset() noexcept 688 | { 689 | /* this will trigger setup() to take effect next time */ 690 | _lastType = (InputType)-1; 691 | } 692 | 693 | bool CvCudaPreprocessor::process(const int& index, 694 | const cv::Mat& input, const bool& last) noexcept 695 | { 696 | #ifdef YOLOV5_OPENCV_HAS_CUDA 697 | try 698 | { 699 | _buffer0.upload(input); 700 | } 701 | catch(const std::exception& e) 702 | { 703 | _logger->logf(LOGGING_ERROR, "[CvCudaPreprocessor] process() " 704 | "failure: got exception trying to upload input " 705 | "to CUDA device: %s", e.what()); 706 | return false; 707 | } 708 | return process(index, _buffer0, last); 709 | #else 710 | YOLOV5_UNUSED(index); 711 | YOLOV5_UNUSED(input); 712 | YOLOV5_UNUSED(last); 713 | _logger->log(LOGGING_ERROR, "[CvCudaPreprocessor] process() failure: " 714 | "OpenCV without CUDA support"); 715 | return false; 716 | #endif 717 | } 718 | 719 | bool CvCudaPreprocessor::process(const int& index, 720 | const cv::cuda::GpuMat& input, 721 | const bool& last) noexcept 722 | { 723 | #ifdef YOLOV5_OPENCV_HAS_CUDA 724 | if(!Preprocessor::process(index, input, last)) 725 | { 726 | return false; 727 | } 728 | 729 | if(index >= 1) 730 | { 731 | /* If we're processing a batch, first process the previous image 732 | before changing the internal buffers */ 733 | if(!synchronizeCudaStream()) 734 | { 735 | return false; 736 | } 737 | } 738 | 739 | PreprocessorTransform& transform = _transforms[index]; 740 | try 741 | { 742 | if(input.rows == _networkRows && input.cols == _networkCols) 743 | { 744 | input.convertTo(_buffer3, CV_32FC3, 1.0f / 255.0f, _cudaStream); 745 | transform = PreprocessorTransform(input.size(), 1.0, 0, 0); 746 | } 747 | else 748 | { 749 | const double f = 750 | MIN((double)_networkRows / (double)input.rows, 751 | (double)_networkCols / (double)input.cols); 752 | const cv::Size boxSize = cv::Size(input.cols * f, 753 | input.rows * f); 754 | 755 | const int dr = _networkRows - boxSize.height; 756 | const int dc = _networkCols - boxSize.width; 757 | const int topHeight = std::floor(dr / 2.0); 758 | const int bottomHeight = std::ceil(dr / 2.0); 759 | const int leftWidth = std::floor(dc / 2.0); 760 | const int rightWidth = std::ceil(dc / 2.0); 761 | 762 | transform = PreprocessorTransform(input.size(), 763 | f, leftWidth, topHeight); 764 | 765 | cv::cuda::resize(input, _buffer1, boxSize, 0, 0, cv::INTER_LINEAR, 766 | _cudaStream); 767 | cv::cuda::copyMakeBorder(_buffer1, _buffer2, 768 | topHeight, bottomHeight, leftWidth, rightWidth, 769 | cv::BORDER_CONSTANT, cv::Scalar(0, 0, 0), 770 | _cudaStream); 771 | _buffer2.convertTo(_buffer3, CV_32FC3, 1.0f / 255.0f, _cudaStream); 772 | } 773 | cv::cuda::split(_buffer3, _inputChannels[index], _cudaStream); 774 | } 775 | catch(const std::exception& e) 776 | { 777 | _logger->logf(LOGGING_ERROR, "[CvCudaPreprocessor] process() " 778 | "failure: got exception setting up input: " 779 | "%s", e.what()); 780 | return false; 781 | } 782 | return true; 783 | #else 784 | YOLOV5_UNUSED(index); 785 | YOLOV5_UNUSED(input); 786 | YOLOV5_UNUSED(last); 787 | _logger->log(LOGGING_ERROR, "[CvCudaPreprocessor] process() failure: " 788 | "OpenCV without CUDA support"); 789 | return false; 790 | #endif 791 | } 792 | 793 | cudaStream_t CvCudaPreprocessor::cudaStream() const noexcept 794 | { 795 | #ifdef YOLOV5_OPENCV_HAS_CUDA 796 | return cv::cuda::StreamAccessor::getStream(_cudaStream); 797 | #else 798 | return nullptr; 799 | #endif 800 | } 801 | 802 | bool CvCudaPreprocessor::synchronizeCudaStream() noexcept 803 | { 804 | try 805 | { 806 | _cudaStream.waitForCompletion(); 807 | } 808 | catch(const std::exception& e) 809 | { 810 | _logger->logf(LOGGING_ERROR, "[CvCudaPreprocessor] " 811 | "synchronizeCudaStream() failure: %s", e.what()); 812 | return false; 813 | } 814 | return true; 815 | } 816 | 817 | 818 | } /* namespace internal */ 819 | 820 | } /* namespace yolov5 */ -------------------------------------------------------------------------------- /src/yolov5_logging.cpp: -------------------------------------------------------------------------------- 1 | /** 2 | * @file 3 | * 4 | * @author Noah van der Meer 5 | * @brief YoloV5 inference through TensorRT (logging) 6 | * 7 | * 8 | * Copyright (c) 2021, Noah van der Meer 9 | * 10 | * Permission is hereby granted, free of charge, to any person obtaining a copy 11 | * of this software and associated documentation files (the "Software"), to 12 | * deal in the Software without restriction, including without limitation the 13 | * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 14 | * sell copies of the Software, and to permit persons to whom the Software is 15 | * furnished to do so, subject to the following conditions: 16 | * 17 | * The above copyright notice and this permission notice shall be included in 18 | * all copies or substantial portions of the Software. 19 | * 20 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 21 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 22 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 23 | * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 24 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 25 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 26 | * IN THE SOFTWARE. 27 | * 28 | */ 29 | 30 | #include "yolov5_logging.hpp" 31 | 32 | #include 33 | #include 34 | 35 | const char* yolov5::loglevel_to_string(const LogLevel& l) noexcept 36 | { 37 | if(l == LOGGING_DEBUG) 38 | { 39 | return "debug"; 40 | } 41 | else if(l == LOGGING_INFO) 42 | { 43 | return "info"; 44 | } 45 | else if(l == LOGGING_WARNING) 46 | { 47 | return "warning"; 48 | } 49 | else if(l == LOGGING_ERROR) 50 | { 51 | return "error"; 52 | } 53 | else 54 | { 55 | return ""; 56 | } 57 | } 58 | 59 | bool yolov5::loglevel_to_string(const LogLevel& l, std::string* out) noexcept 60 | { 61 | const char* str = loglevel_to_string(l); 62 | if(std::strlen(str) == 0) 63 | { 64 | return false; 65 | } 66 | 67 | if(out != nullptr) 68 | { 69 | try 70 | { 71 | *out = std::string(str); 72 | } 73 | catch(const std::exception& e) 74 | { 75 | return false; 76 | } 77 | } 78 | return true; 79 | } 80 | 81 | yolov5::Logger::Logger() noexcept 82 | { 83 | } 84 | 85 | yolov5::Logger::~Logger() 86 | { 87 | } 88 | 89 | void yolov5::Logger::print(const LogLevel& level, const char* msg) 90 | { 91 | std::printf("|yolov5|%s|%s\n", loglevel_to_string(level), msg); 92 | } 93 | 94 | void yolov5::Logger::log(const LogLevel& level, const char* msg) noexcept 95 | { 96 | try 97 | { 98 | this->print(level, msg); 99 | } 100 | catch(...) 101 | { 102 | } 103 | } 104 | 105 | void yolov5::Logger::logf(const LogLevel& level, 106 | const char* fmt, ...) noexcept 107 | { 108 | char* msg = nullptr; 109 | 110 | va_list args; 111 | va_start(args, fmt); 112 | 113 | if(vasprintf(&msg, fmt, args) < 0) 114 | { 115 | msg = nullptr; 116 | } 117 | va_end(args); 118 | 119 | if(msg) 120 | { 121 | try 122 | { 123 | this->print(level, msg); 124 | } 125 | catch(...) 126 | { 127 | } 128 | std::free(msg); 129 | } 130 | } 131 | 132 | yolov5::TensorRT_Logger::TensorRT_Logger() noexcept 133 | { 134 | } 135 | 136 | yolov5::TensorRT_Logger::TensorRT_Logger( 137 | std::shared_ptr logger) noexcept 138 | : _logger(logger) /* shared_ptr copy constructor is marked noexcept */ 139 | { 140 | } 141 | 142 | yolov5::TensorRT_Logger::~TensorRT_Logger() 143 | { 144 | } 145 | 146 | void yolov5::TensorRT_Logger::setLogger( 147 | std::shared_ptr logger) noexcept 148 | { 149 | _logger = logger; 150 | } 151 | 152 | void yolov5::TensorRT_Logger::log(nvinfer1::ILogger::Severity severity, 153 | const char* msg) noexcept 154 | { 155 | if(!_logger) 156 | { 157 | return; 158 | } 159 | 160 | LogLevel level = LOGGING_DEBUG; 161 | if(severity == nvinfer1::ILogger::Severity::kINTERNAL_ERROR || 162 | severity == nvinfer1::ILogger::Severity::kERROR) 163 | { 164 | level = LOGGING_ERROR; 165 | } 166 | else if(severity == nvinfer1::ILogger::Severity::kWARNING) 167 | { 168 | level = LOGGING_WARNING; 169 | } 170 | else if(severity == nvinfer1::ILogger::Severity::kINFO) 171 | { 172 | level = LOGGING_INFO; 173 | } 174 | else if(severity == nvinfer1::ILogger::Severity::kVERBOSE) 175 | { 176 | /* LOGGING_DEBUG */ 177 | } 178 | 179 | 180 | try 181 | { 182 | _logger->logf(level, "[TensorRT] %s", msg); 183 | } 184 | catch(...) 185 | { 186 | } 187 | } -------------------------------------------------------------------------------- /yolov5-tensorrt.pc.in: -------------------------------------------------------------------------------- 1 | prefix=@CMAKE_INSTALL_PREFIX@ 2 | exec_prefix=@CMAKE_INSTALL_PREFIX@ 3 | includedir=${exec_prefix}/@CMAKE_INSTALL_INCLUDEDIR@ 4 | libdir=${prefix}/@CMAKE_INSTALL_LIBDIR@ 5 | 6 | Name: @PROJECT_NAME@ 7 | Description: @PROJECT_DESCRIPTION@ 8 | URL: https://github.com/noahmr/yolov5-tensorrt 9 | Version: @PROJECT_VERSION@ 10 | Cflags: -I${includedir}/yolov5-tensorrt 11 | Libs: -L${libdir} -lyolov5-tensorrt --------------------------------------------------------------------------------