├── .gitignore ├── .pre-commit-config.yaml ├── LICENSE ├── README.md ├── build.py ├── config.py ├── csrc ├── cls │ └── normal │ │ ├── CMakeLists.txt │ │ ├── cmake │ │ ├── FindTensorRT.cmake │ │ └── Function.cmake │ │ ├── include │ │ ├── common.hpp │ │ ├── filesystem.hpp │ │ └── yolov8-cls.hpp │ │ └── main.cpp ├── deepstream │ ├── CMakeLists.txt │ ├── README.md │ ├── config_yoloV8.txt │ ├── custom_bbox_parser │ │ └── nvdsparsebbox_yoloV8.cpp │ ├── deepstream_app_config.txt │ └── labels.txt ├── detect │ ├── end2end │ │ ├── CMakeLists.txt │ │ ├── cmake │ │ │ ├── FindTensorRT.cmake │ │ │ └── Function.cmake │ │ ├── include │ │ │ ├── common.hpp │ │ │ ├── filesystem.hpp │ │ │ └── yolov8.hpp │ │ └── main.cpp │ └── normal │ │ ├── CMakeLists.txt │ │ ├── cmake │ │ ├── FindTensorRT.cmake │ │ └── Function.cmake │ │ ├── include │ │ ├── common.hpp │ │ ├── filesystem.hpp │ │ └── yolov8.hpp │ │ └── main.cpp ├── jetson │ ├── detect │ │ ├── CMakeLists.txt │ │ ├── include │ │ │ ├── common.hpp │ │ │ ├── filesystem.hpp │ │ │ └── yolov8.hpp │ │ └── main.cpp │ ├── pose │ │ ├── CMakeLists.txt │ │ ├── include │ │ │ ├── common.hpp │ │ │ ├── filesystem.hpp │ │ │ └── yolov8-pose.hpp │ │ └── main.cpp │ └── segment │ │ ├── CMakeLists.txt │ │ ├── include │ │ ├── common.hpp │ │ ├── filesystem.hpp │ │ └── yolov8-seg.hpp │ │ └── main.cpp ├── obb │ └── normal │ │ ├── CMakeLists.txt │ │ ├── cmake │ │ ├── FindTensorRT.cmake │ │ └── Function.cmake │ │ ├── include │ │ ├── common.hpp │ │ ├── filesystem.hpp │ │ └── yolov8-obb.hpp │ │ └── main.cpp ├── pose │ └── normal │ │ ├── CMakeLists.txt │ │ ├── cmake │ │ ├── FindTensorRT.cmake │ │ └── Function.cmake │ │ ├── include │ │ ├── common.hpp │ │ ├── filesystem.hpp │ │ └── yolov8-pose.hpp │ │ └── main.cpp └── segment │ ├── normal │ ├── CMakeLists.txt │ ├── cmake │ │ ├── FindTensorRT.cmake │ │ └── Function.cmake │ ├── include │ │ ├── common.hpp │ │ ├── filesystem.hpp │ │ └── yolov8-seg.hpp │ └── main.cpp │ └── simple │ ├── CMakeLists.txt │ ├── cmake │ ├── FindTensorRT.cmake │ └── Function.cmake │ ├── include │ ├── common.hpp │ ├── filesystem.hpp │ └── yolov8-seg.hpp │ └── main.cpp ├── data ├── bus.jpg └── zidane.jpg ├── docs ├── API-Build.md ├── Cls.md ├── Jetson.md ├── Normal.md ├── Obb.md ├── Pose.md ├── Segment.md └── star.md ├── export-det.py ├── export-seg.py ├── gen_pkl.py ├── infer-cls-without-torch.py ├── infer-cls.py ├── infer-det-without-torch.py ├── infer-det.py ├── infer-obb-without-torch.py ├── infer-obb.py ├── infer-pose-without-torch.py ├── infer-pose.py ├── infer-seg-without-torch.py ├── infer-seg.py ├── models ├── __init__.py ├── api.py ├── common.py ├── cudart_api.py ├── engine.py ├── pycuda_api.py ├── torch_utils.py └── utils.py ├── requirements.txt └── trt-profile.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Personal 2 | .idea 3 | *.engine 4 | *.pt 5 | *.pth 6 | *.onnx 7 | *.jpg 8 | *.jpeg 9 | *.png 10 | *.bmp 11 | *.plan 12 | .clang-format 13 | 14 | # Byte-compiled / optimized / DLL files 15 | __pycache__/ 16 | *.py[cod] 17 | *$py.class 18 | 19 | # C extensions 20 | *.so 21 | 22 | # Distribution / packaging 23 | .Python 24 | build/ 25 | cmake-build-debug 26 | develop-eggs/ 27 | dist/ 28 | downloads/ 29 | eggs/ 30 | .eggs/ 31 | lib/ 32 | lib64/ 33 | parts/ 34 | sdist/ 35 | var/ 36 | wheels/ 37 | pip-wheel-metadata/ 38 | share/python-wheels/ 39 | *.egg-info/ 40 | .installed.cfg 41 | *.egg 42 | MANIFEST 43 | 44 | # PyInstaller 45 | # Usually these files are written by a python script from a template 46 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 47 | *.manifest 48 | *.spec 49 | 50 | # Installer logs 51 | pip-log.txt 52 | pip-delete-this-directory.txt 53 | 54 | # Unit test / coverage reports 55 | htmlcov/ 56 | .tox/ 57 | .nox/ 58 | .coverage 59 | .coverage.* 60 | .cache 61 | nosetests.xml 62 | coverage.xml 63 | *.cover 64 | *.py,cover 65 | .hypothesis/ 66 | .pytest_cache/ 67 | 68 | # Translations 69 | *.mo 70 | *.pot 71 | 72 | # Django stuff: 73 | *.log 74 | local_settings.py 75 | db.sqlite3 76 | db.sqlite3-journal 77 | 78 | # Flask stuff: 79 | instance/ 80 | .webassets-cache 81 | 82 | # Scrapy stuff: 83 | .scrapy 84 | 85 | # Sphinx documentation 86 | docs/_build/ 87 | 88 | # PyBuilder 89 | target/ 90 | 91 | # Jupyter Notebook 92 | .ipynb_checkpoints 93 | 94 | # IPython 95 | profile_default/ 96 | ipython_config.py 97 | 98 | # pyenv 99 | .python-version 100 | 101 | # pipenv 102 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 103 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 104 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 105 | # install all needed dependencies. 106 | #Pipfile.lock 107 | 108 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 109 | __pypackages__/ 110 | 111 | # Celery stuff 112 | celerybeat-schedule 113 | celerybeat.pid 114 | 115 | # SageMath parsed files 116 | *.sage.py 117 | 118 | # Environments 119 | .env 120 | .venv 121 | env/ 122 | venv/ 123 | ENV/ 124 | env.bak/ 125 | venv.bak/ 126 | 127 | # Spyder project settings 128 | .spyderproject 129 | .spyproject 130 | 131 | # Rope project settings 132 | .ropeproject 133 | 134 | # mkdocs documentation 135 | /site 136 | 137 | # mypy 138 | .mypy_cache/ 139 | .dmypy.json 140 | dmypy.json 141 | 142 | # Pyre type checker 143 | .pyre/ 144 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | repos: 2 | - repo: https://github.com/PyCQA/flake8 3 | rev: 5.0.4 4 | hooks: 5 | - id: flake8 6 | - repo: https://github.com/PyCQA/isort 7 | rev: 5.11.5 8 | hooks: 9 | - id: isort 10 | - repo: https://github.com/pre-commit/mirrors-yapf 11 | rev: v0.32.0 12 | hooks: 13 | - id: yapf 14 | - repo: https://github.com/pre-commit/pre-commit-hooks 15 | rev: v4.3.0 16 | hooks: 17 | - id: trailing-whitespace 18 | - id: end-of-file-fixer 19 | - id: requirements-txt-fixer 20 | - id: double-quote-string-fixer 21 | - id: check-merge-conflict 22 | - id: fix-encoding-pragma 23 | args: ["--remove"] 24 | - id: mixed-line-ending 25 | args: ["--fix=lf"] 26 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 tripleMu 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 deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | 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 all 13 | 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 FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /build.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | from models import EngineBuilder 4 | 5 | 6 | def parse_args(): 7 | parser = argparse.ArgumentParser() 8 | parser.add_argument('--weights', 9 | type=str, 10 | required=True, 11 | help='Weights file') 12 | parser.add_argument('--iou-thres', 13 | type=float, 14 | default=0.65, 15 | help='IOU threshoud for NMS plugin') 16 | parser.add_argument('--conf-thres', 17 | type=float, 18 | default=0.25, 19 | help='CONF threshoud for NMS plugin') 20 | parser.add_argument('--topk', 21 | type=int, 22 | default=100, 23 | help='Max number of detection bboxes') 24 | parser.add_argument('--input-shape', 25 | nargs='+', 26 | type=int, 27 | default=[1, 3, 640, 640], 28 | help='Model input shape only for api builder') 29 | parser.add_argument('--fp16', 30 | action='store_true', 31 | help='Build model with fp16 mode') 32 | parser.add_argument('--device', 33 | type=str, 34 | default='cuda:0', 35 | help='TensorRT builder device') 36 | parser.add_argument('--seg', 37 | action='store_true', 38 | help='Build seg model by onnx') 39 | args = parser.parse_args() 40 | assert len(args.input_shape) == 4 41 | return args 42 | 43 | 44 | def main(args): 45 | builder = EngineBuilder(args.weights, args.device) 46 | builder.seg = args.seg 47 | builder.build(fp16=args.fp16, 48 | input_shape=args.input_shape, 49 | iou_thres=args.iou_thres, 50 | conf_thres=args.conf_thres, 51 | topk=args.topk) 52 | 53 | 54 | if __name__ == '__main__': 55 | args = parse_args() 56 | main(args) 57 | -------------------------------------------------------------------------------- /csrc/cls/normal/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.12) 2 | 3 | set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86 89 90) 4 | set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc) 5 | 6 | project(yolov8-cls LANGUAGES CXX CUDA) 7 | 8 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -O3") 9 | set(CMAKE_CXX_STANDARD 14) 10 | set(CMAKE_BUILD_TYPE Release) 11 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 12 | 13 | list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake") 14 | include(Function) 15 | 16 | # CUDA 17 | find_package(CUDA REQUIRED) 18 | print_var(CUDA_LIBRARIES) 19 | print_var(CUDA_INCLUDE_DIRS) 20 | get_filename_component(CUDA_LIB_DIR ${CUDA_LIBRARIES} DIRECTORY) 21 | print_var(CUDA_LIB_DIR) 22 | 23 | # OpenCV 24 | find_package(OpenCV REQUIRED) 25 | print_var(OpenCV_LIBS) 26 | print_var(OpenCV_LIBRARIES) 27 | print_var(OpenCV_INCLUDE_DIRS) 28 | 29 | # TensorRT 30 | find_package(TensorRT REQUIRED) 31 | print_var(TensorRT_LIBRARIES) 32 | print_var(TensorRT_INCLUDE_DIRS) 33 | print_var(TensorRT_LIB_DIR) 34 | if (TensorRT_VERSION_MAJOR GREATER_EQUAL 10) 35 | message(STATUS "Build with -DTRT_10") 36 | add_definitions(-DTRT_10) 37 | endif () 38 | 39 | list(APPEND ALL_INCLUDE_DIRS 40 | ${CUDA_INCLUDE_DIRS} 41 | ${OpenCV_INCLUDE_DIRS} 42 | ${TensorRT_INCLUDE_DIRS} 43 | ${CMAKE_CURRENT_SOURCE_DIR}/include 44 | ) 45 | 46 | list(APPEND ALL_LIBS 47 | ${CUDA_LIBRARIES} 48 | ${OpenCV_LIBRARIES} 49 | ${TensorRT_LIBRARIES} 50 | ) 51 | 52 | list(APPEND ALL_LIB_DIRS 53 | ${CUDA_LIB_DIR} 54 | ${TensorRT_LIB_DIR} 55 | ) 56 | 57 | print_var(ALL_INCLUDE_DIRS) 58 | print_var(ALL_LIBS) 59 | print_var(ALL_LIB_DIRS) 60 | 61 | 62 | add_executable( 63 | ${PROJECT_NAME} 64 | ${CMAKE_CURRENT_SOURCE_DIR}/main.cpp 65 | ${CMAKE_CURRENT_SOURCE_DIR}/include/yolov8-cls.hpp 66 | ${CMAKE_CURRENT_SOURCE_DIR}/include/common.hpp 67 | ) 68 | 69 | target_include_directories( 70 | ${PROJECT_NAME} 71 | PUBLIC 72 | ${ALL_INCLUDE_DIRS} 73 | ) 74 | 75 | target_link_directories( 76 | ${PROJECT_NAME} 77 | PUBLIC 78 | ${ALL_LIB_DIRS} 79 | ) 80 | 81 | target_link_libraries( 82 | ${PROJECT_NAME} 83 | PRIVATE 84 | ${ALL_LIBS} 85 | ) 86 | -------------------------------------------------------------------------------- /csrc/cls/normal/cmake/FindTensorRT.cmake: -------------------------------------------------------------------------------- 1 | # This module defines the following variables: 2 | # 3 | # :: 4 | # 5 | # TensorRT_INCLUDE_DIRS 6 | # TensorRT_LIBRARIES 7 | # TensorRT_FOUND 8 | # 9 | # :: 10 | # 11 | # TensorRT_VERSION_STRING - version (x.y.z) 12 | # TensorRT_VERSION_MAJOR - major version (x) 13 | # TensorRT_VERSION_MINOR - minor version (y) 14 | # TensorRT_VERSION_PATCH - patch version (z) 15 | # 16 | # Hints 17 | # ^^^^^ 18 | # A user may set ``TensorRT_ROOT`` to an installation root to tell this module where to look. 19 | # 20 | set(_TensorRT_SEARCHES) 21 | 22 | if(TensorRT_ROOT) 23 | set(_TensorRT_SEARCH_ROOT PATHS ${TensorRT_ROOT} NO_DEFAULT_PATH) 24 | list(APPEND _TensorRT_SEARCHES _TensorRT_SEARCH_ROOT) 25 | endif() 26 | 27 | # appends some common paths 28 | set(_TensorRT_SEARCH_NORMAL 29 | PATHS "/usr" 30 | ) 31 | list(APPEND _TensorRT_SEARCHES _TensorRT_SEARCH_NORMAL) 32 | 33 | # Include dir 34 | foreach(search ${_TensorRT_SEARCHES}) 35 | find_path(TensorRT_INCLUDE_DIR NAMES NvInfer.h ${${search}} PATH_SUFFIXES include) 36 | endforeach() 37 | 38 | if(NOT TensorRT_LIBRARY) 39 | foreach(search ${_TensorRT_SEARCHES}) 40 | find_library(TensorRT_LIBRARY NAMES nvinfer ${${search}} PATH_SUFFIXES lib) 41 | if(NOT TensorRT_LIB_DIR) 42 | get_filename_component(TensorRT_LIB_DIR ${TensorRT_LIBRARY} DIRECTORY) 43 | endif () 44 | endforeach() 45 | endif() 46 | 47 | if(NOT TensorRT_nvinfer_plugin_LIBRARY) 48 | foreach(search ${_TensorRT_SEARCHES}) 49 | find_library(TensorRT_nvinfer_plugin_LIBRARY NAMES nvinfer_plugin ${${search}} PATH_SUFFIXES lib) 50 | endforeach() 51 | endif() 52 | 53 | mark_as_advanced(TensorRT_INCLUDE_DIR) 54 | 55 | if(TensorRT_INCLUDE_DIR AND EXISTS "${TensorRT_INCLUDE_DIR}/NvInfer.h") 56 | if(EXISTS "${TensorRT_INCLUDE_DIR}/NvInferVersion.h") 57 | set(_VersionSearchFile "${TensorRT_INCLUDE_DIR}/NvInferVersion.h") 58 | else () 59 | set(_VersionSearchFile "${TensorRT_INCLUDE_DIR}/NvInfer.h") 60 | endif () 61 | file(STRINGS "${_VersionSearchFile}" TensorRT_MAJOR REGEX "^#define NV_TENSORRT_MAJOR [0-9]+.*$") 62 | file(STRINGS "${_VersionSearchFile}" TensorRT_MINOR REGEX "^#define NV_TENSORRT_MINOR [0-9]+.*$") 63 | file(STRINGS "${_VersionSearchFile}" TensorRT_PATCH REGEX "^#define NV_TENSORRT_PATCH [0-9]+.*$") 64 | 65 | string(REGEX REPLACE "^#define NV_TENSORRT_MAJOR ([0-9]+).*$" "\\1" TensorRT_VERSION_MAJOR "${TensorRT_MAJOR}") 66 | string(REGEX REPLACE "^#define NV_TENSORRT_MINOR ([0-9]+).*$" "\\1" TensorRT_VERSION_MINOR "${TensorRT_MINOR}") 67 | string(REGEX REPLACE "^#define NV_TENSORRT_PATCH ([0-9]+).*$" "\\1" TensorRT_VERSION_PATCH "${TensorRT_PATCH}") 68 | set(TensorRT_VERSION_STRING "${TensorRT_VERSION_MAJOR}.${TensorRT_VERSION_MINOR}.${TensorRT_VERSION_PATCH}") 69 | endif() 70 | 71 | include(FindPackageHandleStandardArgs) 72 | FIND_PACKAGE_HANDLE_STANDARD_ARGS(TensorRT REQUIRED_VARS TensorRT_LIBRARY TensorRT_INCLUDE_DIR VERSION_VAR TensorRT_VERSION_STRING) 73 | 74 | if(TensorRT_FOUND) 75 | set(TensorRT_INCLUDE_DIRS ${TensorRT_INCLUDE_DIR}) 76 | 77 | if(NOT TensorRT_LIBRARIES) 78 | set(TensorRT_LIBRARIES ${TensorRT_LIBRARY}) 79 | if (TensorRT_nvinfer_plugin_LIBRARY) 80 | list(APPEND TensorRT_LIBRARIES ${TensorRT_nvinfer_plugin_LIBRARY}) 81 | endif() 82 | endif() 83 | 84 | if(NOT TARGET TensorRT::TensorRT) 85 | add_library(TensorRT INTERFACE IMPORTED) 86 | add_library(TensorRT::TensorRT ALIAS TensorRT) 87 | endif() 88 | 89 | if(NOT TARGET TensorRT::nvinfer) 90 | add_library(TensorRT::nvinfer SHARED IMPORTED) 91 | if (WIN32) 92 | foreach(search ${_TensorRT_SEARCHES}) 93 | find_file(TensorRT_LIBRARY_DLL 94 | NAMES nvinfer.dll 95 | PATHS ${${search}} 96 | PATH_SUFFIXES bin 97 | ) 98 | endforeach() 99 | 100 | set_target_properties(TensorRT::nvinfer PROPERTIES 101 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 102 | IMPORTED_LOCATION "${TensorRT_LIBRARY_DLL}" 103 | IMPORTED_IMPLIB "${TensorRT_LIBRARY}" 104 | ) 105 | else() 106 | set_target_properties(TensorRT::nvinfer PROPERTIES 107 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 108 | IMPORTED_LOCATION "${TensorRT_LIBRARY}" 109 | ) 110 | endif() 111 | target_link_libraries(TensorRT INTERFACE TensorRT::nvinfer) 112 | endif() 113 | 114 | if(NOT TARGET TensorRT::nvinfer_plugin AND TensorRT_nvinfer_plugin_LIBRARY) 115 | add_library(TensorRT::nvinfer_plugin SHARED IMPORTED) 116 | if (WIN32) 117 | foreach(search ${_TensorRT_SEARCHES}) 118 | find_file(TensorRT_nvinfer_plugin_LIBRARY_DLL 119 | NAMES nvinfer_plugin.dll 120 | PATHS ${${search}} 121 | PATH_SUFFIXES bin 122 | ) 123 | endforeach() 124 | 125 | set_target_properties(TensorRT::nvinfer_plugin PROPERTIES 126 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 127 | IMPORTED_LOCATION "${TensorRT_nvinfer_plugin_LIBRARY_DLL}" 128 | IMPORTED_IMPLIB "${TensorRT_nvinfer_plugin_LIBRARY}" 129 | ) 130 | else() 131 | set_target_properties(TensorRT::nvinfer_plugin PROPERTIES 132 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 133 | IMPORTED_LOCATION "${TensorRT_nvinfer_plugin_LIBRARY}" 134 | ) 135 | endif() 136 | target_link_libraries(TensorRT INTERFACE TensorRT::nvinfer_plugin) 137 | endif() 138 | endif() 139 | -------------------------------------------------------------------------------- /csrc/cls/normal/cmake/Function.cmake: -------------------------------------------------------------------------------- 1 | 2 | function(print_var var) 3 | set(value "${${var}}") 4 | string(LENGTH "${value}" value_length) 5 | if(value_length GREATER 0) 6 | math(EXPR last_index "${value_length} - 1") 7 | string(SUBSTRING "${value}" ${last_index} ${last_index} last_char) 8 | endif() 9 | 10 | if(NOT "${last_char}" STREQUAL "\n") 11 | set(value "${value}\n") 12 | endif() 13 | message(STATUS "${var}:\n ${value}") 14 | endfunction() 15 | -------------------------------------------------------------------------------- /csrc/cls/normal/include/common.hpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 4/27/24. 3 | // 4 | 5 | #ifndef CLS_NORMAL_COMMON_HPP 6 | #define CLS_NORMAL_COMMON_HPP 7 | #include "NvInfer.h" 8 | #include "filesystem.hpp" 9 | #include "opencv2/opencv.hpp" 10 | 11 | #define CHECK(call) \ 12 | do { \ 13 | const cudaError_t error_code = call; \ 14 | if (error_code != cudaSuccess) { \ 15 | printf("CUDA Error:\n"); \ 16 | printf(" File: %s\n", __FILE__); \ 17 | printf(" Line: %d\n", __LINE__); \ 18 | printf(" Error code: %d\n", error_code); \ 19 | printf(" Error text: %s\n", cudaGetErrorString(error_code)); \ 20 | exit(1); \ 21 | } \ 22 | } while (0) 23 | 24 | class Logger: public nvinfer1::ILogger { 25 | public: 26 | nvinfer1::ILogger::Severity reportableSeverity; 27 | 28 | explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO): 29 | reportableSeverity(severity) 30 | { 31 | } 32 | 33 | void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override 34 | { 35 | if (severity > reportableSeverity) { 36 | return; 37 | } 38 | switch (severity) { 39 | case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: 40 | std::cerr << "INTERNAL_ERROR: "; 41 | break; 42 | case nvinfer1::ILogger::Severity::kERROR: 43 | std::cerr << "ERROR: "; 44 | break; 45 | case nvinfer1::ILogger::Severity::kWARNING: 46 | std::cerr << "WARNING: "; 47 | break; 48 | case nvinfer1::ILogger::Severity::kINFO: 49 | std::cerr << "INFO: "; 50 | break; 51 | default: 52 | std::cerr << "VERBOSE: "; 53 | break; 54 | } 55 | std::cerr << msg << std::endl; 56 | } 57 | }; 58 | 59 | inline int get_size_by_dims(const nvinfer1::Dims& dims) 60 | { 61 | int size = 1; 62 | for (int i = 0; i < dims.nbDims; i++) { 63 | size *= dims.d[i]; 64 | } 65 | return size; 66 | } 67 | 68 | inline int type_to_size(const nvinfer1::DataType& dataType) 69 | { 70 | switch (dataType) { 71 | case nvinfer1::DataType::kFLOAT: 72 | return 4; 73 | case nvinfer1::DataType::kHALF: 74 | return 2; 75 | case nvinfer1::DataType::kINT32: 76 | return 4; 77 | case nvinfer1::DataType::kINT8: 78 | return 1; 79 | case nvinfer1::DataType::kBOOL: 80 | return 1; 81 | default: 82 | return 4; 83 | } 84 | } 85 | 86 | inline static float clamp(float val, float min, float max) 87 | { 88 | return val > min ? (val < max ? val : max) : min; 89 | } 90 | 91 | namespace cls { 92 | struct Binding { 93 | size_t size = 1; 94 | size_t dsize = 1; 95 | nvinfer1::Dims dims; 96 | std::string name; 97 | }; 98 | 99 | struct Object { 100 | int label = 0; 101 | float prob = 0.0; 102 | }; 103 | } // namespace cls 104 | #endif // CLS_NORMAL_COMMON_HPP 105 | -------------------------------------------------------------------------------- /csrc/deepstream/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.1) 2 | 3 | set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86 89 90) 4 | set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc) 5 | 6 | project(nvdsinfer_custom_bbox_yoloV8 LANGUAGES CXX CUDA) 7 | 8 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -O3 -Wall -Werror -shared -fPIC") 9 | set(CMAKE_CXX_STANDARD 14) 10 | set(CMAKE_BUILD_TYPE Release) 11 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 12 | 13 | # CUDA 14 | find_package(CUDA REQUIRED) 15 | message(STATUS "CUDA Libs: \n${CUDA_LIBRARIES}\n") 16 | get_filename_component(CUDA_LIB_DIR ${CUDA_LIBRARIES} DIRECTORY) 17 | message(STATUS "CUDA Headers: \n${CUDA_INCLUDE_DIRS}\n") 18 | 19 | 20 | # TensorRT 21 | set(TensorRT_INCLUDE_DIRS /usr/include/x86_64-linux-gnu) 22 | set(TensorRT_LIBRARIES /usr/lib/x86_64-linux-gnu) 23 | 24 | message(STATUS "TensorRT Libs: \n${TensorRT_LIBRARIES}\n") 25 | message(STATUS "TensorRT Headers: \n${TensorRT_INCLUDE_DIRS}\n") 26 | 27 | 28 | set(DEEPSTREAM /opt/nvidia/deepstream/deepstream) 29 | set(DS_LIBRARIES ${DEEPSTREAM}/lib) 30 | set(DS_INCLUDE_DIRS ${DEEPSTREAM}/sources/includes) 31 | message(STATUS "DS Libs: \n${DS_LIBRARIES}\n") 32 | message(STATUS "DS Headers: \n${DS_INCLUDE_DIRS}\n") 33 | 34 | 35 | 36 | list(APPEND INCLUDE_DIRS 37 | ${CUDA_INCLUDE_DIRS} 38 | ${TensorRT_INCLUDE_DIRS} 39 | ${DS_INCLUDE_DIRS} 40 | ) 41 | 42 | list(APPEND ALL_LIBS 43 | ${CUDA_LIBRARIES} 44 | ${CUDA_LIB_DIR} 45 | ${TensorRT_LIBRARIES} 46 | ${DS_LIBRARIES} 47 | ) 48 | 49 | include_directories(${INCLUDE_DIRS}) 50 | 51 | add_library( 52 | ${PROJECT_NAME} 53 | SHARED 54 | custom_bbox_parser/nvdsparsebbox_yoloV8.cpp 55 | ) 56 | 57 | target_link_libraries(${PROJECT_NAME} PRIVATE nvinfer nvinfer_plugin) 58 | -------------------------------------------------------------------------------- /csrc/deepstream/README.md: -------------------------------------------------------------------------------- 1 | # Start DeepStream Using the engine build from [`YOLOv8-TensorRT`](https://github.com/triple-Mu/YOLOv8-TensorRT) 2 | 3 | ## 1. Build you own TensorRT engine from `trtexec` or [`build.py`](https://github.com/triple-Mu/YOLOv8-TensorRT/blob/main/build.py) 4 | 5 | For example, if you have built an engine named `yolov8s.engine`. 6 | 7 | ## 2. Compile deepstream plugin 8 | 9 | First, modify the [`CMakeLists.txt`](https://github.com/triple-Mu/YOLOv8-TensorRT/blob/main/csrc/deepstream/CMakeLists.txt) 10 | 11 | ```cmake 12 | # Set your own TensorRT path 13 | set(TensorRT_INCLUDE_DIRS /usr/include/x86_64-linux-gnu) 14 | set(TensorRT_LIBRARIES /usr/lib/x86_64-linux-gnu) 15 | # Set your own DeepStream path 16 | set(DEEPSTREAM /opt/nvidia/deepstream/deepstream) 17 | ``` 18 | 19 | Second, build deepstream plugin 20 | 21 | ```shell 22 | mkdir build 23 | cd build 24 | cmake .. 25 | make 26 | ``` 27 | You will get a lib `libnvdsinfer_custom_bbox_yoloV8.so` in `build`. 28 | 29 | 30 | ## 3. Modify the deepstream config 31 | 32 | The net config is [`config_yoloV8.txt`](config_yoloV8.txt). Please modify by your own model. 33 | 34 | ```text 35 | net-scale-factor=0.0039215697906911373 # the normalize param == 1/255 36 | model-engine-file=./yolov8s.engine # the engine path you build 37 | labelfile-path=./labels.txt # the class name path 38 | num-detected-classes=80 # the number of classes 39 | output-blob-names=num_dets;bboxes;scores;labels # the model output names 40 | custom-lib-path=./build/libnvdsinfer_custom_bbox_yoloV8.so # the deepstream plugin you build 41 | ``` 42 | 43 | The deepstream config is [`deepstream_app_config.txt`](deepstream_app_config.txt). 44 | 45 | ```text 46 | **** 47 | [source0] 48 | enable=1 49 | #Type - 1=CameraV4L2 2=URI 3=MultiURI 50 | type=3 51 | uri=file://./sample_1080p_h264.mp4 # the video path or stream you want to detect 52 | **** 53 | **** 54 | config-file=config_yoloV8.txt # the net config path 55 | ``` 56 | 57 | You can get more information from [`deepstream offical`](https://developer.nvidia.com/deepstream-sdk). 58 | 59 | ## 4. Runing detector ! 60 | 61 | ```shell 62 | deepstream-app -c deepstream_app_config.txt 63 | ``` 64 | -------------------------------------------------------------------------------- /csrc/deepstream/config_yoloV8.txt: -------------------------------------------------------------------------------- 1 | ################################################################################ 2 | # Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. 3 | # 4 | # Permission is hereby granted, free of charge, to any person obtaining a 5 | # copy of this software and associated documentation files (the "Software"), 6 | # to deal in the Software without restriction, including without limitation 7 | # the rights to use, copy, modify, merge, publish, distribute, sublicense, 8 | # and/or sell copies of the Software, and to permit persons to whom the 9 | # Software is furnished to do so, subject to the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be included in 12 | # all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 17 | # THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 20 | # DEALINGS IN THE SOFTWARE. 21 | ################################################################################ 22 | 23 | # Following properties are mandatory when engine files are not specified: 24 | # int8-calib-file(Only in INT8), model-file-format 25 | # Caffemodel mandatory properties: model-file, proto-file, output-blob-names 26 | # UFF: uff-file, input-dims, uff-input-blob-name, output-blob-names 27 | # ONNX: onnx-file 28 | # 29 | # Mandatory properties for detectors: 30 | # num-detected-classes 31 | # 32 | # Optional properties for detectors: 33 | # enable-dbscan(Default=false), interval(Primary mode only, Default=0) 34 | # custom-lib-path 35 | # parse-bbox-func-name 36 | # 37 | # Mandatory properties for classifiers: 38 | # classifier-threshold, is-classifier 39 | # 40 | # Optional properties for classifiers: 41 | # classifier-async-mode(Secondary mode only, Default=false) 42 | # 43 | # Optional properties in secondary mode: 44 | # operate-on-gie-id(Default=0), operate-on-class-ids(Defaults to all classes), 45 | # input-object-min-width, input-object-min-height, input-object-max-width, 46 | # input-object-max-height 47 | # 48 | # Following properties are always recommended: 49 | # batch-size(Default=1) 50 | # 51 | # Other optional properties: 52 | # net-scale-factor(Default=1), network-mode(Default=0 i.e FP32), 53 | # model-color-format(Default=0 i.e. RGB) model-engine-file, labelfile-path, 54 | # mean-file, gie-unique-id(Default=0), offsets, gie-mode (Default=1 i.e. primary), 55 | # custom-lib-path, network-mode(Default=0 i.e FP32) 56 | # 57 | # The values in the config file are overridden by values set through GObject 58 | # properties. 59 | 60 | [property] 61 | net-scale-factor=0.0039215697906911373 62 | gpu-id=0 63 | #0=RGB, 1=BGR 64 | model-color-format=0 65 | model-engine-file=./yolov8s.engine 66 | labelfile-path=./labels.txt 67 | ## 0=FP32, 1=INT8, 2=FP16 mode 68 | network-mode=2 69 | num-detected-classes=80 70 | gie-unique-id=1 71 | is-classifier=0 72 | maintain-aspect-ratio=1 73 | output-blob-names=num_dets;bboxes;scores;labels 74 | parse-bbox-func-name=NvDsInferParseCustomYoloV8 75 | custom-lib-path=./build/libnvdsinfer_custom_bbox_yoloV8.so 76 | -------------------------------------------------------------------------------- /csrc/deepstream/custom_bbox_parser/nvdsparsebbox_yoloV8.cpp: -------------------------------------------------------------------------------- 1 | #include "nvdsinfer_custom_impl.h" 2 | #include 3 | #include 4 | 5 | /** 6 | * Function expected by DeepStream for decoding the TinyYOLOv2 output. 7 | * 8 | * C-linkage [extern "C"] was written to prevent name-mangling. This function must return true after 9 | * adding all bounding boxes to the objectList vector. 10 | * 11 | * @param [outputLayersInfo] std::vector of NvDsInferLayerInfo objects with information about the output layer. 12 | * @param [networkInfo] NvDsInferNetworkInfo object with information about the TinyYOLOv2 network. 13 | * @param [detectionParams] NvDsInferParseDetectionParams with information about some config params. 14 | * @param [objectList] std::vector of NvDsInferParseObjectInfo objects to which bounding box information must 15 | * be stored. 16 | * 17 | * @return true 18 | */ 19 | 20 | // This is just the function prototype. The definition is written at the end of the file. 21 | extern "C" bool NvDsInferParseCustomYoloV8(std::vector const& outputLayersInfo, 22 | NvDsInferNetworkInfo const& networkInfo, 23 | NvDsInferParseDetectionParams const& detectionParams, 24 | std::vector& objectList); 25 | 26 | static __inline__ float bbox_clip(const float& val, const float& minVal = 0.f, const float& maxVal = 1280.f) 27 | { 28 | assert(minVal <= maxVal); 29 | return std::max(std::min(val, (maxVal - 1)), minVal); 30 | } 31 | 32 | static std::vector decodeYoloV8Tensor(const int* num_dets, 33 | const float* bboxes, 34 | const float* scores, 35 | const int* labels, 36 | const unsigned int& img_w, 37 | const unsigned int& img_h) 38 | { 39 | std::vector bboxInfo; 40 | size_t nums = num_dets[0]; 41 | for (size_t i = 0; i < nums; i++) { 42 | float x0 = (bboxes[i * 4]); 43 | float y0 = (bboxes[i * 4 + 1]); 44 | float x1 = (bboxes[i * 4 + 2]); 45 | float y1 = (bboxes[i * 4 + 3]); 46 | x0 = bbox_clip(x0, 0.f, img_w); 47 | y0 = bbox_clip(y0, 0.f, img_h); 48 | x1 = bbox_clip(x1, 0.f, img_w); 49 | y1 = bbox_clip(y1, 0.f, img_h); 50 | NvDsInferParseObjectInfo obj; 51 | obj.left = x0; 52 | obj.top = y0; 53 | obj.width = x1 - x0; 54 | obj.height = y1 - y0; 55 | obj.detectionConfidence = scores[i]; 56 | obj.classId = labels[i]; 57 | bboxInfo.push_back(obj); 58 | } 59 | 60 | return bboxInfo; 61 | } 62 | 63 | /* C-linkage to prevent name-mangling */ 64 | extern "C" bool NvDsInferParseCustomYoloV8(std::vector const& outputLayersInfo, 65 | NvDsInferNetworkInfo const& networkInfo, 66 | NvDsInferParseDetectionParams const& detectionParams, 67 | std::vector& objectList) 68 | { 69 | 70 | // Some assertions and error checking. 71 | 72 | if (outputLayersInfo.empty() || outputLayersInfo.size() != 4) { 73 | std::cerr << "Could not find output layer in bbox parsing" << std::endl; 74 | return false; 75 | } 76 | 77 | // Obtaining the output layer. 78 | const NvDsInferLayerInfo& num_dets = outputLayersInfo[0]; 79 | const NvDsInferLayerInfo& bboxes = outputLayersInfo[1]; 80 | const NvDsInferLayerInfo& scores = outputLayersInfo[2]; 81 | const NvDsInferLayerInfo& labels = outputLayersInfo[3]; 82 | 83 | // num_dets(int) bboxes(float) scores(float) labels(int) 84 | assert(num_dets.dims.numDims == 2); 85 | assert(bboxes.dims.numDims == 3); 86 | assert(scores.dims.numDims == 2); 87 | assert(labels.dims.numDims == 2); 88 | 89 | // Decoding the output tensor of YOLOv8 to the NvDsInferParseObjectInfo format. 90 | std::vector objects = decodeYoloV8Tensor((const int*)(num_dets.buffer), 91 | (const float*)(bboxes.buffer), 92 | (const float*)(scores.buffer), 93 | (const int*)(labels.buffer), 94 | networkInfo.width, 95 | networkInfo.height); 96 | 97 | objectList.clear(); 98 | objectList = objects; 99 | return true; 100 | } 101 | 102 | /* Check that the custom function has been defined correctly */ 103 | CHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseCustomYoloV8); 104 | -------------------------------------------------------------------------------- /csrc/deepstream/deepstream_app_config.txt: -------------------------------------------------------------------------------- 1 | [application] 2 | enable-perf-measurement=1 3 | perf-measurement-interval-sec=5 4 | #gie-kitti-output-dir=streamscl 5 | 6 | [tiled-display] 7 | enable=1 8 | rows=1 9 | columns=1 10 | # show 11 | width=1920 12 | height=1080 13 | gpu-id=0 14 | #(0): nvbuf-mem-default - Default memory allocated, specific to particular platform 15 | #(1): nvbuf-mem-cuda-pinned - Allocate Pinned/Host cuda memory, applicable for Tesla 16 | #(2): nvbuf-mem-cuda-device - Allocate Device cuda memory, applicable for Tesla 17 | #(3): nvbuf-mem-cuda-unified - Allocate Unified cuda memory, applicable for Tesla 18 | #(4): nvbuf-mem-surface-array - Allocate Surface Array memory, applicable for Jetson 19 | nvbuf-memory-type=0 20 | 21 | [source0] 22 | enable=1 23 | #Type - 1=CameraV4L2 2=URI 3=MultiURI 24 | type=3 25 | uri=file://./sample_1080p_h264.mp4 26 | num-sources=1 27 | drop-frame-interval=1 28 | gpu-id=0 29 | # (0): memtype_device - Memory type Device 30 | # (1): memtype_pinned - Memory type Host Pinned 31 | # (2): memtype_unified - Memory type Unified 32 | cudadec-memtype=0 33 | 34 | [sink0] 35 | enable=1 36 | #Type - 1=FakeSink 2=EglSink 3=File 37 | type=2 38 | sync=0 39 | source-id=0 40 | gpu-id=0 41 | nvbuf-memory-type=0 42 | 43 | #[sink1] 44 | #enable=1 45 | #Type - 1=FakeSink 2=EglSink 3=File 4=RTSPStreaming 46 | #type=4 47 | ##1=h264 2=h265 48 | #codec=1 49 | #sync=0 50 | #bitrate=4000000 51 | ## set below properties in case of RTSPStreaming 52 | #rtsp-port=8554 53 | #udp-port=5400 54 | 55 | #[sink2] 56 | #enable=1 57 | ##Type - 1=FakeSink 2=EglSink 3=File 4=RTSPStreaming 58 | #type=3 59 | ##1=h264 2=h265 60 | #codec=1 61 | #sync=0 62 | #bitrate=4000000 63 | #container=1 64 | #output-file=./output.mp4 65 | 66 | [osd] 67 | enable=1 68 | gpu-id=0 69 | border-width=1 70 | text-size=15 71 | text-color=1;1;1;1; 72 | text-bg-color=0.3;0.3;0.3;1 73 | font=Serif 74 | show-clock=0 75 | clock-x-offset=800 76 | clock-y-offset=820 77 | clock-text-size=12 78 | clock-color=1;0;0;1 79 | nvbuf-memory-type=0 80 | 81 | [streammux] 82 | gpu-id=0 83 | ##Boolean property to inform muxer that sources are live 84 | live-source=0 85 | batch-size=1 86 | ##time out in usec, to wait after the first buffer is available 87 | ##to push the batch even if the complete batch is not formed 88 | batched-push-timeout=40000 89 | ## Set muxer output width and height 90 | width=1920 91 | height=1080 92 | ##Enable to maintain aspect ratio wrt source, and allow black borders, works 93 | ##along with width, height properties 94 | enable-padding=0 95 | nvbuf-memory-type=0 96 | 97 | # config-file property is mandatory for any gie section. 98 | # Other properties are optional and if set will override the properties set in 99 | # the infer config file. 100 | [primary-gie] 101 | enable=1 102 | gpu-id=0 103 | model-engine-file=yolov8s.engine 104 | labelfile-path=./labels.txt 105 | batch-size=1 106 | #Required by the app for OSD, not a plugin property 107 | bbox-border-color0=1;0;0;1 108 | bbox-border-color1=0;1;1;1 109 | bbox-border-color2=0;0;1;1 110 | bbox-border-color3=0;1;0;1 111 | gie-unique-id=1 112 | nvbuf-memory-type=0 113 | config-file=config_yoloV8.txt 114 | 115 | [tests] 116 | file-loop=0 117 | -------------------------------------------------------------------------------- /csrc/deepstream/labels.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 81 | -------------------------------------------------------------------------------- /csrc/detect/end2end/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.12) 2 | 3 | set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86 89 90) 4 | set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc) 5 | 6 | project(yolov8 LANGUAGES CXX CUDA) 7 | 8 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -O3") 9 | set(CMAKE_CXX_STANDARD 14) 10 | set(CMAKE_BUILD_TYPE Release) 11 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 12 | 13 | list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake") 14 | include(Function) 15 | 16 | # CUDA 17 | find_package(CUDA REQUIRED) 18 | print_var(CUDA_LIBRARIES) 19 | print_var(CUDA_INCLUDE_DIRS) 20 | get_filename_component(CUDA_LIB_DIR ${CUDA_LIBRARIES} DIRECTORY) 21 | print_var(CUDA_LIB_DIR) 22 | 23 | # OpenCV 24 | find_package(OpenCV REQUIRED) 25 | print_var(OpenCV_LIBS) 26 | print_var(OpenCV_LIBRARIES) 27 | print_var(OpenCV_INCLUDE_DIRS) 28 | 29 | # TensorRT 30 | find_package(TensorRT REQUIRED) 31 | print_var(TensorRT_LIBRARIES) 32 | print_var(TensorRT_INCLUDE_DIRS) 33 | print_var(TensorRT_LIB_DIR) 34 | if (TensorRT_VERSION_MAJOR GREATER_EQUAL 10) 35 | message(STATUS "Build with -DTRT_10") 36 | add_definitions(-DTRT_10) 37 | endif () 38 | 39 | list(APPEND ALL_INCLUDE_DIRS 40 | ${CUDA_INCLUDE_DIRS} 41 | ${OpenCV_INCLUDE_DIRS} 42 | ${TensorRT_INCLUDE_DIRS} 43 | ${CMAKE_CURRENT_SOURCE_DIR}/include 44 | ) 45 | 46 | list(APPEND ALL_LIBS 47 | ${CUDA_LIBRARIES} 48 | ${OpenCV_LIBRARIES} 49 | ${TensorRT_LIBRARIES} 50 | ) 51 | 52 | list(APPEND ALL_LIB_DIRS 53 | ${CUDA_LIB_DIR} 54 | ${TensorRT_LIB_DIR} 55 | ) 56 | 57 | print_var(ALL_INCLUDE_DIRS) 58 | print_var(ALL_LIBS) 59 | print_var(ALL_LIB_DIRS) 60 | 61 | 62 | add_executable( 63 | ${PROJECT_NAME} 64 | ${CMAKE_CURRENT_SOURCE_DIR}/main.cpp 65 | ${CMAKE_CURRENT_SOURCE_DIR}/include/yolov8.hpp 66 | ${CMAKE_CURRENT_SOURCE_DIR}/include/common.hpp 67 | ) 68 | 69 | target_include_directories( 70 | ${PROJECT_NAME} 71 | PUBLIC 72 | ${ALL_INCLUDE_DIRS} 73 | ) 74 | 75 | target_link_directories( 76 | ${PROJECT_NAME} 77 | PUBLIC 78 | ${ALL_LIB_DIRS} 79 | ) 80 | 81 | target_link_libraries( 82 | ${PROJECT_NAME} 83 | PRIVATE 84 | ${ALL_LIBS} 85 | ) 86 | -------------------------------------------------------------------------------- /csrc/detect/end2end/cmake/FindTensorRT.cmake: -------------------------------------------------------------------------------- 1 | # This module defines the following variables: 2 | # 3 | # :: 4 | # 5 | # TensorRT_INCLUDE_DIRS 6 | # TensorRT_LIBRARIES 7 | # TensorRT_FOUND 8 | # 9 | # :: 10 | # 11 | # TensorRT_VERSION_STRING - version (x.y.z) 12 | # TensorRT_VERSION_MAJOR - major version (x) 13 | # TensorRT_VERSION_MINOR - minor version (y) 14 | # TensorRT_VERSION_PATCH - patch version (z) 15 | # 16 | # Hints 17 | # ^^^^^ 18 | # A user may set ``TensorRT_ROOT`` to an installation root to tell this module where to look. 19 | # 20 | set(_TensorRT_SEARCHES) 21 | 22 | if(TensorRT_ROOT) 23 | set(_TensorRT_SEARCH_ROOT PATHS ${TensorRT_ROOT} NO_DEFAULT_PATH) 24 | list(APPEND _TensorRT_SEARCHES _TensorRT_SEARCH_ROOT) 25 | endif() 26 | 27 | # appends some common paths 28 | set(_TensorRT_SEARCH_NORMAL 29 | PATHS "/usr" 30 | ) 31 | list(APPEND _TensorRT_SEARCHES _TensorRT_SEARCH_NORMAL) 32 | 33 | # Include dir 34 | foreach(search ${_TensorRT_SEARCHES}) 35 | find_path(TensorRT_INCLUDE_DIR NAMES NvInfer.h ${${search}} PATH_SUFFIXES include) 36 | endforeach() 37 | 38 | if(NOT TensorRT_LIBRARY) 39 | foreach(search ${_TensorRT_SEARCHES}) 40 | find_library(TensorRT_LIBRARY NAMES nvinfer ${${search}} PATH_SUFFIXES lib) 41 | if(NOT TensorRT_LIB_DIR) 42 | get_filename_component(TensorRT_LIB_DIR ${TensorRT_LIBRARY} DIRECTORY) 43 | endif () 44 | endforeach() 45 | endif() 46 | 47 | if(NOT TensorRT_nvinfer_plugin_LIBRARY) 48 | foreach(search ${_TensorRT_SEARCHES}) 49 | find_library(TensorRT_nvinfer_plugin_LIBRARY NAMES nvinfer_plugin ${${search}} PATH_SUFFIXES lib) 50 | endforeach() 51 | endif() 52 | 53 | mark_as_advanced(TensorRT_INCLUDE_DIR) 54 | 55 | if(TensorRT_INCLUDE_DIR AND EXISTS "${TensorRT_INCLUDE_DIR}/NvInfer.h") 56 | if(EXISTS "${TensorRT_INCLUDE_DIR}/NvInferVersion.h") 57 | set(_VersionSearchFile "${TensorRT_INCLUDE_DIR}/NvInferVersion.h") 58 | else () 59 | set(_VersionSearchFile "${TensorRT_INCLUDE_DIR}/NvInfer.h") 60 | endif () 61 | file(STRINGS "${_VersionSearchFile}" TensorRT_MAJOR REGEX "^#define NV_TENSORRT_MAJOR [0-9]+.*$") 62 | file(STRINGS "${_VersionSearchFile}" TensorRT_MINOR REGEX "^#define NV_TENSORRT_MINOR [0-9]+.*$") 63 | file(STRINGS "${_VersionSearchFile}" TensorRT_PATCH REGEX "^#define NV_TENSORRT_PATCH [0-9]+.*$") 64 | 65 | string(REGEX REPLACE "^#define NV_TENSORRT_MAJOR ([0-9]+).*$" "\\1" TensorRT_VERSION_MAJOR "${TensorRT_MAJOR}") 66 | string(REGEX REPLACE "^#define NV_TENSORRT_MINOR ([0-9]+).*$" "\\1" TensorRT_VERSION_MINOR "${TensorRT_MINOR}") 67 | string(REGEX REPLACE "^#define NV_TENSORRT_PATCH ([0-9]+).*$" "\\1" TensorRT_VERSION_PATCH "${TensorRT_PATCH}") 68 | set(TensorRT_VERSION_STRING "${TensorRT_VERSION_MAJOR}.${TensorRT_VERSION_MINOR}.${TensorRT_VERSION_PATCH}") 69 | endif() 70 | 71 | include(FindPackageHandleStandardArgs) 72 | FIND_PACKAGE_HANDLE_STANDARD_ARGS(TensorRT REQUIRED_VARS TensorRT_LIBRARY TensorRT_INCLUDE_DIR VERSION_VAR TensorRT_VERSION_STRING) 73 | 74 | if(TensorRT_FOUND) 75 | set(TensorRT_INCLUDE_DIRS ${TensorRT_INCLUDE_DIR}) 76 | 77 | if(NOT TensorRT_LIBRARIES) 78 | set(TensorRT_LIBRARIES ${TensorRT_LIBRARY}) 79 | if (TensorRT_nvinfer_plugin_LIBRARY) 80 | list(APPEND TensorRT_LIBRARIES ${TensorRT_nvinfer_plugin_LIBRARY}) 81 | endif() 82 | endif() 83 | 84 | if(NOT TARGET TensorRT::TensorRT) 85 | add_library(TensorRT INTERFACE IMPORTED) 86 | add_library(TensorRT::TensorRT ALIAS TensorRT) 87 | endif() 88 | 89 | if(NOT TARGET TensorRT::nvinfer) 90 | add_library(TensorRT::nvinfer SHARED IMPORTED) 91 | if (WIN32) 92 | foreach(search ${_TensorRT_SEARCHES}) 93 | find_file(TensorRT_LIBRARY_DLL 94 | NAMES nvinfer.dll 95 | PATHS ${${search}} 96 | PATH_SUFFIXES bin 97 | ) 98 | endforeach() 99 | 100 | set_target_properties(TensorRT::nvinfer PROPERTIES 101 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 102 | IMPORTED_LOCATION "${TensorRT_LIBRARY_DLL}" 103 | IMPORTED_IMPLIB "${TensorRT_LIBRARY}" 104 | ) 105 | else() 106 | set_target_properties(TensorRT::nvinfer PROPERTIES 107 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 108 | IMPORTED_LOCATION "${TensorRT_LIBRARY}" 109 | ) 110 | endif() 111 | target_link_libraries(TensorRT INTERFACE TensorRT::nvinfer) 112 | endif() 113 | 114 | if(NOT TARGET TensorRT::nvinfer_plugin AND TensorRT_nvinfer_plugin_LIBRARY) 115 | add_library(TensorRT::nvinfer_plugin SHARED IMPORTED) 116 | if (WIN32) 117 | foreach(search ${_TensorRT_SEARCHES}) 118 | find_file(TensorRT_nvinfer_plugin_LIBRARY_DLL 119 | NAMES nvinfer_plugin.dll 120 | PATHS ${${search}} 121 | PATH_SUFFIXES bin 122 | ) 123 | endforeach() 124 | 125 | set_target_properties(TensorRT::nvinfer_plugin PROPERTIES 126 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 127 | IMPORTED_LOCATION "${TensorRT_nvinfer_plugin_LIBRARY_DLL}" 128 | IMPORTED_IMPLIB "${TensorRT_nvinfer_plugin_LIBRARY}" 129 | ) 130 | else() 131 | set_target_properties(TensorRT::nvinfer_plugin PROPERTIES 132 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 133 | IMPORTED_LOCATION "${TensorRT_nvinfer_plugin_LIBRARY}" 134 | ) 135 | endif() 136 | target_link_libraries(TensorRT INTERFACE TensorRT::nvinfer_plugin) 137 | endif() 138 | endif() 139 | -------------------------------------------------------------------------------- /csrc/detect/end2end/cmake/Function.cmake: -------------------------------------------------------------------------------- 1 | 2 | function(print_var var) 3 | set(value "${${var}}") 4 | string(LENGTH "${value}" value_length) 5 | if(value_length GREATER 0) 6 | math(EXPR last_index "${value_length} - 1") 7 | string(SUBSTRING "${value}" ${last_index} ${last_index} last_char) 8 | endif() 9 | 10 | if(NOT "${last_char}" STREQUAL "\n") 11 | set(value "${value}\n") 12 | endif() 13 | message(STATUS "${var}:\n ${value}") 14 | endfunction() 15 | -------------------------------------------------------------------------------- /csrc/detect/end2end/include/common.hpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 1/24/23. 3 | // 4 | 5 | #ifndef DETECT_END2END_COMMON_HPP 6 | #define DETECT_END2END_COMMON_HPP 7 | #include "NvInfer.h" 8 | #include "filesystem.hpp" 9 | #include "opencv2/opencv.hpp" 10 | 11 | #define CHECK(call) \ 12 | do { \ 13 | const cudaError_t error_code = call; \ 14 | if (error_code != cudaSuccess) { \ 15 | printf("CUDA Error:\n"); \ 16 | printf(" File: %s\n", __FILE__); \ 17 | printf(" Line: %d\n", __LINE__); \ 18 | printf(" Error code: %d\n", error_code); \ 19 | printf(" Error text: %s\n", cudaGetErrorString(error_code)); \ 20 | exit(1); \ 21 | } \ 22 | } while (0) 23 | 24 | class Logger: public nvinfer1::ILogger { 25 | public: 26 | nvinfer1::ILogger::Severity reportableSeverity; 27 | 28 | explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO): 29 | reportableSeverity(severity) 30 | { 31 | } 32 | 33 | void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override 34 | { 35 | if (severity > reportableSeverity) { 36 | return; 37 | } 38 | switch (severity) { 39 | case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: 40 | std::cerr << "INTERNAL_ERROR: "; 41 | break; 42 | case nvinfer1::ILogger::Severity::kERROR: 43 | std::cerr << "ERROR: "; 44 | break; 45 | case nvinfer1::ILogger::Severity::kWARNING: 46 | std::cerr << "WARNING: "; 47 | break; 48 | case nvinfer1::ILogger::Severity::kINFO: 49 | std::cerr << "INFO: "; 50 | break; 51 | default: 52 | std::cerr << "VERBOSE: "; 53 | break; 54 | } 55 | std::cerr << msg << std::endl; 56 | } 57 | }; 58 | 59 | inline int get_size_by_dims(const nvinfer1::Dims& dims) 60 | { 61 | int size = 1; 62 | for (int i = 0; i < dims.nbDims; i++) { 63 | size *= dims.d[i]; 64 | } 65 | return size; 66 | } 67 | 68 | inline int type_to_size(const nvinfer1::DataType& dataType) 69 | { 70 | switch (dataType) { 71 | case nvinfer1::DataType::kFLOAT: 72 | return 4; 73 | case nvinfer1::DataType::kHALF: 74 | return 2; 75 | case nvinfer1::DataType::kINT32: 76 | return 4; 77 | case nvinfer1::DataType::kINT8: 78 | return 1; 79 | case nvinfer1::DataType::kBOOL: 80 | return 1; 81 | default: 82 | return 4; 83 | } 84 | } 85 | 86 | inline static float clamp(float val, float min, float max) 87 | { 88 | return val > min ? (val < max ? val : max) : min; 89 | } 90 | 91 | namespace det { 92 | struct Binding { 93 | size_t size = 1; 94 | size_t dsize = 1; 95 | nvinfer1::Dims dims; 96 | std::string name; 97 | }; 98 | 99 | struct Object { 100 | cv::Rect_ rect; 101 | int label = 0; 102 | float prob = 0.0; 103 | }; 104 | 105 | struct PreParam { 106 | float ratio = 1.0f; 107 | float dw = 0.0f; 108 | float dh = 0.0f; 109 | float height = 0; 110 | float width = 0; 111 | }; 112 | } // namespace det 113 | #endif // DETECT_END2END_COMMON_HPP 114 | -------------------------------------------------------------------------------- /csrc/detect/normal/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.12) 2 | 3 | set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86 89 90) 4 | set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc) 5 | 6 | project(yolov8 LANGUAGES CXX CUDA) 7 | 8 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -O3") 9 | set(CMAKE_CXX_STANDARD 14) 10 | set(CMAKE_BUILD_TYPE Release) 11 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 12 | 13 | list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake") 14 | include(Function) 15 | 16 | # CUDA 17 | find_package(CUDA REQUIRED) 18 | print_var(CUDA_LIBRARIES) 19 | print_var(CUDA_INCLUDE_DIRS) 20 | get_filename_component(CUDA_LIB_DIR ${CUDA_LIBRARIES} DIRECTORY) 21 | print_var(CUDA_LIB_DIR) 22 | 23 | # OpenCV 24 | find_package(OpenCV REQUIRED) 25 | print_var(OpenCV_LIBS) 26 | print_var(OpenCV_LIBRARIES) 27 | print_var(OpenCV_INCLUDE_DIRS) 28 | 29 | # TensorRT 30 | find_package(TensorRT REQUIRED) 31 | print_var(TensorRT_LIBRARIES) 32 | print_var(TensorRT_INCLUDE_DIRS) 33 | print_var(TensorRT_LIB_DIR) 34 | if (TensorRT_VERSION_MAJOR GREATER_EQUAL 10) 35 | message(STATUS "Build with -DTRT_10") 36 | add_definitions(-DTRT_10) 37 | endif () 38 | 39 | list(APPEND ALL_INCLUDE_DIRS 40 | ${CUDA_INCLUDE_DIRS} 41 | ${OpenCV_INCLUDE_DIRS} 42 | ${TensorRT_INCLUDE_DIRS} 43 | ${CMAKE_CURRENT_SOURCE_DIR}/include 44 | ) 45 | 46 | list(APPEND ALL_LIBS 47 | ${CUDA_LIBRARIES} 48 | ${OpenCV_LIBRARIES} 49 | ${TensorRT_LIBRARIES} 50 | ) 51 | 52 | list(APPEND ALL_LIB_DIRS 53 | ${CUDA_LIB_DIR} 54 | ${TensorRT_LIB_DIR} 55 | ) 56 | 57 | print_var(ALL_INCLUDE_DIRS) 58 | print_var(ALL_LIBS) 59 | print_var(ALL_LIB_DIRS) 60 | 61 | 62 | add_executable( 63 | ${PROJECT_NAME} 64 | ${CMAKE_CURRENT_SOURCE_DIR}/main.cpp 65 | ${CMAKE_CURRENT_SOURCE_DIR}/include/yolov8.hpp 66 | ${CMAKE_CURRENT_SOURCE_DIR}/include/common.hpp 67 | ) 68 | 69 | target_include_directories( 70 | ${PROJECT_NAME} 71 | PUBLIC 72 | ${ALL_INCLUDE_DIRS} 73 | ) 74 | 75 | target_link_directories( 76 | ${PROJECT_NAME} 77 | PUBLIC 78 | ${ALL_LIB_DIRS} 79 | ) 80 | 81 | target_link_libraries( 82 | ${PROJECT_NAME} 83 | PRIVATE 84 | ${ALL_LIBS} 85 | ) 86 | 87 | if (${OpenCV_VERSION} VERSION_GREATER_EQUAL 4.7.0) 88 | message(STATUS "Build with -DBATCHED_NMS") 89 | add_definitions(-DBATCHED_NMS) 90 | endif () 91 | -------------------------------------------------------------------------------- /csrc/detect/normal/cmake/FindTensorRT.cmake: -------------------------------------------------------------------------------- 1 | # This module defines the following variables: 2 | # 3 | # :: 4 | # 5 | # TensorRT_INCLUDE_DIRS 6 | # TensorRT_LIBRARIES 7 | # TensorRT_FOUND 8 | # 9 | # :: 10 | # 11 | # TensorRT_VERSION_STRING - version (x.y.z) 12 | # TensorRT_VERSION_MAJOR - major version (x) 13 | # TensorRT_VERSION_MINOR - minor version (y) 14 | # TensorRT_VERSION_PATCH - patch version (z) 15 | # 16 | # Hints 17 | # ^^^^^ 18 | # A user may set ``TensorRT_ROOT`` to an installation root to tell this module where to look. 19 | # 20 | set(_TensorRT_SEARCHES) 21 | 22 | if(TensorRT_ROOT) 23 | set(_TensorRT_SEARCH_ROOT PATHS ${TensorRT_ROOT} NO_DEFAULT_PATH) 24 | list(APPEND _TensorRT_SEARCHES _TensorRT_SEARCH_ROOT) 25 | endif() 26 | 27 | # appends some common paths 28 | set(_TensorRT_SEARCH_NORMAL 29 | PATHS "/usr" 30 | ) 31 | list(APPEND _TensorRT_SEARCHES _TensorRT_SEARCH_NORMAL) 32 | 33 | # Include dir 34 | foreach(search ${_TensorRT_SEARCHES}) 35 | find_path(TensorRT_INCLUDE_DIR NAMES NvInfer.h ${${search}} PATH_SUFFIXES include) 36 | endforeach() 37 | 38 | if(NOT TensorRT_LIBRARY) 39 | foreach(search ${_TensorRT_SEARCHES}) 40 | find_library(TensorRT_LIBRARY NAMES nvinfer ${${search}} PATH_SUFFIXES lib) 41 | if(NOT TensorRT_LIB_DIR) 42 | get_filename_component(TensorRT_LIB_DIR ${TensorRT_LIBRARY} DIRECTORY) 43 | endif () 44 | endforeach() 45 | endif() 46 | 47 | if(NOT TensorRT_nvinfer_plugin_LIBRARY) 48 | foreach(search ${_TensorRT_SEARCHES}) 49 | find_library(TensorRT_nvinfer_plugin_LIBRARY NAMES nvinfer_plugin ${${search}} PATH_SUFFIXES lib) 50 | endforeach() 51 | endif() 52 | 53 | mark_as_advanced(TensorRT_INCLUDE_DIR) 54 | 55 | if(TensorRT_INCLUDE_DIR AND EXISTS "${TensorRT_INCLUDE_DIR}/NvInfer.h") 56 | if(EXISTS "${TensorRT_INCLUDE_DIR}/NvInferVersion.h") 57 | set(_VersionSearchFile "${TensorRT_INCLUDE_DIR}/NvInferVersion.h") 58 | else () 59 | set(_VersionSearchFile "${TensorRT_INCLUDE_DIR}/NvInfer.h") 60 | endif () 61 | file(STRINGS "${_VersionSearchFile}" TensorRT_MAJOR REGEX "^#define NV_TENSORRT_MAJOR [0-9]+.*$") 62 | file(STRINGS "${_VersionSearchFile}" TensorRT_MINOR REGEX "^#define NV_TENSORRT_MINOR [0-9]+.*$") 63 | file(STRINGS "${_VersionSearchFile}" TensorRT_PATCH REGEX "^#define NV_TENSORRT_PATCH [0-9]+.*$") 64 | 65 | string(REGEX REPLACE "^#define NV_TENSORRT_MAJOR ([0-9]+).*$" "\\1" TensorRT_VERSION_MAJOR "${TensorRT_MAJOR}") 66 | string(REGEX REPLACE "^#define NV_TENSORRT_MINOR ([0-9]+).*$" "\\1" TensorRT_VERSION_MINOR "${TensorRT_MINOR}") 67 | string(REGEX REPLACE "^#define NV_TENSORRT_PATCH ([0-9]+).*$" "\\1" TensorRT_VERSION_PATCH "${TensorRT_PATCH}") 68 | set(TensorRT_VERSION_STRING "${TensorRT_VERSION_MAJOR}.${TensorRT_VERSION_MINOR}.${TensorRT_VERSION_PATCH}") 69 | endif() 70 | 71 | include(FindPackageHandleStandardArgs) 72 | FIND_PACKAGE_HANDLE_STANDARD_ARGS(TensorRT REQUIRED_VARS TensorRT_LIBRARY TensorRT_INCLUDE_DIR VERSION_VAR TensorRT_VERSION_STRING) 73 | 74 | if(TensorRT_FOUND) 75 | set(TensorRT_INCLUDE_DIRS ${TensorRT_INCLUDE_DIR}) 76 | 77 | if(NOT TensorRT_LIBRARIES) 78 | set(TensorRT_LIBRARIES ${TensorRT_LIBRARY}) 79 | if (TensorRT_nvinfer_plugin_LIBRARY) 80 | list(APPEND TensorRT_LIBRARIES ${TensorRT_nvinfer_plugin_LIBRARY}) 81 | endif() 82 | endif() 83 | 84 | if(NOT TARGET TensorRT::TensorRT) 85 | add_library(TensorRT INTERFACE IMPORTED) 86 | add_library(TensorRT::TensorRT ALIAS TensorRT) 87 | endif() 88 | 89 | if(NOT TARGET TensorRT::nvinfer) 90 | add_library(TensorRT::nvinfer SHARED IMPORTED) 91 | if (WIN32) 92 | foreach(search ${_TensorRT_SEARCHES}) 93 | find_file(TensorRT_LIBRARY_DLL 94 | NAMES nvinfer.dll 95 | PATHS ${${search}} 96 | PATH_SUFFIXES bin 97 | ) 98 | endforeach() 99 | 100 | set_target_properties(TensorRT::nvinfer PROPERTIES 101 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 102 | IMPORTED_LOCATION "${TensorRT_LIBRARY_DLL}" 103 | IMPORTED_IMPLIB "${TensorRT_LIBRARY}" 104 | ) 105 | else() 106 | set_target_properties(TensorRT::nvinfer PROPERTIES 107 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 108 | IMPORTED_LOCATION "${TensorRT_LIBRARY}" 109 | ) 110 | endif() 111 | target_link_libraries(TensorRT INTERFACE TensorRT::nvinfer) 112 | endif() 113 | 114 | if(NOT TARGET TensorRT::nvinfer_plugin AND TensorRT_nvinfer_plugin_LIBRARY) 115 | add_library(TensorRT::nvinfer_plugin SHARED IMPORTED) 116 | if (WIN32) 117 | foreach(search ${_TensorRT_SEARCHES}) 118 | find_file(TensorRT_nvinfer_plugin_LIBRARY_DLL 119 | NAMES nvinfer_plugin.dll 120 | PATHS ${${search}} 121 | PATH_SUFFIXES bin 122 | ) 123 | endforeach() 124 | 125 | set_target_properties(TensorRT::nvinfer_plugin PROPERTIES 126 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 127 | IMPORTED_LOCATION "${TensorRT_nvinfer_plugin_LIBRARY_DLL}" 128 | IMPORTED_IMPLIB "${TensorRT_nvinfer_plugin_LIBRARY}" 129 | ) 130 | else() 131 | set_target_properties(TensorRT::nvinfer_plugin PROPERTIES 132 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 133 | IMPORTED_LOCATION "${TensorRT_nvinfer_plugin_LIBRARY}" 134 | ) 135 | endif() 136 | target_link_libraries(TensorRT INTERFACE TensorRT::nvinfer_plugin) 137 | endif() 138 | endif() 139 | -------------------------------------------------------------------------------- /csrc/detect/normal/cmake/Function.cmake: -------------------------------------------------------------------------------- 1 | 2 | function(print_var var) 3 | set(value "${${var}}") 4 | string(LENGTH "${value}" value_length) 5 | if(value_length GREATER 0) 6 | math(EXPR last_index "${value_length} - 1") 7 | string(SUBSTRING "${value}" ${last_index} ${last_index} last_char) 8 | endif() 9 | 10 | if(NOT "${last_char}" STREQUAL "\n") 11 | set(value "${value}\n") 12 | endif() 13 | message(STATUS "${var}:\n ${value}") 14 | endfunction() 15 | -------------------------------------------------------------------------------- /csrc/detect/normal/include/common.hpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 1/24/23. 3 | // 4 | 5 | #ifndef DETECT_NORMAL_COMMON_HPP 6 | #define DETECT_NORMAL_COMMON_HPP 7 | #include "NvInfer.h" 8 | #include "filesystem.hpp" 9 | #include "opencv2/opencv.hpp" 10 | 11 | #define CHECK(call) \ 12 | do { \ 13 | const cudaError_t error_code = call; \ 14 | if (error_code != cudaSuccess) { \ 15 | printf("CUDA Error:\n"); \ 16 | printf(" File: %s\n", __FILE__); \ 17 | printf(" Line: %d\n", __LINE__); \ 18 | printf(" Error code: %d\n", error_code); \ 19 | printf(" Error text: %s\n", cudaGetErrorString(error_code)); \ 20 | exit(1); \ 21 | } \ 22 | } while (0) 23 | 24 | class Logger: public nvinfer1::ILogger { 25 | public: 26 | nvinfer1::ILogger::Severity reportableSeverity; 27 | 28 | explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO): 29 | reportableSeverity(severity) 30 | { 31 | } 32 | 33 | void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override 34 | { 35 | if (severity > reportableSeverity) { 36 | return; 37 | } 38 | switch (severity) { 39 | case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: 40 | std::cerr << "INTERNAL_ERROR: "; 41 | break; 42 | case nvinfer1::ILogger::Severity::kERROR: 43 | std::cerr << "ERROR: "; 44 | break; 45 | case nvinfer1::ILogger::Severity::kWARNING: 46 | std::cerr << "WARNING: "; 47 | break; 48 | case nvinfer1::ILogger::Severity::kINFO: 49 | std::cerr << "INFO: "; 50 | break; 51 | default: 52 | std::cerr << "VERBOSE: "; 53 | break; 54 | } 55 | std::cerr << msg << std::endl; 56 | } 57 | }; 58 | 59 | inline int get_size_by_dims(const nvinfer1::Dims& dims) 60 | { 61 | int size = 1; 62 | for (int i = 0; i < dims.nbDims; i++) { 63 | size *= dims.d[i]; 64 | } 65 | return size; 66 | } 67 | 68 | inline int type_to_size(const nvinfer1::DataType& dataType) 69 | { 70 | switch (dataType) { 71 | case nvinfer1::DataType::kFLOAT: 72 | return 4; 73 | case nvinfer1::DataType::kHALF: 74 | return 2; 75 | case nvinfer1::DataType::kINT32: 76 | return 4; 77 | case nvinfer1::DataType::kINT8: 78 | return 1; 79 | case nvinfer1::DataType::kBOOL: 80 | return 1; 81 | default: 82 | return 4; 83 | } 84 | } 85 | 86 | inline static float clamp(float val, float min, float max) 87 | { 88 | return val > min ? (val < max ? val : max) : min; 89 | } 90 | 91 | namespace det { 92 | struct Binding { 93 | size_t size = 1; 94 | size_t dsize = 1; 95 | nvinfer1::Dims dims; 96 | std::string name; 97 | }; 98 | 99 | struct Object { 100 | cv::Rect_ rect; 101 | int label = 0; 102 | float prob = 0.0; 103 | }; 104 | 105 | struct PreParam { 106 | float ratio = 1.0f; 107 | float dw = 0.0f; 108 | float dh = 0.0f; 109 | float height = 0; 110 | float width = 0; 111 | }; 112 | } // namespace det 113 | #endif // DETECT_NORMAL_COMMON_HPP 114 | -------------------------------------------------------------------------------- /csrc/jetson/detect/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.1) 2 | 3 | set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86) 4 | set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc) 5 | 6 | project(yolov8 LANGUAGES CXX CUDA) 7 | 8 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -O3") 9 | set(CMAKE_CXX_STANDARD 14) 10 | set(CMAKE_BUILD_TYPE Release) 11 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 12 | 13 | # CUDA 14 | find_package(CUDA REQUIRED) 15 | message(STATUS "CUDA Libs: \n${CUDA_LIBRARIES}\n") 16 | get_filename_component(CUDA_LIB_DIR ${CUDA_LIBRARIES} DIRECTORY) 17 | message(STATUS "CUDA Headers: \n${CUDA_INCLUDE_DIRS}\n") 18 | 19 | # OpenCV 20 | find_package(OpenCV REQUIRED) 21 | message(STATUS "OpenCV Libs: \n${OpenCV_LIBS}\n") 22 | message(STATUS "OpenCV Libraries: \n${OpenCV_LIBRARIES}\n") 23 | message(STATUS "OpenCV Headers: \n${OpenCV_INCLUDE_DIRS}\n") 24 | 25 | # TensorRT 26 | set(TensorRT_INCLUDE_DIRS /usr/include/aarch64-linux-gnu) 27 | set(TensorRT_LIBRARIES /usr/lib/aarch64-linux-gnu) 28 | 29 | 30 | message(STATUS "TensorRT Libs: \n${TensorRT_LIBRARIES}\n") 31 | message(STATUS "TensorRT Headers: \n${TensorRT_INCLUDE_DIRS}\n") 32 | 33 | list(APPEND INCLUDE_DIRS 34 | ${CUDA_INCLUDE_DIRS} 35 | ${OpenCV_INCLUDE_DIRS} 36 | ${TensorRT_INCLUDE_DIRS} 37 | include 38 | ) 39 | 40 | list(APPEND ALL_LIBS 41 | ${CUDA_LIBRARIES} 42 | ${CUDA_LIB_DIR} 43 | ${OpenCV_LIBRARIES} 44 | ${TensorRT_LIBRARIES} 45 | ) 46 | 47 | include_directories(${INCLUDE_DIRS}) 48 | 49 | add_executable(${PROJECT_NAME} 50 | main.cpp 51 | include/yolov8.hpp 52 | include/common.hpp 53 | ) 54 | 55 | link_directories(${ALL_LIBS}) 56 | target_link_libraries(${PROJECT_NAME} PRIVATE nvinfer nvinfer_plugin ${CUDA_LIBRARIES} ${OpenCV_LIBS}) 57 | -------------------------------------------------------------------------------- /csrc/jetson/detect/include/common.hpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 3/16/23. 3 | // 4 | 5 | #ifndef JETSON_DETECT_COMMON_HPP 6 | #define JETSON_DETECT_COMMON_HPP 7 | #include "NvInfer.h" 8 | #include "filesystem.hpp" 9 | #include "opencv2/opencv.hpp" 10 | 11 | #define CHECK(call) \ 12 | do { \ 13 | const cudaError_t error_code = call; \ 14 | if (error_code != cudaSuccess) { \ 15 | printf("CUDA Error:\n"); \ 16 | printf(" File: %s\n", __FILE__); \ 17 | printf(" Line: %d\n", __LINE__); \ 18 | printf(" Error code: %d\n", error_code); \ 19 | printf(" Error text: %s\n", cudaGetErrorString(error_code)); \ 20 | exit(1); \ 21 | } \ 22 | } while (0) 23 | 24 | class Logger: public nvinfer1::ILogger { 25 | public: 26 | nvinfer1::ILogger::Severity reportableSeverity; 27 | 28 | explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO): 29 | reportableSeverity(severity) 30 | { 31 | } 32 | 33 | void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override 34 | { 35 | if (severity > reportableSeverity) { 36 | return; 37 | } 38 | switch (severity) { 39 | case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: 40 | std::cerr << "INTERNAL_ERROR: "; 41 | break; 42 | case nvinfer1::ILogger::Severity::kERROR: 43 | std::cerr << "ERROR: "; 44 | break; 45 | case nvinfer1::ILogger::Severity::kWARNING: 46 | std::cerr << "WARNING: "; 47 | break; 48 | case nvinfer1::ILogger::Severity::kINFO: 49 | std::cerr << "INFO: "; 50 | break; 51 | default: 52 | std::cerr << "VERBOSE: "; 53 | break; 54 | } 55 | std::cerr << msg << std::endl; 56 | } 57 | }; 58 | 59 | inline int get_size_by_dims(const nvinfer1::Dims& dims) 60 | { 61 | int size = 1; 62 | for (int i = 0; i < dims.nbDims; i++) { 63 | size *= dims.d[i]; 64 | } 65 | return size; 66 | } 67 | 68 | inline int type_to_size(const nvinfer1::DataType& dataType) 69 | { 70 | switch (dataType) { 71 | case nvinfer1::DataType::kFLOAT: 72 | return 4; 73 | case nvinfer1::DataType::kHALF: 74 | return 2; 75 | case nvinfer1::DataType::kINT32: 76 | return 4; 77 | case nvinfer1::DataType::kINT8: 78 | return 1; 79 | case nvinfer1::DataType::kBOOL: 80 | return 1; 81 | default: 82 | return 4; 83 | } 84 | } 85 | 86 | inline static float clamp(float val, float min, float max) 87 | { 88 | return val > min ? (val < max ? val : max) : min; 89 | } 90 | 91 | namespace det { 92 | struct Binding { 93 | size_t size = 1; 94 | size_t dsize = 1; 95 | nvinfer1::Dims dims; 96 | std::string name; 97 | }; 98 | 99 | struct Object { 100 | cv::Rect_ rect; 101 | int label = 0; 102 | float prob = 0.0; 103 | }; 104 | 105 | struct PreParam { 106 | float ratio = 1.0f; 107 | float dw = 0.0f; 108 | float dh = 0.0f; 109 | float height = 0; 110 | float width = 0; 111 | }; 112 | } // namespace det 113 | #endif // JETSON_DETECT_COMMON_HPP 114 | -------------------------------------------------------------------------------- /csrc/jetson/detect/main.cpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 3/16/23. 3 | // 4 | #include "opencv2/opencv.hpp" 5 | #include "yolov8.hpp" 6 | #include 7 | 8 | namespace fs = ghc::filesystem; 9 | 10 | const std::vector CLASS_NAMES = { 11 | "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", 12 | "truck", "boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", 13 | "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", 14 | "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", 15 | "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat", 16 | "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup", 17 | "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", 18 | "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", 19 | "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", 20 | "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven", 21 | "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", 22 | "teddy bear", "hair drier", "toothbrush"}; 23 | 24 | const std::vector> COLORS = { 25 | {0, 114, 189}, {217, 83, 25}, {237, 177, 32}, {126, 47, 142}, {119, 172, 48}, {77, 190, 238}, 26 | {162, 20, 47}, {76, 76, 76}, {153, 153, 153}, {255, 0, 0}, {255, 128, 0}, {191, 191, 0}, 27 | {0, 255, 0}, {0, 0, 255}, {170, 0, 255}, {85, 85, 0}, {85, 170, 0}, {85, 255, 0}, 28 | {170, 85, 0}, {170, 170, 0}, {170, 255, 0}, {255, 85, 0}, {255, 170, 0}, {255, 255, 0}, 29 | {0, 85, 128}, {0, 170, 128}, {0, 255, 128}, {85, 0, 128}, {85, 85, 128}, {85, 170, 128}, 30 | {85, 255, 128}, {170, 0, 128}, {170, 85, 128}, {170, 170, 128}, {170, 255, 128}, {255, 0, 128}, 31 | {255, 85, 128}, {255, 170, 128}, {255, 255, 128}, {0, 85, 255}, {0, 170, 255}, {0, 255, 255}, 32 | {85, 0, 255}, {85, 85, 255}, {85, 170, 255}, {85, 255, 255}, {170, 0, 255}, {170, 85, 255}, 33 | {170, 170, 255}, {170, 255, 255}, {255, 0, 255}, {255, 85, 255}, {255, 170, 255}, {85, 0, 0}, 34 | {128, 0, 0}, {170, 0, 0}, {212, 0, 0}, {255, 0, 0}, {0, 43, 0}, {0, 85, 0}, 35 | {0, 128, 0}, {0, 170, 0}, {0, 212, 0}, {0, 255, 0}, {0, 0, 43}, {0, 0, 85}, 36 | {0, 0, 128}, {0, 0, 170}, {0, 0, 212}, {0, 0, 255}, {0, 0, 0}, {36, 36, 36}, 37 | {73, 73, 73}, {109, 109, 109}, {146, 146, 146}, {182, 182, 182}, {219, 219, 219}, {0, 114, 189}, 38 | {80, 183, 189}, {128, 128, 0}}; 39 | 40 | int main(int argc, char** argv) 41 | { 42 | if (argc != 3) { 43 | fprintf(stderr, "Usage: %s [engine_path] [image_path/image_dir/video_path]\n", argv[0]); 44 | return -1; 45 | } 46 | 47 | // cuda:0 48 | cudaSetDevice(0); 49 | 50 | const std::string engine_file_path{argv[1]}; 51 | const fs::path path{argv[2]}; 52 | 53 | std::vector imagePathList; 54 | bool isVideo{false}; 55 | 56 | auto yolov8 = new YOLOv8(engine_file_path); 57 | yolov8->make_pipe(true); 58 | 59 | if (fs::exists(path)) { 60 | std::string suffix = path.extension(); 61 | if (suffix == ".jpg" || suffix == ".jpeg" || suffix == ".png") { 62 | imagePathList.push_back(path); 63 | } 64 | else if (suffix == ".mp4" || suffix == ".avi" || suffix == ".m4v" || suffix == ".mpeg" || suffix == ".mov" 65 | || suffix == ".mkv") { 66 | isVideo = true; 67 | } 68 | else { 69 | printf("suffix %s is wrong !!!\n", suffix.c_str()); 70 | std::abort(); 71 | } 72 | } 73 | else if (fs::is_directory(path)) { 74 | cv::glob(path.string() + "/*.jpg", imagePathList); 75 | } 76 | 77 | cv::Mat res, image; 78 | cv::Size size = cv::Size{640, 640}; 79 | std::vector objs; 80 | 81 | cv::namedWindow("result", cv::WINDOW_AUTOSIZE); 82 | 83 | if (isVideo) { 84 | cv::VideoCapture cap(path); 85 | 86 | if (!cap.isOpened()) { 87 | printf("can not open %s\n", path.c_str()); 88 | return -1; 89 | } 90 | while (cap.read(image)) { 91 | objs.clear(); 92 | yolov8->copy_from_Mat(image, size); 93 | auto start = std::chrono::system_clock::now(); 94 | yolov8->infer(); 95 | auto end = std::chrono::system_clock::now(); 96 | yolov8->postprocess(objs); 97 | yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS); 98 | auto tc = (double)std::chrono::duration_cast(end - start).count() / 1000.; 99 | printf("cost %2.4lf ms\n", tc); 100 | cv::imshow("result", res); 101 | if (cv::waitKey(10) == 'q') { 102 | break; 103 | } 104 | } 105 | } 106 | else { 107 | for (auto& p : imagePathList) { 108 | objs.clear(); 109 | image = cv::imread(p); 110 | yolov8->copy_from_Mat(image, size); 111 | auto start = std::chrono::system_clock::now(); 112 | yolov8->infer(); 113 | auto end = std::chrono::system_clock::now(); 114 | yolov8->postprocess(objs); 115 | yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS); 116 | auto tc = (double)std::chrono::duration_cast(end - start).count() / 1000.; 117 | printf("cost %2.4lf ms\n", tc); 118 | cv::imshow("result", res); 119 | cv::waitKey(0); 120 | } 121 | } 122 | cv::destroyAllWindows(); 123 | delete yolov8; 124 | return 0; 125 | } 126 | -------------------------------------------------------------------------------- /csrc/jetson/pose/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.1) 2 | 3 | set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86) 4 | set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc) 5 | 6 | project(yolov8-pose LANGUAGES CXX CUDA) 7 | 8 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -O3") 9 | set(CMAKE_CXX_STANDARD 14) 10 | set(CMAKE_BUILD_TYPE Release) 11 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 12 | 13 | # CUDA 14 | find_package(CUDA REQUIRED) 15 | message(STATUS "CUDA Libs: \n${CUDA_LIBRARIES}\n") 16 | get_filename_component(CUDA_LIB_DIR ${CUDA_LIBRARIES} DIRECTORY) 17 | message(STATUS "CUDA Headers: \n${CUDA_INCLUDE_DIRS}\n") 18 | 19 | # OpenCV 20 | find_package(OpenCV REQUIRED) 21 | message(STATUS "OpenCV Libs: \n${OpenCV_LIBS}\n") 22 | message(STATUS "OpenCV Libraries: \n${OpenCV_LIBRARIES}\n") 23 | message(STATUS "OpenCV Headers: \n${OpenCV_INCLUDE_DIRS}\n") 24 | 25 | # TensorRT 26 | set(TensorRT_INCLUDE_DIRS /usr/include/aarch64-linux-gnu) 27 | set(TensorRT_LIBRARIES /usr/lib/aarch64-linux-gnu) 28 | 29 | 30 | message(STATUS "TensorRT Libs: \n${TensorRT_LIBRARIES}\n") 31 | message(STATUS "TensorRT Headers: \n${TensorRT_INCLUDE_DIRS}\n") 32 | 33 | list(APPEND INCLUDE_DIRS 34 | ${CUDA_INCLUDE_DIRS} 35 | ${OpenCV_INCLUDE_DIRS} 36 | ${TensorRT_INCLUDE_DIRS} 37 | include 38 | ) 39 | 40 | list(APPEND ALL_LIBS 41 | ${CUDA_LIBRARIES} 42 | ${CUDA_LIB_DIR} 43 | ${OpenCV_LIBRARIES} 44 | ${TensorRT_LIBRARIES} 45 | ) 46 | 47 | include_directories(${INCLUDE_DIRS}) 48 | 49 | add_executable(${PROJECT_NAME} 50 | main.cpp 51 | include/yolov8-pose.hpp 52 | include/common.hpp 53 | ) 54 | 55 | link_directories(${ALL_LIBS}) 56 | target_link_libraries(${PROJECT_NAME} PRIVATE nvinfer nvinfer_plugin ${CUDA_LIBRARIES} ${OpenCV_LIBS}) 57 | 58 | 59 | if(${OpenCV_VERSION} VERSION_GREATER_EQUAL 4.7.0) 60 | message(STATUS "Build with -DBATCHED_NMS") 61 | add_definitions(-DBATCHED_NMS) 62 | endif() 63 | -------------------------------------------------------------------------------- /csrc/jetson/pose/include/common.hpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 5/15/23. 3 | // 4 | 5 | #ifndef JETSON_POSE_COMMON_HPP 6 | #define JETSON_POSE_COMMON_HPP 7 | #include "NvInfer.h" 8 | #include "filesystem.hpp" 9 | #include "opencv2/opencv.hpp" 10 | 11 | #define CHECK(call) \ 12 | do { \ 13 | const cudaError_t error_code = call; \ 14 | if (error_code != cudaSuccess) { \ 15 | printf("CUDA Error:\n"); \ 16 | printf(" File: %s\n", __FILE__); \ 17 | printf(" Line: %d\n", __LINE__); \ 18 | printf(" Error code: %d\n", error_code); \ 19 | printf(" Error text: %s\n", cudaGetErrorString(error_code)); \ 20 | exit(1); \ 21 | } \ 22 | } while (0) 23 | 24 | class Logger: public nvinfer1::ILogger { 25 | public: 26 | nvinfer1::ILogger::Severity reportableSeverity; 27 | 28 | explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO): 29 | reportableSeverity(severity) 30 | { 31 | } 32 | 33 | void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override 34 | { 35 | if (severity > reportableSeverity) { 36 | return; 37 | } 38 | switch (severity) { 39 | case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: 40 | std::cerr << "INTERNAL_ERROR: "; 41 | break; 42 | case nvinfer1::ILogger::Severity::kERROR: 43 | std::cerr << "ERROR: "; 44 | break; 45 | case nvinfer1::ILogger::Severity::kWARNING: 46 | std::cerr << "WARNING: "; 47 | break; 48 | case nvinfer1::ILogger::Severity::kINFO: 49 | std::cerr << "INFO: "; 50 | break; 51 | default: 52 | std::cerr << "VERBOSE: "; 53 | break; 54 | } 55 | std::cerr << msg << std::endl; 56 | } 57 | }; 58 | 59 | inline int get_size_by_dims(const nvinfer1::Dims& dims) 60 | { 61 | int size = 1; 62 | for (int i = 0; i < dims.nbDims; i++) { 63 | size *= dims.d[i]; 64 | } 65 | return size; 66 | } 67 | 68 | inline int type_to_size(const nvinfer1::DataType& dataType) 69 | { 70 | switch (dataType) { 71 | case nvinfer1::DataType::kFLOAT: 72 | return 4; 73 | case nvinfer1::DataType::kHALF: 74 | return 2; 75 | case nvinfer1::DataType::kINT32: 76 | return 4; 77 | case nvinfer1::DataType::kINT8: 78 | return 1; 79 | case nvinfer1::DataType::kBOOL: 80 | return 1; 81 | default: 82 | return 4; 83 | } 84 | } 85 | 86 | inline static float clamp(float val, float min, float max) 87 | { 88 | return val > min ? (val < max ? val : max) : min; 89 | } 90 | 91 | namespace pose { 92 | struct Binding { 93 | size_t size = 1; 94 | size_t dsize = 1; 95 | nvinfer1::Dims dims; 96 | std::string name; 97 | }; 98 | 99 | struct Object { 100 | cv::Rect_ rect; 101 | int label = 0; 102 | float prob = 0.0; 103 | std::vector kps; 104 | }; 105 | 106 | struct PreParam { 107 | float ratio = 1.0f; 108 | float dw = 0.0f; 109 | float dh = 0.0f; 110 | float height = 0; 111 | float width = 0; 112 | }; 113 | } // namespace pose 114 | #endif // JETSON_POSE_COMMON_HPP 115 | -------------------------------------------------------------------------------- /csrc/jetson/segment/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.1) 2 | 3 | set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86) 4 | set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc) 5 | 6 | project(yolov8-seg LANGUAGES CXX CUDA) 7 | 8 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -O3") 9 | set(CMAKE_CXX_STANDARD 14) 10 | set(CMAKE_BUILD_TYPE Release) 11 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 12 | 13 | # CUDA 14 | find_package(CUDA REQUIRED) 15 | message(STATUS "CUDA Libs: \n${CUDA_LIBRARIES}\n") 16 | get_filename_component(CUDA_LIB_DIR ${CUDA_LIBRARIES} DIRECTORY) 17 | message(STATUS "CUDA Headers: \n${CUDA_INCLUDE_DIRS}\n") 18 | 19 | # OpenCV 20 | find_package(OpenCV REQUIRED) 21 | message(STATUS "OpenCV Libs: \n${OpenCV_LIBS}\n") 22 | message(STATUS "OpenCV Libraries: \n${OpenCV_LIBRARIES}\n") 23 | message(STATUS "OpenCV Headers: \n${OpenCV_INCLUDE_DIRS}\n") 24 | 25 | # TensorRT 26 | set(TensorRT_INCLUDE_DIRS /usr/include/aarch64-linux-gnu) 27 | set(TensorRT_LIBRARIES /usr/lib/aarch64-linux-gnu) 28 | 29 | 30 | message(STATUS "TensorRT Libs: \n${TensorRT_LIBRARIES}\n") 31 | message(STATUS "TensorRT Headers: \n${TensorRT_INCLUDE_DIRS}\n") 32 | 33 | list(APPEND INCLUDE_DIRS 34 | ${CUDA_INCLUDE_DIRS} 35 | ${OpenCV_INCLUDE_DIRS} 36 | ${TensorRT_INCLUDE_DIRS} 37 | include 38 | ) 39 | 40 | list(APPEND ALL_LIBS 41 | ${CUDA_LIBRARIES} 42 | ${CUDA_LIB_DIR} 43 | ${OpenCV_LIBRARIES} 44 | ${TensorRT_LIBRARIES} 45 | ) 46 | 47 | include_directories(${INCLUDE_DIRS}) 48 | 49 | add_executable(${PROJECT_NAME} 50 | main.cpp 51 | include/yolov8-seg.hpp 52 | include/common.hpp 53 | ) 54 | 55 | link_directories(${ALL_LIBS}) 56 | target_link_libraries(${PROJECT_NAME} PRIVATE nvinfer nvinfer_plugin ${CUDA_LIBRARIES} ${OpenCV_LIBS}) 57 | 58 | 59 | if(${OpenCV_VERSION} VERSION_GREATER_EQUAL 4.7.0) 60 | message(STATUS "Build with -DBATCHED_NMS") 61 | add_definitions(-DBATCHED_NMS) 62 | endif() 63 | -------------------------------------------------------------------------------- /csrc/jetson/segment/include/common.hpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 3/16/23. 3 | // 4 | 5 | #ifndef JETSON_SEGMENT_COMMON_HPP 6 | #define JETSON_SEGMENT_COMMON_HPP 7 | #include "NvInfer.h" 8 | #include "filesystem.hpp" 9 | #include "opencv2/opencv.hpp" 10 | 11 | #define CHECK(call) \ 12 | do { \ 13 | const cudaError_t error_code = call; \ 14 | if (error_code != cudaSuccess) { \ 15 | printf("CUDA Error:\n"); \ 16 | printf(" File: %s\n", __FILE__); \ 17 | printf(" Line: %d\n", __LINE__); \ 18 | printf(" Error code: %d\n", error_code); \ 19 | printf(" Error text: %s\n", cudaGetErrorString(error_code)); \ 20 | exit(1); \ 21 | } \ 22 | } while (0) 23 | 24 | class Logger: public nvinfer1::ILogger { 25 | public: 26 | nvinfer1::ILogger::Severity reportableSeverity; 27 | 28 | explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO): 29 | reportableSeverity(severity) 30 | { 31 | } 32 | 33 | void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override 34 | { 35 | if (severity > reportableSeverity) { 36 | return; 37 | } 38 | switch (severity) { 39 | case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: 40 | std::cerr << "INTERNAL_ERROR: "; 41 | break; 42 | case nvinfer1::ILogger::Severity::kERROR: 43 | std::cerr << "ERROR: "; 44 | break; 45 | case nvinfer1::ILogger::Severity::kWARNING: 46 | std::cerr << "WARNING: "; 47 | break; 48 | case nvinfer1::ILogger::Severity::kINFO: 49 | std::cerr << "INFO: "; 50 | break; 51 | default: 52 | std::cerr << "VERBOSE: "; 53 | break; 54 | } 55 | std::cerr << msg << std::endl; 56 | } 57 | }; 58 | 59 | inline int get_size_by_dims(const nvinfer1::Dims& dims) 60 | { 61 | int size = 1; 62 | for (int i = 0; i < dims.nbDims; i++) { 63 | size *= dims.d[i]; 64 | } 65 | return size; 66 | } 67 | 68 | inline int type_to_size(const nvinfer1::DataType& dataType) 69 | { 70 | switch (dataType) { 71 | case nvinfer1::DataType::kFLOAT: 72 | return 4; 73 | case nvinfer1::DataType::kHALF: 74 | return 2; 75 | case nvinfer1::DataType::kINT32: 76 | return 4; 77 | case nvinfer1::DataType::kINT8: 78 | return 1; 79 | case nvinfer1::DataType::kBOOL: 80 | return 1; 81 | default: 82 | return 4; 83 | } 84 | } 85 | 86 | inline static float clamp(float val, float min, float max) 87 | { 88 | return val > min ? (val < max ? val : max) : min; 89 | } 90 | 91 | namespace seg { 92 | struct Binding { 93 | size_t size = 1; 94 | size_t dsize = 1; 95 | nvinfer1::Dims dims; 96 | std::string name; 97 | }; 98 | 99 | struct Object { 100 | cv::Rect_ rect; 101 | int label = 0; 102 | float prob = 0.0; 103 | cv::Mat boxMask; 104 | }; 105 | 106 | struct PreParam { 107 | float ratio = 1.0f; 108 | float dw = 0.0f; 109 | float dh = 0.0f; 110 | float height = 0; 111 | float width = 0; 112 | }; 113 | } // namespace seg 114 | #endif // JETSON_SEGMENT_COMMON_HPP 115 | -------------------------------------------------------------------------------- /csrc/obb/normal/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.1) 2 | 3 | set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86 89 90) 4 | set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc) 5 | 6 | project(yolov8-obb LANGUAGES CXX CUDA) 7 | 8 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -O3") 9 | set(CMAKE_CXX_STANDARD 14) 10 | set(CMAKE_BUILD_TYPE Release) 11 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 12 | 13 | list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake") 14 | include(Function) 15 | 16 | # CUDA 17 | find_package(CUDA REQUIRED) 18 | print_var(CUDA_LIBRARIES) 19 | print_var(CUDA_INCLUDE_DIRS) 20 | get_filename_component(CUDA_LIB_DIR ${CUDA_LIBRARIES} DIRECTORY) 21 | print_var(CUDA_LIB_DIR) 22 | 23 | # OpenCV 24 | find_package(OpenCV REQUIRED) 25 | print_var(OpenCV_LIBS) 26 | print_var(OpenCV_LIBRARIES) 27 | print_var(OpenCV_INCLUDE_DIRS) 28 | 29 | # TensorRT 30 | find_package(TensorRT REQUIRED) 31 | print_var(TensorRT_LIBRARIES) 32 | print_var(TensorRT_INCLUDE_DIRS) 33 | print_var(TensorRT_LIB_DIR) 34 | if (TensorRT_VERSION_MAJOR GREATER_EQUAL 10) 35 | message(STATUS "Build with -DTRT_10") 36 | add_definitions(-DTRT_10) 37 | endif () 38 | 39 | list(APPEND ALL_INCLUDE_DIRS 40 | ${CUDA_INCLUDE_DIRS} 41 | ${OpenCV_INCLUDE_DIRS} 42 | ${TensorRT_INCLUDE_DIRS} 43 | ${CMAKE_CURRENT_SOURCE_DIR}/include 44 | ) 45 | 46 | list(APPEND ALL_LIBS 47 | ${CUDA_LIBRARIES} 48 | ${OpenCV_LIBRARIES} 49 | ${TensorRT_LIBRARIES} 50 | ) 51 | 52 | list(APPEND ALL_LIB_DIRS 53 | ${CUDA_LIB_DIR} 54 | ${TensorRT_LIB_DIR} 55 | ) 56 | 57 | print_var(ALL_INCLUDE_DIRS) 58 | print_var(ALL_LIBS) 59 | print_var(ALL_LIB_DIRS) 60 | 61 | 62 | add_executable(${PROJECT_NAME} 63 | ${CMAKE_CURRENT_SOURCE_DIR}/main.cpp 64 | ${CMAKE_CURRENT_SOURCE_DIR}/include/yolov8-obb.hpp 65 | ${CMAKE_CURRENT_SOURCE_DIR}/include/common.hpp 66 | ) 67 | 68 | target_include_directories( 69 | ${PROJECT_NAME} 70 | PUBLIC 71 | ${ALL_INCLUDE_DIRS} 72 | ) 73 | 74 | target_link_directories( 75 | ${PROJECT_NAME} 76 | PUBLIC 77 | ${ALL_LIB_DIRS} 78 | ) 79 | 80 | target_link_libraries( 81 | ${PROJECT_NAME} 82 | PRIVATE 83 | ${ALL_LIBS} 84 | ) 85 | -------------------------------------------------------------------------------- /csrc/obb/normal/cmake/FindTensorRT.cmake: -------------------------------------------------------------------------------- 1 | # This module defines the following variables: 2 | # 3 | # :: 4 | # 5 | # TensorRT_INCLUDE_DIRS 6 | # TensorRT_LIBRARIES 7 | # TensorRT_FOUND 8 | # 9 | # :: 10 | # 11 | # TensorRT_VERSION_STRING - version (x.y.z) 12 | # TensorRT_VERSION_MAJOR - major version (x) 13 | # TensorRT_VERSION_MINOR - minor version (y) 14 | # TensorRT_VERSION_PATCH - patch version (z) 15 | # 16 | # Hints 17 | # ^^^^^ 18 | # A user may set ``TensorRT_ROOT`` to an installation root to tell this module where to look. 19 | # 20 | set(_TensorRT_SEARCHES) 21 | 22 | if(TensorRT_ROOT) 23 | set(_TensorRT_SEARCH_ROOT PATHS ${TensorRT_ROOT} NO_DEFAULT_PATH) 24 | list(APPEND _TensorRT_SEARCHES _TensorRT_SEARCH_ROOT) 25 | endif() 26 | 27 | # appends some common paths 28 | set(_TensorRT_SEARCH_NORMAL 29 | PATHS "/usr" 30 | ) 31 | list(APPEND _TensorRT_SEARCHES _TensorRT_SEARCH_NORMAL) 32 | 33 | # Include dir 34 | foreach(search ${_TensorRT_SEARCHES}) 35 | find_path(TensorRT_INCLUDE_DIR NAMES NvInfer.h ${${search}} PATH_SUFFIXES include) 36 | endforeach() 37 | 38 | if(NOT TensorRT_LIBRARY) 39 | foreach(search ${_TensorRT_SEARCHES}) 40 | find_library(TensorRT_LIBRARY NAMES nvinfer ${${search}} PATH_SUFFIXES lib) 41 | if(NOT TensorRT_LIB_DIR) 42 | get_filename_component(TensorRT_LIB_DIR ${TensorRT_LIBRARY} DIRECTORY) 43 | endif () 44 | endforeach() 45 | endif() 46 | 47 | if(NOT TensorRT_nvinfer_plugin_LIBRARY) 48 | foreach(search ${_TensorRT_SEARCHES}) 49 | find_library(TensorRT_nvinfer_plugin_LIBRARY NAMES nvinfer_plugin ${${search}} PATH_SUFFIXES lib) 50 | endforeach() 51 | endif() 52 | 53 | mark_as_advanced(TensorRT_INCLUDE_DIR) 54 | 55 | if(TensorRT_INCLUDE_DIR AND EXISTS "${TensorRT_INCLUDE_DIR}/NvInfer.h") 56 | if(EXISTS "${TensorRT_INCLUDE_DIR}/NvInferVersion.h") 57 | set(_VersionSearchFile "${TensorRT_INCLUDE_DIR}/NvInferVersion.h") 58 | else () 59 | set(_VersionSearchFile "${TensorRT_INCLUDE_DIR}/NvInfer.h") 60 | endif () 61 | file(STRINGS "${_VersionSearchFile}" TensorRT_MAJOR REGEX "^#define NV_TENSORRT_MAJOR [0-9]+.*$") 62 | file(STRINGS "${_VersionSearchFile}" TensorRT_MINOR REGEX "^#define NV_TENSORRT_MINOR [0-9]+.*$") 63 | file(STRINGS "${_VersionSearchFile}" TensorRT_PATCH REGEX "^#define NV_TENSORRT_PATCH [0-9]+.*$") 64 | 65 | string(REGEX REPLACE "^#define NV_TENSORRT_MAJOR ([0-9]+).*$" "\\1" TensorRT_VERSION_MAJOR "${TensorRT_MAJOR}") 66 | string(REGEX REPLACE "^#define NV_TENSORRT_MINOR ([0-9]+).*$" "\\1" TensorRT_VERSION_MINOR "${TensorRT_MINOR}") 67 | string(REGEX REPLACE "^#define NV_TENSORRT_PATCH ([0-9]+).*$" "\\1" TensorRT_VERSION_PATCH "${TensorRT_PATCH}") 68 | set(TensorRT_VERSION_STRING "${TensorRT_VERSION_MAJOR}.${TensorRT_VERSION_MINOR}.${TensorRT_VERSION_PATCH}") 69 | endif() 70 | 71 | include(FindPackageHandleStandardArgs) 72 | FIND_PACKAGE_HANDLE_STANDARD_ARGS(TensorRT REQUIRED_VARS TensorRT_LIBRARY TensorRT_INCLUDE_DIR VERSION_VAR TensorRT_VERSION_STRING) 73 | 74 | if(TensorRT_FOUND) 75 | set(TensorRT_INCLUDE_DIRS ${TensorRT_INCLUDE_DIR}) 76 | 77 | if(NOT TensorRT_LIBRARIES) 78 | set(TensorRT_LIBRARIES ${TensorRT_LIBRARY}) 79 | if (TensorRT_nvinfer_plugin_LIBRARY) 80 | list(APPEND TensorRT_LIBRARIES ${TensorRT_nvinfer_plugin_LIBRARY}) 81 | endif() 82 | endif() 83 | 84 | if(NOT TARGET TensorRT::TensorRT) 85 | add_library(TensorRT INTERFACE IMPORTED) 86 | add_library(TensorRT::TensorRT ALIAS TensorRT) 87 | endif() 88 | 89 | if(NOT TARGET TensorRT::nvinfer) 90 | add_library(TensorRT::nvinfer SHARED IMPORTED) 91 | if (WIN32) 92 | foreach(search ${_TensorRT_SEARCHES}) 93 | find_file(TensorRT_LIBRARY_DLL 94 | NAMES nvinfer.dll 95 | PATHS ${${search}} 96 | PATH_SUFFIXES bin 97 | ) 98 | endforeach() 99 | 100 | set_target_properties(TensorRT::nvinfer PROPERTIES 101 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 102 | IMPORTED_LOCATION "${TensorRT_LIBRARY_DLL}" 103 | IMPORTED_IMPLIB "${TensorRT_LIBRARY}" 104 | ) 105 | else() 106 | set_target_properties(TensorRT::nvinfer PROPERTIES 107 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 108 | IMPORTED_LOCATION "${TensorRT_LIBRARY}" 109 | ) 110 | endif() 111 | target_link_libraries(TensorRT INTERFACE TensorRT::nvinfer) 112 | endif() 113 | 114 | if(NOT TARGET TensorRT::nvinfer_plugin AND TensorRT_nvinfer_plugin_LIBRARY) 115 | add_library(TensorRT::nvinfer_plugin SHARED IMPORTED) 116 | if (WIN32) 117 | foreach(search ${_TensorRT_SEARCHES}) 118 | find_file(TensorRT_nvinfer_plugin_LIBRARY_DLL 119 | NAMES nvinfer_plugin.dll 120 | PATHS ${${search}} 121 | PATH_SUFFIXES bin 122 | ) 123 | endforeach() 124 | 125 | set_target_properties(TensorRT::nvinfer_plugin PROPERTIES 126 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 127 | IMPORTED_LOCATION "${TensorRT_nvinfer_plugin_LIBRARY_DLL}" 128 | IMPORTED_IMPLIB "${TensorRT_nvinfer_plugin_LIBRARY}" 129 | ) 130 | else() 131 | set_target_properties(TensorRT::nvinfer_plugin PROPERTIES 132 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 133 | IMPORTED_LOCATION "${TensorRT_nvinfer_plugin_LIBRARY}" 134 | ) 135 | endif() 136 | target_link_libraries(TensorRT INTERFACE TensorRT::nvinfer_plugin) 137 | endif() 138 | endif() 139 | -------------------------------------------------------------------------------- /csrc/obb/normal/cmake/Function.cmake: -------------------------------------------------------------------------------- 1 | 2 | function(print_var var) 3 | set(value "${${var}}") 4 | string(LENGTH "${value}" value_length) 5 | if(value_length GREATER 0) 6 | math(EXPR last_index "${value_length} - 1") 7 | string(SUBSTRING "${value}" ${last_index} ${last_index} last_char) 8 | endif() 9 | 10 | if(NOT "${last_char}" STREQUAL "\n") 11 | set(value "${value}\n") 12 | endif() 13 | message(STATUS "${var}:\n ${value}") 14 | endfunction() 15 | -------------------------------------------------------------------------------- /csrc/obb/normal/include/common.hpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 4/7/23. 3 | // 4 | 5 | #ifndef POSE_NORMAL_COMMON_HPP 6 | #define POSE_NORMAL_COMMON_HPP 7 | #include "NvInfer.h" 8 | #include "filesystem.hpp" 9 | #include "opencv2/opencv.hpp" 10 | 11 | #define CHECK(call) \ 12 | do { \ 13 | const cudaError_t error_code = call; \ 14 | if (error_code != cudaSuccess) { \ 15 | printf("CUDA Error:\n"); \ 16 | printf(" File: %s\n", __FILE__); \ 17 | printf(" Line: %d\n", __LINE__); \ 18 | printf(" Error code: %d\n", error_code); \ 19 | printf(" Error text: %s\n", cudaGetErrorString(error_code)); \ 20 | exit(1); \ 21 | } \ 22 | } while (0) 23 | 24 | class Logger: public nvinfer1::ILogger { 25 | public: 26 | nvinfer1::ILogger::Severity reportableSeverity; 27 | 28 | explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO): 29 | reportableSeverity(severity) 30 | { 31 | } 32 | 33 | void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override 34 | { 35 | if (severity > reportableSeverity) { 36 | return; 37 | } 38 | switch (severity) { 39 | case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: 40 | std::cerr << "INTERNAL_ERROR: "; 41 | break; 42 | case nvinfer1::ILogger::Severity::kERROR: 43 | std::cerr << "ERROR: "; 44 | break; 45 | case nvinfer1::ILogger::Severity::kWARNING: 46 | std::cerr << "WARNING: "; 47 | break; 48 | case nvinfer1::ILogger::Severity::kINFO: 49 | std::cerr << "INFO: "; 50 | break; 51 | default: 52 | std::cerr << "VERBOSE: "; 53 | break; 54 | } 55 | std::cerr << msg << std::endl; 56 | } 57 | }; 58 | 59 | inline int get_size_by_dims(const nvinfer1::Dims& dims) 60 | { 61 | int size = 1; 62 | for (int i = 0; i < dims.nbDims; i++) { 63 | size *= dims.d[i]; 64 | } 65 | return size; 66 | } 67 | 68 | inline int type_to_size(const nvinfer1::DataType& dataType) 69 | { 70 | switch (dataType) { 71 | case nvinfer1::DataType::kFLOAT: 72 | return 4; 73 | case nvinfer1::DataType::kHALF: 74 | return 2; 75 | case nvinfer1::DataType::kINT32: 76 | return 4; 77 | case nvinfer1::DataType::kINT8: 78 | return 1; 79 | case nvinfer1::DataType::kBOOL: 80 | return 1; 81 | default: 82 | return 4; 83 | } 84 | } 85 | 86 | inline static float clamp(float val, float min, float max) 87 | { 88 | return val > min ? (val < max ? val : max) : min; 89 | } 90 | 91 | namespace obb { 92 | struct Binding { 93 | size_t size = 1; 94 | size_t dsize = 1; 95 | nvinfer1::Dims dims; 96 | std::string name; 97 | }; 98 | 99 | struct Object { 100 | cv::RotatedRect rect; 101 | int label = 0; 102 | float prob = 0.0; 103 | }; 104 | 105 | struct PreParam { 106 | float ratio = 1.0f; 107 | float dw = 0.0f; 108 | float dh = 0.0f; 109 | float height = 0; 110 | float width = 0; 111 | }; 112 | } // namespace obb 113 | #endif // POSE_NORMAL_COMMON_HPP 114 | -------------------------------------------------------------------------------- /csrc/obb/normal/main.cpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 4/7/23. 3 | // 4 | #include "opencv2/opencv.hpp" 5 | #include "yolov8-obb.hpp" 6 | #include 7 | 8 | namespace fs = ghc::filesystem; 9 | 10 | const std::vector CLASS_NAMES = {"plane", 11 | "ship", 12 | "storage tank", 13 | "baseball diamond", 14 | "tennis court", 15 | "basketball court", 16 | "ground track field", 17 | "harbor", 18 | "bridge", 19 | "large vehicle", 20 | "small vehicle", 21 | "helicopter", 22 | "roundabout", 23 | "soccer ball field", 24 | "swimming pool"}; 25 | 26 | const std::vector> COLORS = {{0, 114, 189}, 27 | {217, 83, 25}, 28 | {237, 177, 32}, 29 | {126, 47, 142}, 30 | {119, 172, 48}, 31 | {77, 190, 238}, 32 | {162, 20, 47}, 33 | {76, 76, 76}, 34 | {153, 153, 153}, 35 | {255, 0, 0}, 36 | {255, 128, 0}, 37 | {191, 191, 0}, 38 | {0, 255, 0}, 39 | {0, 0, 255}, 40 | {170, 0, 255}}; 41 | 42 | int main(int argc, char** argv) 43 | { 44 | if (argc != 3) { 45 | fprintf(stderr, "Usage: %s [engine_path] [image_path/image_dir/video_path]\n", argv[0]); 46 | return -1; 47 | } 48 | 49 | // cuda:0 50 | cudaSetDevice(0); 51 | 52 | const std::string engine_file_path{argv[1]}; 53 | const fs::path path{argv[2]}; 54 | 55 | std::vector imagePathList; 56 | bool isVideo{false}; 57 | 58 | assert(argc == 3); 59 | 60 | auto yolov8_obb = new YOLOv8_obb(engine_file_path); 61 | yolov8_obb->make_pipe(true); 62 | 63 | if (fs::exists(path)) { 64 | std::string suffix = path.extension(); 65 | if (suffix == ".jpg" || suffix == ".jpeg" || suffix == ".png") { 66 | imagePathList.push_back(path); 67 | } 68 | else if (suffix == ".mp4" || suffix == ".avi" || suffix == ".m4v" || suffix == ".mpeg" || suffix == ".mov" 69 | || suffix == ".mkv") { 70 | isVideo = true; 71 | } 72 | else { 73 | printf("suffix %s is wrong !!!\n", suffix.c_str()); 74 | std::abort(); 75 | } 76 | } 77 | else if (fs::is_directory(path)) { 78 | cv::glob(path.string() + "/*.jpg", imagePathList); 79 | } 80 | 81 | cv::Mat res, image; 82 | cv::Size size = cv::Size{1024, 1024}; 83 | int num_labels = 15; 84 | int topk = 100; 85 | float score_thres = 0.25f; 86 | float iou_thres = 0.65f; 87 | 88 | std::vector objs; 89 | 90 | cv::namedWindow("result", cv::WINDOW_AUTOSIZE); 91 | 92 | if (isVideo) { 93 | cv::VideoCapture cap(path); 94 | 95 | if (!cap.isOpened()) { 96 | printf("can not open %s\n", path.c_str()); 97 | return -1; 98 | } 99 | while (cap.read(image)) { 100 | objs.clear(); 101 | yolov8_obb->copy_from_Mat(image, size); 102 | auto start = std::chrono::system_clock::now(); 103 | yolov8_obb->infer(); 104 | auto end = std::chrono::system_clock::now(); 105 | yolov8_obb->postprocess(objs, score_thres, iou_thres, topk, num_labels); 106 | yolov8_obb->draw_objects(image, res, objs, CLASS_NAMES, COLORS); 107 | auto tc = (double)std::chrono::duration_cast(end - start).count() / 1000.; 108 | printf("cost %2.4lf ms\n", tc); 109 | cv::imshow("result", res); 110 | if (cv::waitKey(10) == 'q') { 111 | break; 112 | } 113 | } 114 | } 115 | else { 116 | for (auto& p : imagePathList) { 117 | objs.clear(); 118 | image = cv::imread(p); 119 | yolov8_obb->copy_from_Mat(image, size); 120 | auto start = std::chrono::system_clock::now(); 121 | yolov8_obb->infer(); 122 | auto end = std::chrono::system_clock::now(); 123 | yolov8_obb->postprocess(objs, score_thres, iou_thres, topk, num_labels); 124 | yolov8_obb->draw_objects(image, res, objs, CLASS_NAMES, COLORS); 125 | auto tc = (double)std::chrono::duration_cast(end - start).count() / 1000.; 126 | printf("cost %2.4lf ms\n", tc); 127 | cv::imshow("result", res); 128 | cv::waitKey(0); 129 | } 130 | } 131 | cv::destroyAllWindows(); 132 | delete yolov8_obb; 133 | return 0; 134 | } 135 | -------------------------------------------------------------------------------- /csrc/pose/normal/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.1) 2 | 3 | set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86 89 90) 4 | set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc) 5 | 6 | project(yolov8-pose LANGUAGES CXX CUDA) 7 | 8 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -O3") 9 | set(CMAKE_CXX_STANDARD 14) 10 | set(CMAKE_BUILD_TYPE Release) 11 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 12 | 13 | list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake") 14 | include(Function) 15 | 16 | # CUDA 17 | find_package(CUDA REQUIRED) 18 | print_var(CUDA_LIBRARIES) 19 | print_var(CUDA_INCLUDE_DIRS) 20 | get_filename_component(CUDA_LIB_DIR ${CUDA_LIBRARIES} DIRECTORY) 21 | print_var(CUDA_LIB_DIR) 22 | 23 | # OpenCV 24 | find_package(OpenCV REQUIRED) 25 | print_var(OpenCV_LIBS) 26 | print_var(OpenCV_LIBRARIES) 27 | print_var(OpenCV_INCLUDE_DIRS) 28 | 29 | # TensorRT 30 | find_package(TensorRT REQUIRED) 31 | print_var(TensorRT_LIBRARIES) 32 | print_var(TensorRT_INCLUDE_DIRS) 33 | print_var(TensorRT_LIB_DIR) 34 | if (TensorRT_VERSION_MAJOR GREATER_EQUAL 10) 35 | message(STATUS "Build with -DTRT_10") 36 | add_definitions(-DTRT_10) 37 | endif () 38 | 39 | list(APPEND ALL_INCLUDE_DIRS 40 | ${CUDA_INCLUDE_DIRS} 41 | ${OpenCV_INCLUDE_DIRS} 42 | ${TensorRT_INCLUDE_DIRS} 43 | ${CMAKE_CURRENT_SOURCE_DIR}/include 44 | ) 45 | 46 | list(APPEND ALL_LIBS 47 | ${CUDA_LIBRARIES} 48 | ${OpenCV_LIBRARIES} 49 | ${TensorRT_LIBRARIES} 50 | ) 51 | 52 | list(APPEND ALL_LIB_DIRS 53 | ${CUDA_LIB_DIR} 54 | ${TensorRT_LIB_DIR} 55 | ) 56 | 57 | print_var(ALL_INCLUDE_DIRS) 58 | print_var(ALL_LIBS) 59 | print_var(ALL_LIB_DIRS) 60 | 61 | 62 | add_executable(${PROJECT_NAME} 63 | ${CMAKE_CURRENT_SOURCE_DIR}/main.cpp 64 | ${CMAKE_CURRENT_SOURCE_DIR}/include/yolov8-pose.hpp 65 | ${CMAKE_CURRENT_SOURCE_DIR}/include/common.hpp 66 | ) 67 | 68 | target_include_directories( 69 | ${PROJECT_NAME} 70 | PUBLIC 71 | ${ALL_INCLUDE_DIRS} 72 | ) 73 | 74 | target_link_directories( 75 | ${PROJECT_NAME} 76 | PUBLIC 77 | ${ALL_LIB_DIRS} 78 | ) 79 | 80 | target_link_libraries( 81 | ${PROJECT_NAME} 82 | PRIVATE 83 | ${ALL_LIBS} 84 | ) 85 | 86 | if (${OpenCV_VERSION} VERSION_GREATER_EQUAL 4.7.0) 87 | message(STATUS "Build with -DBATCHED_NMS") 88 | add_definitions(-DBATCHED_NMS) 89 | endif () 90 | -------------------------------------------------------------------------------- /csrc/pose/normal/cmake/FindTensorRT.cmake: -------------------------------------------------------------------------------- 1 | # This module defines the following variables: 2 | # 3 | # :: 4 | # 5 | # TensorRT_INCLUDE_DIRS 6 | # TensorRT_LIBRARIES 7 | # TensorRT_FOUND 8 | # 9 | # :: 10 | # 11 | # TensorRT_VERSION_STRING - version (x.y.z) 12 | # TensorRT_VERSION_MAJOR - major version (x) 13 | # TensorRT_VERSION_MINOR - minor version (y) 14 | # TensorRT_VERSION_PATCH - patch version (z) 15 | # 16 | # Hints 17 | # ^^^^^ 18 | # A user may set ``TensorRT_ROOT`` to an installation root to tell this module where to look. 19 | # 20 | set(_TensorRT_SEARCHES) 21 | 22 | if(TensorRT_ROOT) 23 | set(_TensorRT_SEARCH_ROOT PATHS ${TensorRT_ROOT} NO_DEFAULT_PATH) 24 | list(APPEND _TensorRT_SEARCHES _TensorRT_SEARCH_ROOT) 25 | endif() 26 | 27 | # appends some common paths 28 | set(_TensorRT_SEARCH_NORMAL 29 | PATHS "/usr" 30 | ) 31 | list(APPEND _TensorRT_SEARCHES _TensorRT_SEARCH_NORMAL) 32 | 33 | # Include dir 34 | foreach(search ${_TensorRT_SEARCHES}) 35 | find_path(TensorRT_INCLUDE_DIR NAMES NvInfer.h ${${search}} PATH_SUFFIXES include) 36 | endforeach() 37 | 38 | if(NOT TensorRT_LIBRARY) 39 | foreach(search ${_TensorRT_SEARCHES}) 40 | find_library(TensorRT_LIBRARY NAMES nvinfer ${${search}} PATH_SUFFIXES lib) 41 | if(NOT TensorRT_LIB_DIR) 42 | get_filename_component(TensorRT_LIB_DIR ${TensorRT_LIBRARY} DIRECTORY) 43 | endif () 44 | endforeach() 45 | endif() 46 | 47 | if(NOT TensorRT_nvinfer_plugin_LIBRARY) 48 | foreach(search ${_TensorRT_SEARCHES}) 49 | find_library(TensorRT_nvinfer_plugin_LIBRARY NAMES nvinfer_plugin ${${search}} PATH_SUFFIXES lib) 50 | endforeach() 51 | endif() 52 | 53 | mark_as_advanced(TensorRT_INCLUDE_DIR) 54 | 55 | if(TensorRT_INCLUDE_DIR AND EXISTS "${TensorRT_INCLUDE_DIR}/NvInfer.h") 56 | if(EXISTS "${TensorRT_INCLUDE_DIR}/NvInferVersion.h") 57 | set(_VersionSearchFile "${TensorRT_INCLUDE_DIR}/NvInferVersion.h") 58 | else () 59 | set(_VersionSearchFile "${TensorRT_INCLUDE_DIR}/NvInfer.h") 60 | endif () 61 | file(STRINGS "${_VersionSearchFile}" TensorRT_MAJOR REGEX "^#define NV_TENSORRT_MAJOR [0-9]+.*$") 62 | file(STRINGS "${_VersionSearchFile}" TensorRT_MINOR REGEX "^#define NV_TENSORRT_MINOR [0-9]+.*$") 63 | file(STRINGS "${_VersionSearchFile}" TensorRT_PATCH REGEX "^#define NV_TENSORRT_PATCH [0-9]+.*$") 64 | 65 | string(REGEX REPLACE "^#define NV_TENSORRT_MAJOR ([0-9]+).*$" "\\1" TensorRT_VERSION_MAJOR "${TensorRT_MAJOR}") 66 | string(REGEX REPLACE "^#define NV_TENSORRT_MINOR ([0-9]+).*$" "\\1" TensorRT_VERSION_MINOR "${TensorRT_MINOR}") 67 | string(REGEX REPLACE "^#define NV_TENSORRT_PATCH ([0-9]+).*$" "\\1" TensorRT_VERSION_PATCH "${TensorRT_PATCH}") 68 | set(TensorRT_VERSION_STRING "${TensorRT_VERSION_MAJOR}.${TensorRT_VERSION_MINOR}.${TensorRT_VERSION_PATCH}") 69 | endif() 70 | 71 | include(FindPackageHandleStandardArgs) 72 | FIND_PACKAGE_HANDLE_STANDARD_ARGS(TensorRT REQUIRED_VARS TensorRT_LIBRARY TensorRT_INCLUDE_DIR VERSION_VAR TensorRT_VERSION_STRING) 73 | 74 | if(TensorRT_FOUND) 75 | set(TensorRT_INCLUDE_DIRS ${TensorRT_INCLUDE_DIR}) 76 | 77 | if(NOT TensorRT_LIBRARIES) 78 | set(TensorRT_LIBRARIES ${TensorRT_LIBRARY}) 79 | if (TensorRT_nvinfer_plugin_LIBRARY) 80 | list(APPEND TensorRT_LIBRARIES ${TensorRT_nvinfer_plugin_LIBRARY}) 81 | endif() 82 | endif() 83 | 84 | if(NOT TARGET TensorRT::TensorRT) 85 | add_library(TensorRT INTERFACE IMPORTED) 86 | add_library(TensorRT::TensorRT ALIAS TensorRT) 87 | endif() 88 | 89 | if(NOT TARGET TensorRT::nvinfer) 90 | add_library(TensorRT::nvinfer SHARED IMPORTED) 91 | if (WIN32) 92 | foreach(search ${_TensorRT_SEARCHES}) 93 | find_file(TensorRT_LIBRARY_DLL 94 | NAMES nvinfer.dll 95 | PATHS ${${search}} 96 | PATH_SUFFIXES bin 97 | ) 98 | endforeach() 99 | 100 | set_target_properties(TensorRT::nvinfer PROPERTIES 101 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 102 | IMPORTED_LOCATION "${TensorRT_LIBRARY_DLL}" 103 | IMPORTED_IMPLIB "${TensorRT_LIBRARY}" 104 | ) 105 | else() 106 | set_target_properties(TensorRT::nvinfer PROPERTIES 107 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 108 | IMPORTED_LOCATION "${TensorRT_LIBRARY}" 109 | ) 110 | endif() 111 | target_link_libraries(TensorRT INTERFACE TensorRT::nvinfer) 112 | endif() 113 | 114 | if(NOT TARGET TensorRT::nvinfer_plugin AND TensorRT_nvinfer_plugin_LIBRARY) 115 | add_library(TensorRT::nvinfer_plugin SHARED IMPORTED) 116 | if (WIN32) 117 | foreach(search ${_TensorRT_SEARCHES}) 118 | find_file(TensorRT_nvinfer_plugin_LIBRARY_DLL 119 | NAMES nvinfer_plugin.dll 120 | PATHS ${${search}} 121 | PATH_SUFFIXES bin 122 | ) 123 | endforeach() 124 | 125 | set_target_properties(TensorRT::nvinfer_plugin PROPERTIES 126 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 127 | IMPORTED_LOCATION "${TensorRT_nvinfer_plugin_LIBRARY_DLL}" 128 | IMPORTED_IMPLIB "${TensorRT_nvinfer_plugin_LIBRARY}" 129 | ) 130 | else() 131 | set_target_properties(TensorRT::nvinfer_plugin PROPERTIES 132 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 133 | IMPORTED_LOCATION "${TensorRT_nvinfer_plugin_LIBRARY}" 134 | ) 135 | endif() 136 | target_link_libraries(TensorRT INTERFACE TensorRT::nvinfer_plugin) 137 | endif() 138 | endif() 139 | -------------------------------------------------------------------------------- /csrc/pose/normal/cmake/Function.cmake: -------------------------------------------------------------------------------- 1 | 2 | function(print_var var) 3 | set(value "${${var}}") 4 | string(LENGTH "${value}" value_length) 5 | if(value_length GREATER 0) 6 | math(EXPR last_index "${value_length} - 1") 7 | string(SUBSTRING "${value}" ${last_index} ${last_index} last_char) 8 | endif() 9 | 10 | if(NOT "${last_char}" STREQUAL "\n") 11 | set(value "${value}\n") 12 | endif() 13 | message(STATUS "${var}:\n ${value}") 14 | endfunction() 15 | -------------------------------------------------------------------------------- /csrc/pose/normal/include/common.hpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 4/7/23. 3 | // 4 | 5 | #ifndef POSE_NORMAL_COMMON_HPP 6 | #define POSE_NORMAL_COMMON_HPP 7 | #include "NvInfer.h" 8 | #include "filesystem.hpp" 9 | #include "opencv2/opencv.hpp" 10 | 11 | #define CHECK(call) \ 12 | do { \ 13 | const cudaError_t error_code = call; \ 14 | if (error_code != cudaSuccess) { \ 15 | printf("CUDA Error:\n"); \ 16 | printf(" File: %s\n", __FILE__); \ 17 | printf(" Line: %d\n", __LINE__); \ 18 | printf(" Error code: %d\n", error_code); \ 19 | printf(" Error text: %s\n", cudaGetErrorString(error_code)); \ 20 | exit(1); \ 21 | } \ 22 | } while (0) 23 | 24 | class Logger: public nvinfer1::ILogger { 25 | public: 26 | nvinfer1::ILogger::Severity reportableSeverity; 27 | 28 | explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO): 29 | reportableSeverity(severity) 30 | { 31 | } 32 | 33 | void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override 34 | { 35 | if (severity > reportableSeverity) { 36 | return; 37 | } 38 | switch (severity) { 39 | case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: 40 | std::cerr << "INTERNAL_ERROR: "; 41 | break; 42 | case nvinfer1::ILogger::Severity::kERROR: 43 | std::cerr << "ERROR: "; 44 | break; 45 | case nvinfer1::ILogger::Severity::kWARNING: 46 | std::cerr << "WARNING: "; 47 | break; 48 | case nvinfer1::ILogger::Severity::kINFO: 49 | std::cerr << "INFO: "; 50 | break; 51 | default: 52 | std::cerr << "VERBOSE: "; 53 | break; 54 | } 55 | std::cerr << msg << std::endl; 56 | } 57 | }; 58 | 59 | inline int get_size_by_dims(const nvinfer1::Dims& dims) 60 | { 61 | int size = 1; 62 | for (int i = 0; i < dims.nbDims; i++) { 63 | size *= dims.d[i]; 64 | } 65 | return size; 66 | } 67 | 68 | inline int type_to_size(const nvinfer1::DataType& dataType) 69 | { 70 | switch (dataType) { 71 | case nvinfer1::DataType::kFLOAT: 72 | return 4; 73 | case nvinfer1::DataType::kHALF: 74 | return 2; 75 | case nvinfer1::DataType::kINT32: 76 | return 4; 77 | case nvinfer1::DataType::kINT8: 78 | return 1; 79 | case nvinfer1::DataType::kBOOL: 80 | return 1; 81 | default: 82 | return 4; 83 | } 84 | } 85 | 86 | inline static float clamp(float val, float min, float max) 87 | { 88 | return val > min ? (val < max ? val : max) : min; 89 | } 90 | 91 | namespace pose { 92 | struct Binding { 93 | size_t size = 1; 94 | size_t dsize = 1; 95 | nvinfer1::Dims dims; 96 | std::string name; 97 | }; 98 | 99 | struct Object { 100 | cv::Rect_ rect; 101 | int label = 0; 102 | float prob = 0.0; 103 | std::vector kps; 104 | }; 105 | 106 | struct PreParam { 107 | float ratio = 1.0f; 108 | float dw = 0.0f; 109 | float dh = 0.0f; 110 | float height = 0; 111 | float width = 0; 112 | }; 113 | } // namespace pose 114 | #endif // POSE_NORMAL_COMMON_HPP 115 | -------------------------------------------------------------------------------- /csrc/segment/normal/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.1) 2 | 3 | set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86 89 90) 4 | set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc) 5 | 6 | project(yolov8-seg LANGUAGES CXX CUDA) 7 | 8 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -O3") 9 | set(CMAKE_CXX_STANDARD 14) 10 | set(CMAKE_BUILD_TYPE Release) 11 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 12 | 13 | list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake") 14 | include(Function) 15 | 16 | # CUDA 17 | find_package(CUDA REQUIRED) 18 | print_var(CUDA_LIBRARIES) 19 | print_var(CUDA_INCLUDE_DIRS) 20 | get_filename_component(CUDA_LIB_DIR ${CUDA_LIBRARIES} DIRECTORY) 21 | print_var(CUDA_LIB_DIR) 22 | 23 | # OpenCV 24 | find_package(OpenCV REQUIRED) 25 | print_var(OpenCV_LIBS) 26 | print_var(OpenCV_LIBRARIES) 27 | print_var(OpenCV_INCLUDE_DIRS) 28 | 29 | # TensorRT 30 | find_package(TensorRT REQUIRED) 31 | print_var(TensorRT_LIBRARIES) 32 | print_var(TensorRT_INCLUDE_DIRS) 33 | print_var(TensorRT_LIB_DIR) 34 | if (TensorRT_VERSION_MAJOR GREATER_EQUAL 10) 35 | message(STATUS "Build with -DTRT_10") 36 | add_definitions(-DTRT_10) 37 | endif () 38 | 39 | list(APPEND ALL_INCLUDE_DIRS 40 | ${CUDA_INCLUDE_DIRS} 41 | ${OpenCV_INCLUDE_DIRS} 42 | ${TensorRT_INCLUDE_DIRS} 43 | ${CMAKE_CURRENT_SOURCE_DIR}/include 44 | ) 45 | 46 | list(APPEND ALL_LIBS 47 | ${CUDA_LIBRARIES} 48 | ${OpenCV_LIBRARIES} 49 | ${TensorRT_LIBRARIES} 50 | ) 51 | 52 | list(APPEND ALL_LIB_DIRS 53 | ${CUDA_LIB_DIR} 54 | ${TensorRT_LIB_DIR} 55 | ) 56 | 57 | print_var(ALL_INCLUDE_DIRS) 58 | print_var(ALL_LIBS) 59 | print_var(ALL_LIB_DIRS) 60 | 61 | 62 | add_executable(${PROJECT_NAME} 63 | ${CMAKE_CURRENT_SOURCE_DIR}/main.cpp 64 | ${CMAKE_CURRENT_SOURCE_DIR}/include/yolov8-seg.hpp 65 | ${CMAKE_CURRENT_SOURCE_DIR}/include/common.hpp 66 | ) 67 | 68 | target_include_directories( 69 | ${PROJECT_NAME} 70 | PUBLIC 71 | ${ALL_INCLUDE_DIRS} 72 | ) 73 | 74 | target_link_directories( 75 | ${PROJECT_NAME} 76 | PUBLIC 77 | ${ALL_LIB_DIRS} 78 | ) 79 | 80 | target_link_libraries( 81 | ${PROJECT_NAME} 82 | PRIVATE 83 | ${ALL_LIBS} 84 | ) 85 | 86 | if (${OpenCV_VERSION} VERSION_GREATER_EQUAL 4.7.0) 87 | message(STATUS "Build with -DBATCHED_NMS") 88 | add_definitions(-DBATCHED_NMS) 89 | endif () 90 | -------------------------------------------------------------------------------- /csrc/segment/normal/cmake/FindTensorRT.cmake: -------------------------------------------------------------------------------- 1 | # This module defines the following variables: 2 | # 3 | # :: 4 | # 5 | # TensorRT_INCLUDE_DIRS 6 | # TensorRT_LIBRARIES 7 | # TensorRT_FOUND 8 | # 9 | # :: 10 | # 11 | # TensorRT_VERSION_STRING - version (x.y.z) 12 | # TensorRT_VERSION_MAJOR - major version (x) 13 | # TensorRT_VERSION_MINOR - minor version (y) 14 | # TensorRT_VERSION_PATCH - patch version (z) 15 | # 16 | # Hints 17 | # ^^^^^ 18 | # A user may set ``TensorRT_ROOT`` to an installation root to tell this module where to look. 19 | # 20 | set(_TensorRT_SEARCHES) 21 | 22 | if(TensorRT_ROOT) 23 | set(_TensorRT_SEARCH_ROOT PATHS ${TensorRT_ROOT} NO_DEFAULT_PATH) 24 | list(APPEND _TensorRT_SEARCHES _TensorRT_SEARCH_ROOT) 25 | endif() 26 | 27 | # appends some common paths 28 | set(_TensorRT_SEARCH_NORMAL 29 | PATHS "/usr" 30 | ) 31 | list(APPEND _TensorRT_SEARCHES _TensorRT_SEARCH_NORMAL) 32 | 33 | # Include dir 34 | foreach(search ${_TensorRT_SEARCHES}) 35 | find_path(TensorRT_INCLUDE_DIR NAMES NvInfer.h ${${search}} PATH_SUFFIXES include) 36 | endforeach() 37 | 38 | if(NOT TensorRT_LIBRARY) 39 | foreach(search ${_TensorRT_SEARCHES}) 40 | find_library(TensorRT_LIBRARY NAMES nvinfer ${${search}} PATH_SUFFIXES lib) 41 | if(NOT TensorRT_LIB_DIR) 42 | get_filename_component(TensorRT_LIB_DIR ${TensorRT_LIBRARY} DIRECTORY) 43 | endif () 44 | endforeach() 45 | endif() 46 | 47 | if(NOT TensorRT_nvinfer_plugin_LIBRARY) 48 | foreach(search ${_TensorRT_SEARCHES}) 49 | find_library(TensorRT_nvinfer_plugin_LIBRARY NAMES nvinfer_plugin ${${search}} PATH_SUFFIXES lib) 50 | endforeach() 51 | endif() 52 | 53 | mark_as_advanced(TensorRT_INCLUDE_DIR) 54 | 55 | if(TensorRT_INCLUDE_DIR AND EXISTS "${TensorRT_INCLUDE_DIR}/NvInfer.h") 56 | if(EXISTS "${TensorRT_INCLUDE_DIR}/NvInferVersion.h") 57 | set(_VersionSearchFile "${TensorRT_INCLUDE_DIR}/NvInferVersion.h") 58 | else () 59 | set(_VersionSearchFile "${TensorRT_INCLUDE_DIR}/NvInfer.h") 60 | endif () 61 | file(STRINGS "${_VersionSearchFile}" TensorRT_MAJOR REGEX "^#define NV_TENSORRT_MAJOR [0-9]+.*$") 62 | file(STRINGS "${_VersionSearchFile}" TensorRT_MINOR REGEX "^#define NV_TENSORRT_MINOR [0-9]+.*$") 63 | file(STRINGS "${_VersionSearchFile}" TensorRT_PATCH REGEX "^#define NV_TENSORRT_PATCH [0-9]+.*$") 64 | 65 | string(REGEX REPLACE "^#define NV_TENSORRT_MAJOR ([0-9]+).*$" "\\1" TensorRT_VERSION_MAJOR "${TensorRT_MAJOR}") 66 | string(REGEX REPLACE "^#define NV_TENSORRT_MINOR ([0-9]+).*$" "\\1" TensorRT_VERSION_MINOR "${TensorRT_MINOR}") 67 | string(REGEX REPLACE "^#define NV_TENSORRT_PATCH ([0-9]+).*$" "\\1" TensorRT_VERSION_PATCH "${TensorRT_PATCH}") 68 | set(TensorRT_VERSION_STRING "${TensorRT_VERSION_MAJOR}.${TensorRT_VERSION_MINOR}.${TensorRT_VERSION_PATCH}") 69 | endif() 70 | 71 | include(FindPackageHandleStandardArgs) 72 | FIND_PACKAGE_HANDLE_STANDARD_ARGS(TensorRT REQUIRED_VARS TensorRT_LIBRARY TensorRT_INCLUDE_DIR VERSION_VAR TensorRT_VERSION_STRING) 73 | 74 | if(TensorRT_FOUND) 75 | set(TensorRT_INCLUDE_DIRS ${TensorRT_INCLUDE_DIR}) 76 | 77 | if(NOT TensorRT_LIBRARIES) 78 | set(TensorRT_LIBRARIES ${TensorRT_LIBRARY}) 79 | if (TensorRT_nvinfer_plugin_LIBRARY) 80 | list(APPEND TensorRT_LIBRARIES ${TensorRT_nvinfer_plugin_LIBRARY}) 81 | endif() 82 | endif() 83 | 84 | if(NOT TARGET TensorRT::TensorRT) 85 | add_library(TensorRT INTERFACE IMPORTED) 86 | add_library(TensorRT::TensorRT ALIAS TensorRT) 87 | endif() 88 | 89 | if(NOT TARGET TensorRT::nvinfer) 90 | add_library(TensorRT::nvinfer SHARED IMPORTED) 91 | if (WIN32) 92 | foreach(search ${_TensorRT_SEARCHES}) 93 | find_file(TensorRT_LIBRARY_DLL 94 | NAMES nvinfer.dll 95 | PATHS ${${search}} 96 | PATH_SUFFIXES bin 97 | ) 98 | endforeach() 99 | 100 | set_target_properties(TensorRT::nvinfer PROPERTIES 101 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 102 | IMPORTED_LOCATION "${TensorRT_LIBRARY_DLL}" 103 | IMPORTED_IMPLIB "${TensorRT_LIBRARY}" 104 | ) 105 | else() 106 | set_target_properties(TensorRT::nvinfer PROPERTIES 107 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 108 | IMPORTED_LOCATION "${TensorRT_LIBRARY}" 109 | ) 110 | endif() 111 | target_link_libraries(TensorRT INTERFACE TensorRT::nvinfer) 112 | endif() 113 | 114 | if(NOT TARGET TensorRT::nvinfer_plugin AND TensorRT_nvinfer_plugin_LIBRARY) 115 | add_library(TensorRT::nvinfer_plugin SHARED IMPORTED) 116 | if (WIN32) 117 | foreach(search ${_TensorRT_SEARCHES}) 118 | find_file(TensorRT_nvinfer_plugin_LIBRARY_DLL 119 | NAMES nvinfer_plugin.dll 120 | PATHS ${${search}} 121 | PATH_SUFFIXES bin 122 | ) 123 | endforeach() 124 | 125 | set_target_properties(TensorRT::nvinfer_plugin PROPERTIES 126 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 127 | IMPORTED_LOCATION "${TensorRT_nvinfer_plugin_LIBRARY_DLL}" 128 | IMPORTED_IMPLIB "${TensorRT_nvinfer_plugin_LIBRARY}" 129 | ) 130 | else() 131 | set_target_properties(TensorRT::nvinfer_plugin PROPERTIES 132 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 133 | IMPORTED_LOCATION "${TensorRT_nvinfer_plugin_LIBRARY}" 134 | ) 135 | endif() 136 | target_link_libraries(TensorRT INTERFACE TensorRT::nvinfer_plugin) 137 | endif() 138 | endif() 139 | -------------------------------------------------------------------------------- /csrc/segment/normal/cmake/Function.cmake: -------------------------------------------------------------------------------- 1 | 2 | function(print_var var) 3 | set(value "${${var}}") 4 | string(LENGTH "${value}" value_length) 5 | if(value_length GREATER 0) 6 | math(EXPR last_index "${value_length} - 1") 7 | string(SUBSTRING "${value}" ${last_index} ${last_index} last_char) 8 | endif() 9 | 10 | if(NOT "${last_char}" STREQUAL "\n") 11 | set(value "${value}\n") 12 | endif() 13 | message(STATUS "${var}:\n ${value}") 14 | endfunction() 15 | -------------------------------------------------------------------------------- /csrc/segment/normal/include/common.hpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 2/8/23. 3 | // 4 | 5 | #ifndef SEGMENT_NORMAL_COMMON_HPP 6 | #define SEGMENT_NORMAL_COMMON_HPP 7 | #include "NvInfer.h" 8 | #include "filesystem.hpp" 9 | #include "opencv2/opencv.hpp" 10 | 11 | #define CHECK(call) \ 12 | do { \ 13 | const cudaError_t error_code = call; \ 14 | if (error_code != cudaSuccess) { \ 15 | printf("CUDA Error:\n"); \ 16 | printf(" File: %s\n", __FILE__); \ 17 | printf(" Line: %d\n", __LINE__); \ 18 | printf(" Error code: %d\n", error_code); \ 19 | printf(" Error text: %s\n", cudaGetErrorString(error_code)); \ 20 | exit(1); \ 21 | } \ 22 | } while (0) 23 | 24 | class Logger: public nvinfer1::ILogger { 25 | public: 26 | nvinfer1::ILogger::Severity reportableSeverity; 27 | 28 | explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO): 29 | reportableSeverity(severity) 30 | { 31 | } 32 | 33 | void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override 34 | { 35 | if (severity > reportableSeverity) { 36 | return; 37 | } 38 | switch (severity) { 39 | case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: 40 | std::cerr << "INTERNAL_ERROR: "; 41 | break; 42 | case nvinfer1::ILogger::Severity::kERROR: 43 | std::cerr << "ERROR: "; 44 | break; 45 | case nvinfer1::ILogger::Severity::kWARNING: 46 | std::cerr << "WARNING: "; 47 | break; 48 | case nvinfer1::ILogger::Severity::kINFO: 49 | std::cerr << "INFO: "; 50 | break; 51 | default: 52 | std::cerr << "VERBOSE: "; 53 | break; 54 | } 55 | std::cerr << msg << std::endl; 56 | } 57 | }; 58 | 59 | inline int get_size_by_dims(const nvinfer1::Dims& dims) 60 | { 61 | int size = 1; 62 | for (int i = 0; i < dims.nbDims; i++) { 63 | size *= dims.d[i]; 64 | } 65 | return size; 66 | } 67 | 68 | inline int type_to_size(const nvinfer1::DataType& dataType) 69 | { 70 | switch (dataType) { 71 | case nvinfer1::DataType::kFLOAT: 72 | return 4; 73 | case nvinfer1::DataType::kHALF: 74 | return 2; 75 | case nvinfer1::DataType::kINT32: 76 | return 4; 77 | case nvinfer1::DataType::kINT8: 78 | return 1; 79 | case nvinfer1::DataType::kBOOL: 80 | return 1; 81 | default: 82 | return 4; 83 | } 84 | } 85 | 86 | inline static float clamp(float val, float min, float max) 87 | { 88 | return val > min ? (val < max ? val : max) : min; 89 | } 90 | 91 | namespace seg { 92 | struct Binding { 93 | size_t size = 1; 94 | size_t dsize = 1; 95 | nvinfer1::Dims dims; 96 | std::string name; 97 | }; 98 | 99 | struct Object { 100 | cv::Rect_ rect; 101 | int label = 0; 102 | float prob = 0.0; 103 | cv::Mat boxMask; 104 | }; 105 | 106 | struct PreParam { 107 | float ratio = 1.0f; 108 | float dw = 0.0f; 109 | float dh = 0.0f; 110 | float height = 0; 111 | float width = 0; 112 | }; 113 | } // namespace seg 114 | #endif // SEGMENT_NORMAL_COMMON_HPP 115 | -------------------------------------------------------------------------------- /csrc/segment/simple/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.1) 2 | 3 | set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86 89 90) 4 | set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc) 5 | 6 | project(yolov8-seg LANGUAGES CXX CUDA) 7 | 8 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -O3") 9 | set(CMAKE_CXX_STANDARD 14) 10 | set(CMAKE_BUILD_TYPE Release) 11 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 12 | 13 | # CUDA 14 | find_package(CUDA REQUIRED) 15 | message(STATUS "CUDA Libs: \n${CUDA_LIBRARIES}\n") 16 | get_filename_component(CUDA_LIB_DIR ${CUDA_LIBRARIES} DIRECTORY) 17 | message(STATUS "CUDA Headers: \n${CUDA_INCLUDE_DIRS}\n") 18 | 19 | # OpenCV 20 | find_package(OpenCV REQUIRED) 21 | message(STATUS "OpenCV Libs: \n${OpenCV_LIBS}\n") 22 | message(STATUS "OpenCV Libraries: \n${OpenCV_LIBRARIES}\n") 23 | message(STATUS "OpenCV Headers: \n${OpenCV_INCLUDE_DIRS}\n") 24 | 25 | # TensorRT 26 | set(TensorRT_INCLUDE_DIRS /usr/include/x86_64-linux-gnu) 27 | set(TensorRT_LIBRARIES /usr/lib/x86_64-linux-gnu) 28 | 29 | 30 | message(STATUS "TensorRT Libs: \n${TensorRT_LIBRARIES}\n") 31 | message(STATUS "TensorRT Headers: \n${TensorRT_INCLUDE_DIRS}\n") 32 | 33 | list(APPEND INCLUDE_DIRS 34 | ${CUDA_INCLUDE_DIRS} 35 | ${OpenCV_INCLUDE_DIRS} 36 | ${TensorRT_INCLUDE_DIRS} 37 | ./include 38 | ) 39 | 40 | list(APPEND ALL_LIBS 41 | ${CUDA_LIBRARIES} 42 | ${CUDA_LIB_DIR} 43 | ${OpenCV_LIBRARIES} 44 | ${TensorRT_LIBRARIES} 45 | ) 46 | 47 | include_directories(${INCLUDE_DIRS}) 48 | 49 | add_executable(${PROJECT_NAME} 50 | main.cpp 51 | include/yolov8-seg.hpp 52 | include/common.hpp 53 | ) 54 | 55 | target_link_directories(${PROJECT_NAME} PUBLIC ${ALL_LIBS}) 56 | target_link_libraries(${PROJECT_NAME} PRIVATE nvinfer nvinfer_plugin cudart ${OpenCV_LIBS}) 57 | 58 | if (${OpenCV_VERSION} VERSION_GREATER_EQUAL 4.7.0) 59 | message(STATUS "Build with -DBATCHED_NMS") 60 | add_definitions(-DBATCHED_NMS) 61 | endif () 62 | -------------------------------------------------------------------------------- /csrc/segment/simple/cmake/FindTensorRT.cmake: -------------------------------------------------------------------------------- 1 | # This module defines the following variables: 2 | # 3 | # :: 4 | # 5 | # TensorRT_INCLUDE_DIRS 6 | # TensorRT_LIBRARIES 7 | # TensorRT_FOUND 8 | # 9 | # :: 10 | # 11 | # TensorRT_VERSION_STRING - version (x.y.z) 12 | # TensorRT_VERSION_MAJOR - major version (x) 13 | # TensorRT_VERSION_MINOR - minor version (y) 14 | # TensorRT_VERSION_PATCH - patch version (z) 15 | # 16 | # Hints 17 | # ^^^^^ 18 | # A user may set ``TensorRT_ROOT`` to an installation root to tell this module where to look. 19 | # 20 | set(_TensorRT_SEARCHES) 21 | 22 | if(TensorRT_ROOT) 23 | set(_TensorRT_SEARCH_ROOT PATHS ${TensorRT_ROOT} NO_DEFAULT_PATH) 24 | list(APPEND _TensorRT_SEARCHES _TensorRT_SEARCH_ROOT) 25 | endif() 26 | 27 | # appends some common paths 28 | set(_TensorRT_SEARCH_NORMAL 29 | PATHS "/usr" 30 | ) 31 | list(APPEND _TensorRT_SEARCHES _TensorRT_SEARCH_NORMAL) 32 | 33 | # Include dir 34 | foreach(search ${_TensorRT_SEARCHES}) 35 | find_path(TensorRT_INCLUDE_DIR NAMES NvInfer.h ${${search}} PATH_SUFFIXES include) 36 | endforeach() 37 | 38 | if(NOT TensorRT_LIBRARY) 39 | foreach(search ${_TensorRT_SEARCHES}) 40 | find_library(TensorRT_LIBRARY NAMES nvinfer ${${search}} PATH_SUFFIXES lib) 41 | if(NOT TensorRT_LIB_DIR) 42 | get_filename_component(TensorRT_LIB_DIR ${TensorRT_LIBRARY} DIRECTORY) 43 | endif () 44 | endforeach() 45 | endif() 46 | 47 | if(NOT TensorRT_nvinfer_plugin_LIBRARY) 48 | foreach(search ${_TensorRT_SEARCHES}) 49 | find_library(TensorRT_nvinfer_plugin_LIBRARY NAMES nvinfer_plugin ${${search}} PATH_SUFFIXES lib) 50 | endforeach() 51 | endif() 52 | 53 | mark_as_advanced(TensorRT_INCLUDE_DIR) 54 | 55 | if(TensorRT_INCLUDE_DIR AND EXISTS "${TensorRT_INCLUDE_DIR}/NvInfer.h") 56 | if(EXISTS "${TensorRT_INCLUDE_DIR}/NvInferVersion.h") 57 | set(_VersionSearchFile "${TensorRT_INCLUDE_DIR}/NvInferVersion.h") 58 | else () 59 | set(_VersionSearchFile "${TensorRT_INCLUDE_DIR}/NvInfer.h") 60 | endif () 61 | file(STRINGS "${_VersionSearchFile}" TensorRT_MAJOR REGEX "^#define NV_TENSORRT_MAJOR [0-9]+.*$") 62 | file(STRINGS "${_VersionSearchFile}" TensorRT_MINOR REGEX "^#define NV_TENSORRT_MINOR [0-9]+.*$") 63 | file(STRINGS "${_VersionSearchFile}" TensorRT_PATCH REGEX "^#define NV_TENSORRT_PATCH [0-9]+.*$") 64 | 65 | string(REGEX REPLACE "^#define NV_TENSORRT_MAJOR ([0-9]+).*$" "\\1" TensorRT_VERSION_MAJOR "${TensorRT_MAJOR}") 66 | string(REGEX REPLACE "^#define NV_TENSORRT_MINOR ([0-9]+).*$" "\\1" TensorRT_VERSION_MINOR "${TensorRT_MINOR}") 67 | string(REGEX REPLACE "^#define NV_TENSORRT_PATCH ([0-9]+).*$" "\\1" TensorRT_VERSION_PATCH "${TensorRT_PATCH}") 68 | set(TensorRT_VERSION_STRING "${TensorRT_VERSION_MAJOR}.${TensorRT_VERSION_MINOR}.${TensorRT_VERSION_PATCH}") 69 | endif() 70 | 71 | include(FindPackageHandleStandardArgs) 72 | FIND_PACKAGE_HANDLE_STANDARD_ARGS(TensorRT REQUIRED_VARS TensorRT_LIBRARY TensorRT_INCLUDE_DIR VERSION_VAR TensorRT_VERSION_STRING) 73 | 74 | if(TensorRT_FOUND) 75 | set(TensorRT_INCLUDE_DIRS ${TensorRT_INCLUDE_DIR}) 76 | 77 | if(NOT TensorRT_LIBRARIES) 78 | set(TensorRT_LIBRARIES ${TensorRT_LIBRARY}) 79 | if (TensorRT_nvinfer_plugin_LIBRARY) 80 | list(APPEND TensorRT_LIBRARIES ${TensorRT_nvinfer_plugin_LIBRARY}) 81 | endif() 82 | endif() 83 | 84 | if(NOT TARGET TensorRT::TensorRT) 85 | add_library(TensorRT INTERFACE IMPORTED) 86 | add_library(TensorRT::TensorRT ALIAS TensorRT) 87 | endif() 88 | 89 | if(NOT TARGET TensorRT::nvinfer) 90 | add_library(TensorRT::nvinfer SHARED IMPORTED) 91 | if (WIN32) 92 | foreach(search ${_TensorRT_SEARCHES}) 93 | find_file(TensorRT_LIBRARY_DLL 94 | NAMES nvinfer.dll 95 | PATHS ${${search}} 96 | PATH_SUFFIXES bin 97 | ) 98 | endforeach() 99 | 100 | set_target_properties(TensorRT::nvinfer PROPERTIES 101 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 102 | IMPORTED_LOCATION "${TensorRT_LIBRARY_DLL}" 103 | IMPORTED_IMPLIB "${TensorRT_LIBRARY}" 104 | ) 105 | else() 106 | set_target_properties(TensorRT::nvinfer PROPERTIES 107 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 108 | IMPORTED_LOCATION "${TensorRT_LIBRARY}" 109 | ) 110 | endif() 111 | target_link_libraries(TensorRT INTERFACE TensorRT::nvinfer) 112 | endif() 113 | 114 | if(NOT TARGET TensorRT::nvinfer_plugin AND TensorRT_nvinfer_plugin_LIBRARY) 115 | add_library(TensorRT::nvinfer_plugin SHARED IMPORTED) 116 | if (WIN32) 117 | foreach(search ${_TensorRT_SEARCHES}) 118 | find_file(TensorRT_nvinfer_plugin_LIBRARY_DLL 119 | NAMES nvinfer_plugin.dll 120 | PATHS ${${search}} 121 | PATH_SUFFIXES bin 122 | ) 123 | endforeach() 124 | 125 | set_target_properties(TensorRT::nvinfer_plugin PROPERTIES 126 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 127 | IMPORTED_LOCATION "${TensorRT_nvinfer_plugin_LIBRARY_DLL}" 128 | IMPORTED_IMPLIB "${TensorRT_nvinfer_plugin_LIBRARY}" 129 | ) 130 | else() 131 | set_target_properties(TensorRT::nvinfer_plugin PROPERTIES 132 | INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}" 133 | IMPORTED_LOCATION "${TensorRT_nvinfer_plugin_LIBRARY}" 134 | ) 135 | endif() 136 | target_link_libraries(TensorRT INTERFACE TensorRT::nvinfer_plugin) 137 | endif() 138 | endif() 139 | -------------------------------------------------------------------------------- /csrc/segment/simple/cmake/Function.cmake: -------------------------------------------------------------------------------- 1 | 2 | function(print_var var) 3 | set(value "${${var}}") 4 | string(LENGTH "${value}" value_length) 5 | if(value_length GREATER 0) 6 | math(EXPR last_index "${value_length} - 1") 7 | string(SUBSTRING "${value}" ${last_index} ${last_index} last_char) 8 | endif() 9 | 10 | if(NOT "${last_char}" STREQUAL "\n") 11 | set(value "${value}\n") 12 | endif() 13 | message(STATUS "${var}:\n ${value}") 14 | endfunction() 15 | -------------------------------------------------------------------------------- /csrc/segment/simple/include/common.hpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 2/9/23. 3 | // 4 | 5 | #ifndef SEGMENT_SIMPLE_COMMON_HPP 6 | #define SEGMENT_SIMPLE_COMMON_HPP 7 | #include "NvInfer.h" 8 | #include "opencv2/opencv.hpp" 9 | #include 10 | #include 11 | 12 | #define CHECK(call) \ 13 | do { \ 14 | const cudaError_t error_code = call; \ 15 | if (error_code != cudaSuccess) { \ 16 | printf("CUDA Error:\n"); \ 17 | printf(" File: %s\n", __FILE__); \ 18 | printf(" Line: %d\n", __LINE__); \ 19 | printf(" Error code: %d\n", error_code); \ 20 | printf(" Error text: %s\n", cudaGetErrorString(error_code)); \ 21 | exit(1); \ 22 | } \ 23 | } while (0) 24 | 25 | class Logger: public nvinfer1::ILogger { 26 | public: 27 | nvinfer1::ILogger::Severity reportableSeverity; 28 | 29 | explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO): 30 | reportableSeverity(severity) 31 | { 32 | } 33 | 34 | void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override 35 | { 36 | if (severity > reportableSeverity) { 37 | return; 38 | } 39 | switch (severity) { 40 | case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: 41 | std::cerr << "INTERNAL_ERROR: "; 42 | break; 43 | case nvinfer1::ILogger::Severity::kERROR: 44 | std::cerr << "ERROR: "; 45 | break; 46 | case nvinfer1::ILogger::Severity::kWARNING: 47 | std::cerr << "WARNING: "; 48 | break; 49 | case nvinfer1::ILogger::Severity::kINFO: 50 | std::cerr << "INFO: "; 51 | break; 52 | default: 53 | std::cerr << "VERBOSE: "; 54 | break; 55 | } 56 | std::cerr << msg << std::endl; 57 | } 58 | }; 59 | 60 | inline int get_size_by_dims(const nvinfer1::Dims& dims) 61 | { 62 | int size = 1; 63 | for (int i = 0; i < dims.nbDims; i++) { 64 | size *= dims.d[i]; 65 | } 66 | return size; 67 | } 68 | 69 | inline int type_to_size(const nvinfer1::DataType& dataType) 70 | { 71 | switch (dataType) { 72 | case nvinfer1::DataType::kFLOAT: 73 | return 4; 74 | case nvinfer1::DataType::kHALF: 75 | return 2; 76 | case nvinfer1::DataType::kINT32: 77 | return 4; 78 | case nvinfer1::DataType::kINT8: 79 | return 1; 80 | case nvinfer1::DataType::kBOOL: 81 | return 1; 82 | default: 83 | return 4; 84 | } 85 | } 86 | 87 | inline static float clamp(float val, float min, float max) 88 | { 89 | return val > min ? (val < max ? val : max) : min; 90 | } 91 | 92 | inline bool IsPathExist(const std::string& path) 93 | { 94 | if (access(path.c_str(), 0) == F_OK) { 95 | return true; 96 | } 97 | return false; 98 | } 99 | 100 | inline bool IsFile(const std::string& path) 101 | { 102 | if (!IsPathExist(path)) { 103 | printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str()); 104 | return false; 105 | } 106 | struct stat buffer; 107 | return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode)); 108 | } 109 | 110 | inline bool IsFolder(const std::string& path) 111 | { 112 | if (!IsPathExist(path)) { 113 | return false; 114 | } 115 | struct stat buffer; 116 | return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode)); 117 | } 118 | 119 | namespace seg { 120 | struct Binding { 121 | size_t size = 1; 122 | size_t dsize = 1; 123 | nvinfer1::Dims dims; 124 | std::string name; 125 | }; 126 | 127 | struct Object { 128 | cv::Rect_ rect; 129 | int label = 0; 130 | float prob = 0.0; 131 | cv::Mat boxMask; 132 | }; 133 | 134 | struct PreParam { 135 | float ratio = 1.0f; 136 | float dw = 0.0f; 137 | float dh = 0.0f; 138 | float height = 0; 139 | float width = 0; 140 | }; 141 | } // namespace seg 142 | #endif // SEGMENT_SIMPLE_COMMON_HPP 143 | -------------------------------------------------------------------------------- /data/bus.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/triple-Mu/YOLOv8-TensorRT/d5f68cd79a589dd6875ccc070c31fbfc2f38227d/data/bus.jpg -------------------------------------------------------------------------------- /data/zidane.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/triple-Mu/YOLOv8-TensorRT/d5f68cd79a589dd6875ccc070c31fbfc2f38227d/data/zidane.jpg -------------------------------------------------------------------------------- /docs/API-Build.md: -------------------------------------------------------------------------------- 1 | # Build TensorRT Engine By TensorRT Python API 2 | 3 | When you want to build engine by API. You should generate the pickle weights parameters first. 4 | 5 | ``` shell 6 | python3 gen_pkl.py -w yolov8s.pt -o yolov8s.pkl 7 | ``` 8 | 9 | You will get a `yolov8s.pkl` which contain the operators' parameters. 10 | 11 | And you can rebuild `yolov8s` model in TensorRT api. 12 | 13 | ``` 14 | python3 build.py \ 15 | --weights yolov8s.pkl \ 16 | --iou-thres 0.65 \ 17 | --conf-thres 0.25 \ 18 | --topk 100 \ 19 | --fp16 \ 20 | --input-shape 1 3 640 640 \ 21 | --device cuda:0 22 | ``` 23 | 24 | ***Notice !!!*** 25 | 26 | Now we only support static input shape model build by TensorRT api. 27 | You'd best give the legal `input-shape`. 28 | 29 | ***Notice !!!*** 30 | 31 | Now we don't support YOLOv8-seg model building by API. It will be supported later. 32 | -------------------------------------------------------------------------------- /docs/Cls.md: -------------------------------------------------------------------------------- 1 | # YOLOv8-cls Model with TensorRT 2 | 3 | The yolov8-cls model conversion route is : 4 | YOLOv8 PyTorch model -> ONNX -> TensorRT Engine 5 | 6 | ***Notice !!!*** We don't support TensorRT API building !!! 7 | 8 | # Export Orin ONNX model by ultralytics 9 | 10 | You can leave this repo and use the original `ultralytics` repo for onnx export. 11 | 12 | ### 1. ONNX -> TensorRT 13 | 14 | You can export your onnx model by `ultralytics` API. 15 | 16 | ``` shell 17 | yolo export model=yolov8s-cls.pt format=onnx opset=11 simplify=True 18 | ``` 19 | 20 | or run this python script: 21 | 22 | ```python 23 | from ultralytics import YOLO 24 | 25 | # Load a model 26 | model = YOLO("yolov8s-cls.pt") # load a pretrained model (recommended for training) 27 | success = model.export(format="onnx", opset=11, simplify=True) # export the model to onnx format 28 | assert success 29 | ``` 30 | 31 | Then build engine by Trtexec Tools. 32 | 33 | You can export TensorRT engine by [`trtexec`](https://github.com/NVIDIA/TensorRT/tree/main/samples/trtexec) tools. 34 | 35 | Usage: 36 | 37 | ``` shell 38 | /usr/src/tensorrt/bin/trtexec \ 39 | --onnx=yolov8s-cls.onnx \ 40 | --saveEngine=yolov8s-cls.engine \ 41 | --fp16 42 | ``` 43 | 44 | ### 2. Direct to TensorRT (NOT RECOMMAND!!) 45 | 46 | Usage: 47 | 48 | ```shell 49 | yolo export model=yolov8s-cls.pt format=engine device=0 50 | ``` 51 | 52 | or run python script: 53 | 54 | ```python 55 | from ultralytics import YOLO 56 | 57 | # Load a model 58 | model = YOLO("yolov8s-cls.pt") # load a pretrained model (recommended for training) 59 | success = model.export(format="engine", device=0) # export the model to engine format 60 | assert success 61 | ``` 62 | 63 | After executing the above script, you will get an engine named `yolov8s-cls.engine` . 64 | 65 | # Inference 66 | 67 | ## Infer with python script 68 | 69 | You can infer images with the engine by [`infer-cls.py`](../infer-cls.py) . 70 | 71 | Usage: 72 | 73 | ``` shell 74 | python3 infer-cls.py \ 75 | --engine yolov8s-cls.engine \ 76 | --imgs data \ 77 | --show \ 78 | --out-dir outputs \ 79 | --device cuda:0 80 | ``` 81 | 82 | #### Description of all arguments 83 | 84 | - `--engine` : The Engine you export. 85 | - `--imgs` : The images path you want to detect. 86 | - `--show` : Whether to show detection results. 87 | - `--out-dir` : Where to save detection results images. It will not work when use `--show` flag. 88 | - `--device` : The CUDA deivce you use. 89 | 90 | ## Inference with c++ 91 | 92 | You can infer with c++ in [`csrc/cls/normal`](../csrc/cls/normal) . 93 | 94 | ### Build: 95 | 96 | Please set you own librarys in [`CMakeLists.txt`](../csrc/cls/normal/CMakeLists.txt) and modify `CLASS_NAMES` in [`main.cpp`](../csrc/cls/normal/main.cpp). 97 | 98 | And build: 99 | 100 | ``` shell 101 | export root=${PWD} 102 | cd src/cls/normal 103 | mkdir build 104 | cmake .. 105 | make 106 | mv yolov8-cls ${root} 107 | cd ${root} 108 | ``` 109 | 110 | Usage: 111 | 112 | ``` shell 113 | # infer image 114 | ./yolov8-cls yolov8s-cls.engine data/bus.jpg 115 | # infer images 116 | ./yolov8-cls yolov8s-cls.engine data 117 | # infer video 118 | ./yolov8-cls yolov8s-cls.engine data/test.mp4 # the video path 119 | ``` 120 | -------------------------------------------------------------------------------- /docs/Jetson.md: -------------------------------------------------------------------------------- 1 | # YOLOv8 on Jetson 2 | 3 | Only test on `Jetson-NX 4GB` 4 | 5 | ENVS: 6 | 7 | - Jetpack 4.6.3 8 | - CUDA-10.2 9 | - CUDNN-8.2.1 10 | - TensorRT-8.2.1 11 | - DeepStream-6.0.1 12 | - OpenCV-4.1.1 13 | - CMake-3.10.2 14 | 15 | If you have other environment-related issues, please discuss in issue. 16 | 17 | ## End2End Detection 18 | 19 | ### 1. Export Detection End2End ONNX 20 | 21 | `yolov8s.pt` is your trained pytorch model, or the official pre-trained model. 22 | 23 | Do not use any model other than pytorch model. 24 | Do not use [`build.py`](../build.py) to export engine if you don't know how to install pytorch and other environments on 25 | jetson. 26 | 27 | ***!!! Please use the PC to execute the following script !!!*** 28 | 29 | ```shell 30 | # Export yolov8s.pt to yolov8s.onnx 31 | python3 export-det.py --weights yolov8s.pt --sim 32 | ``` 33 | 34 | ***!!! Please use the Jetson to execute the following script !!!*** 35 | 36 | ```shell 37 | # Using trtexec tools for export engine 38 | /usr/src/tensorrt/bin/trtexec \ 39 | --onnx=yolov8s.onnx \ 40 | --saveEngine=yolov8s.engine 41 | ``` 42 | 43 | After executing the above command, you will get an engine named `yolov8s.engine` . 44 | 45 | ### 2. Inference with c++ 46 | 47 | It is highly recommended to use C++ inference on Jetson. 48 | Here is a demo: [`csrc/jetson/detect`](../csrc/jetson/detect) . 49 | 50 | #### Build: 51 | 52 | Please modify `CLASS_NAMES` and `COLORS` in [`main.cpp`](../csrc/jetson/detect/main.cpp) for yourself. 53 | 54 | And build: 55 | 56 | ``` shell 57 | export root=${PWD} 58 | cd src/jetson/detect 59 | mkdir build 60 | cmake .. 61 | make 62 | mv yolov8 ${root} 63 | cd ${root} 64 | ``` 65 | 66 | Usage: 67 | 68 | ``` shell 69 | # infer image 70 | ./yolov8 yolov8s.engine data/bus.jpg 71 | # infer images 72 | ./yolov8 yolov8s.engine data 73 | # infer video 74 | ./yolov8 yolov8s.engine data/test.mp4 # the video path 75 | ``` 76 | 77 | ## Speedup Segmention 78 | 79 | ### 1. Export Segmention Speedup ONNX 80 | 81 | `yolov8s-seg.pt` is your trained pytorch model, or the official pre-trained model. 82 | 83 | Do not use any model other than pytorch model. 84 | Do not use [`build.py`](../build.py) to export engine if you don't know how to install pytorch and other environments on 85 | jetson. 86 | 87 | ***!!! Please use the PC to execute the following script !!!*** 88 | 89 | ```shell 90 | # Export yolov8s-seg.pt to yolov8s-seg.onnx 91 | python3 export-seg.py --weights yolov8s-seg.pt --sim 92 | ``` 93 | 94 | ***!!! Please use the Jetson to execute the following script !!!*** 95 | 96 | ```shell 97 | # Using trtexec tools for export engine 98 | /usr/src/tensorrt/bin/trtexec \ 99 | --onnx=yolov8s-seg.onnx \ 100 | --saveEngine=yolov8s-seg.engine 101 | ``` 102 | 103 | After executing the above command, you will get an engine named `yolov8s-seg.engine` . 104 | 105 | ### 2. Inference with c++ 106 | 107 | It is highly recommended to use C++ inference on Jetson. 108 | Here is a demo: [`csrc/jetson/segment`](../csrc/jetson/segment) . 109 | 110 | #### Build: 111 | 112 | Please modify `CLASS_NAMES` and `COLORS` and postprocess parameters in [`main.cpp`](../csrc/jetson/segment/main.cpp) for 113 | yourself. 114 | 115 | ```c++ 116 | int topk = 100; 117 | int seg_h = 160; // yolov8 model proto height 118 | int seg_w = 160; // yolov8 model proto width 119 | int seg_channels = 32; // yolov8 model proto channels 120 | float score_thres = 0.25f; 121 | float iou_thres = 0.65f; 122 | ``` 123 | 124 | And build: 125 | 126 | ``` shell 127 | export root=${PWD} 128 | cd src/jetson/segment 129 | mkdir build 130 | cmake .. 131 | make 132 | mv yolov8-seg ${root} 133 | cd ${root} 134 | ``` 135 | 136 | Usage: 137 | 138 | ``` shell 139 | # infer image 140 | ./yolov8-seg yolov8s-seg.engine data/bus.jpg 141 | # infer images 142 | ./yolov8-seg yolov8s-seg.engine data 143 | # infer video 144 | ./yolov8-seg yolov8s-seg.engine data/test.mp4 # the video path 145 | ``` 146 | 147 | ## Normal Posture 148 | 149 | ### 1. Export Posture Normal ONNX 150 | 151 | `yolov8s-pose.pt` is your trained pytorch model, or the official pre-trained model. 152 | 153 | Do not use any model other than pytorch model. 154 | Do not use [`build.py`](../build.py) to export engine if you don't know how to install pytorch and other environments on 155 | jetson. 156 | 157 | ***!!! Please use the PC to execute the following script !!!*** 158 | 159 | ```shell 160 | # Export yolov8s-pose.pt to yolov8s-pose.onnx 161 | yolo export model=yolov8s-pose.pt format=onnx simplify=True 162 | ``` 163 | 164 | ***!!! Please use the Jetson to execute the following script !!!*** 165 | 166 | ```shell 167 | # Using trtexec tools for export engine 168 | /usr/src/tensorrt/bin/trtexec \ 169 | --onnx=yolov8s-pose.onnx \ 170 | --saveEngine=yolov8s-pose.engine 171 | ``` 172 | 173 | After executing the above command, you will get an engine named `yolov8s-pose.engine` . 174 | 175 | ### 2. Inference with c++ 176 | 177 | It is highly recommended to use C++ inference on Jetson. 178 | Here is a demo: [`csrc/jetson/pose`](../csrc/jetson/pose) . 179 | 180 | #### Build: 181 | 182 | Please modify `KPS_COLORS` and `SKELETON` and `LIMB_COLORS` and postprocess parameters 183 | in [`main.cpp`](../csrc/jetson/pose/main.cpp) for yourself. 184 | 185 | ```c++ 186 | int topk = 100; 187 | float score_thres = 0.25f; 188 | float iou_thres = 0.65f; 189 | ``` 190 | 191 | And build: 192 | 193 | ``` shell 194 | export root=${PWD} 195 | cd src/jetson/pose 196 | mkdir build 197 | cmake .. 198 | make 199 | mv yolov8-pose ${root} 200 | cd ${root} 201 | ``` 202 | 203 | Usage: 204 | 205 | ``` shell 206 | # infer image 207 | ./yolov8-pose yolov8s-pose.engine data/bus.jpg 208 | # infer images 209 | ./yolov8-pose yolov8s-pose.engine data 210 | # infer video 211 | ./yolov8-pose yolov8s-pose.engine data/test.mp4 # the video path 212 | ``` 213 | -------------------------------------------------------------------------------- /docs/Normal.md: -------------------------------------------------------------------------------- 1 | # Normal Usage of [`ultralytics`](https://github.com/ultralytics/ultralytics) 2 | 3 | ## Export TensorRT Engine 4 | 5 | ### 1. ONNX -> TensorRT 6 | 7 | You can export your onnx model by `ultralytics` API. 8 | 9 | ``` shell 10 | yolo export model=yolov8s.pt format=onnx opset=11 simplify=True 11 | ``` 12 | 13 | or run this python script: 14 | 15 | ```python 16 | from ultralytics import YOLO 17 | 18 | # Load a model 19 | model = YOLO("yolov8s.pt") # load a pretrained model (recommended for training) 20 | success = model.export(format="onnx", opset=11, simplify=True) # export the model to onnx format 21 | assert success 22 | ``` 23 | 24 | Then build engine by Trtexec Tools. 25 | 26 | You can export TensorRT engine by [`trtexec`](https://github.com/NVIDIA/TensorRT/tree/main/samples/trtexec) tools. 27 | 28 | Usage: 29 | 30 | ``` shell 31 | /usr/src/tensorrt/bin/trtexec \ 32 | --onnx=yolov8s.onnx \ 33 | --saveEngine=yolov8s.engine \ 34 | --fp16 35 | ``` 36 | 37 | ### 2. Direct to TensorRT (NOT RECOMMAND!!) 38 | 39 | Usage: 40 | 41 | ```shell 42 | yolo export model=yolov8s.pt format=engine device=0 43 | ``` 44 | 45 | or run python script: 46 | 47 | ```python 48 | 49 | from ultralytics import YOLO 50 | 51 | # Load a model 52 | model = YOLO("yolov8s.pt") # load a pretrained model (recommended for training) 53 | success = model.export(format="engine", device=0) # export the model to engine format 54 | assert success 55 | ``` 56 | 57 | After executing the above script, you will get an engine named `yolov8s.engine` . 58 | 59 | ## Inference with c++ 60 | 61 | You can infer with c++ in [`csrc/detect/normal`](../csrc/detect/normal) . 62 | 63 | ### Build: 64 | 65 | Please set you own librarys in [`CMakeLists.txt`](../csrc/detect/normal/CMakeLists.txt) and modify `CLASS_NAMES` 66 | and `COLORS` in [`main.cpp`](../csrc/detect/normal/main.cpp). 67 | 68 | Besides, you can modify the postprocess parameters such as `num_labels` and `score_thres` and `iou_thres` and `topk` 69 | in [`main.cpp`](../csrc/detect/normal/main.cpp). 70 | 71 | ```c++ 72 | int num_labels = 80; 73 | int topk = 100; 74 | float score_thres = 0.25f; 75 | float iou_thres = 0.65f; 76 | ``` 77 | 78 | And build: 79 | 80 | ``` shell 81 | export root=${PWD} 82 | cd src/detect/normal 83 | mkdir build && cd build 84 | cmake .. 85 | make 86 | mv yolov8 ${root} 87 | cd ${root} 88 | ``` 89 | 90 | Usage: 91 | 92 | ``` shell 93 | # infer image 94 | ./yolov8 yolov8s.engine data/bus.jpg 95 | # infer images 96 | ./yolov8 yolov8s.engine data 97 | # infer video 98 | ./yolov8 yolov8s.engine data/test.mp4 # the video path 99 | ``` 100 | -------------------------------------------------------------------------------- /docs/Obb.md: -------------------------------------------------------------------------------- 1 | # YOLOv8-obb Model with TensorRT 2 | 3 | The yolov8-obb model conversion route is : 4 | YOLOv8 PyTorch model -> ONNX -> TensorRT Engine 5 | 6 | ***Notice !!!*** We don't support TensorRT API building !!! 7 | 8 | # Export Orin ONNX model by ultralytics 9 | 10 | You can leave this repo and use the original `ultralytics` repo for onnx export. 11 | 12 | ### 1. ONNX -> TensorRT 13 | 14 | You can export your onnx model by `ultralytics` API. 15 | 16 | ``` shell 17 | yolo export model=yolov8s-obb.pt format=onnx opset=11 simplify=True 18 | ``` 19 | 20 | or run this python script: 21 | 22 | ```python 23 | from ultralytics import YOLO 24 | 25 | # Load a model 26 | model = YOLO("yolov8s-obb.pt") # load a pretrained model (recommended for training) 27 | success = model.export(format="onnx", opset=11, simplify=True) # export the model to onnx format 28 | assert success 29 | ``` 30 | 31 | Then build engine by Trtexec Tools. 32 | 33 | You can export TensorRT engine by [`trtexec`](https://github.com/NVIDIA/TensorRT/tree/main/samples/trtexec) tools. 34 | 35 | Usage: 36 | 37 | ``` shell 38 | /usr/src/tensorrt/bin/trtexec \ 39 | --onnx=yolov8s-obb.onnx \ 40 | --saveEngine=yolov8s-obb.engine \ 41 | --fp16 42 | ``` 43 | 44 | ### 2. Direct to TensorRT (NOT RECOMMAND!!) 45 | 46 | Usage: 47 | 48 | ```shell 49 | yolo export model=yolov8s-obb.pt format=engine device=0 50 | ``` 51 | 52 | or run python script: 53 | 54 | ```python 55 | from ultralytics import YOLO 56 | 57 | # Load a model 58 | model = YOLO("yolov8s-obb.pt") # load a pretrained model (recommended for training) 59 | success = model.export(format="engine", device=0) # export the model to engine format 60 | assert success 61 | ``` 62 | 63 | After executing the above script, you will get an engine named `yolov8s-obb.engine` . 64 | 65 | # Inference 66 | 67 | ## Infer with python script 68 | 69 | You can infer images with the engine by [`infer-obb.py`](../infer-obb.py) . 70 | 71 | Usage: 72 | 73 | ``` shell 74 | python3 infer-obb.py \ 75 | --engine yolov8s-obb.engine \ 76 | --imgs data \ 77 | --show \ 78 | --out-dir outputs \ 79 | --device cuda:0 80 | ``` 81 | 82 | #### Description of all arguments 83 | 84 | - `--engine` : The Engine you export. 85 | - `--imgs` : The images path you want to detect. 86 | - `--show` : Whether to show detection results. 87 | - `--out-dir` : Where to save detection results images. It will not work when use `--show` flag. 88 | - `--device` : The CUDA deivce you use. 89 | 90 | ## Inference with c++ 91 | 92 | You can infer with c++ in [`csrc/obb/normal`](../csrc/obb/normal) . 93 | 94 | ### Build: 95 | 96 | Please set you own librarys in [`CMakeLists.txt`](../csrc/obb/normal/CMakeLists.txt) and modify `CLASS_NAMES` 97 | and `COLORS` in [`main.cpp`](../csrc/obb/normal/main.cpp). 98 | 99 | Besides, you can modify the postprocess parameters such as `num_labels` and `score_thres` and `iou_thres` and `topk` 100 | 101 | And build: 102 | 103 | ``` shell 104 | export root=${PWD} 105 | cd src/obb/normal 106 | mkdir build 107 | cmake .. 108 | make 109 | mv yolov8-obb ${root} 110 | cd ${root} 111 | ``` 112 | 113 | Usage: 114 | 115 | ``` shell 116 | # infer image 117 | ./yolov8-obb yolov8s-obb.engine data/bus.jpg 118 | # infer images 119 | ./yolov8-obb yolov8s-obb.engine data 120 | # infer video 121 | ./yolov8-obb yolov8s-obb.engine data/test.mp4 # the video path 122 | ``` 123 | -------------------------------------------------------------------------------- /docs/Pose.md: -------------------------------------------------------------------------------- 1 | # YOLOv8-pose Model with TensorRT 2 | 3 | The yolov8-pose model conversion route is : 4 | YOLOv8 PyTorch model -> ONNX -> TensorRT Engine 5 | 6 | ***Notice !!!*** We don't support TensorRT API building !!! 7 | 8 | # Export Orin ONNX model by ultralytics 9 | 10 | You can leave this repo and use the original `ultralytics` repo for onnx export. 11 | 12 | ### 1. ONNX -> TensorRT 13 | 14 | You can export your onnx model by `ultralytics` API. 15 | 16 | ``` shell 17 | yolo export model=yolov8s-pose.pt format=onnx opset=11 simplify=True 18 | ``` 19 | 20 | or run this python script: 21 | 22 | ```python 23 | from ultralytics import YOLO 24 | 25 | # Load a model 26 | model = YOLO("yolov8s-pose.pt") # load a pretrained model (recommended for training) 27 | success = model.export(format="onnx", opset=11, simplify=True) # export the model to onnx format 28 | assert success 29 | ``` 30 | 31 | Then build engine by Trtexec Tools. 32 | 33 | You can export TensorRT engine by [`trtexec`](https://github.com/NVIDIA/TensorRT/tree/main/samples/trtexec) tools. 34 | 35 | Usage: 36 | 37 | ``` shell 38 | /usr/src/tensorrt/bin/trtexec \ 39 | --onnx=yolov8s-pose.onnx \ 40 | --saveEngine=yolov8s-pose.engine \ 41 | --fp16 42 | ``` 43 | 44 | ### 2. Direct to TensorRT (NOT RECOMMAND!!) 45 | 46 | Usage: 47 | 48 | ```shell 49 | yolo export model=yolov8s-pose.pt format=engine device=0 50 | ``` 51 | 52 | or run python script: 53 | 54 | ```python 55 | from ultralytics import YOLO 56 | 57 | # Load a model 58 | model = YOLO("yolov8s-pose.pt") # load a pretrained model (recommended for training) 59 | success = model.export(format="engine", device=0) # export the model to engine format 60 | assert success 61 | ``` 62 | 63 | After executing the above script, you will get an engine named `yolov8s-pose.engine` . 64 | 65 | # Inference 66 | 67 | ## Infer with python script 68 | 69 | You can infer images with the engine by [`infer-pose.py`](../infer-pose.py) . 70 | 71 | Usage: 72 | 73 | ``` shell 74 | python3 infer-pose.py \ 75 | --engine yolov8s-pose.engine \ 76 | --imgs data \ 77 | --show \ 78 | --out-dir outputs \ 79 | --device cuda:0 80 | ``` 81 | 82 | #### Description of all arguments 83 | 84 | - `--engine` : The Engine you export. 85 | - `--imgs` : The images path you want to detect. 86 | - `--show` : Whether to show detection results. 87 | - `--out-dir` : Where to save detection results images. It will not work when use `--show` flag. 88 | - `--device` : The CUDA deivce you use. 89 | 90 | ## Inference with c++ 91 | 92 | You can infer with c++ in [`csrc/pose/normal`](../csrc/pose/normal) . 93 | 94 | ### Build: 95 | 96 | Please set you own librarys in [`CMakeLists.txt`](../csrc/pose/normal/CMakeLists.txt) and modify `KPS_COLORS` 97 | and `SKELETON` and `LIMB_COLORS` in [`main.cpp`](../csrc/pose/normal/main.cpp). 98 | 99 | Besides, you can modify the postprocess parameters such as `score_thres` and `iou_thres` and `topk` 100 | in [`main.cpp`](../csrc/pose/normal/main.cpp). 101 | 102 | ```c++ 103 | int topk = 100; 104 | float score_thres = 0.25f; 105 | float iou_thres = 0.65f; 106 | ``` 107 | 108 | And build: 109 | 110 | ``` shell 111 | export root=${PWD} 112 | cd src/pose/normal 113 | mkdir build 114 | cmake .. 115 | make 116 | mv yolov8-pose ${root} 117 | cd ${root} 118 | ``` 119 | 120 | Usage: 121 | 122 | ``` shell 123 | # infer image 124 | ./yolov8-pose yolov8s-pose.engine data/bus.jpg 125 | # infer images 126 | ./yolov8-pose yolov8s-pose.engine data 127 | # infer video 128 | ./yolov8-pose yolov8s-pose.engine data/test.mp4 # the video path 129 | ``` 130 | -------------------------------------------------------------------------------- /docs/star.md: -------------------------------------------------------------------------------- 1 | ## Star History 2 | 3 | [![YOLOv8-TensorRT](https://api.star-history.com/svg?repos=triple-Mu/YOLOv8-TensorRT&type=Date)](https://star-history.com/#triple-Mu/YOLOv8-TensorRT&Date) 4 | -------------------------------------------------------------------------------- /export-det.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from io import BytesIO 3 | 4 | import onnx 5 | import torch 6 | from ultralytics import YOLO 7 | 8 | from models.common import PostDetect, optim 9 | 10 | try: 11 | import onnxsim 12 | except ImportError: 13 | onnxsim = None 14 | 15 | 16 | def parse_args(): 17 | parser = argparse.ArgumentParser() 18 | parser.add_argument('-w', 19 | '--weights', 20 | type=str, 21 | required=True, 22 | help='PyTorch yolov8 weights') 23 | parser.add_argument('--iou-thres', 24 | type=float, 25 | default=0.65, 26 | help='IOU threshoud for NMS plugin') 27 | parser.add_argument('--conf-thres', 28 | type=float, 29 | default=0.25, 30 | help='CONF threshoud for NMS plugin') 31 | parser.add_argument('--topk', 32 | type=int, 33 | default=100, 34 | help='Max number of detection bboxes') 35 | parser.add_argument('--opset', 36 | type=int, 37 | default=11, 38 | help='ONNX opset version') 39 | parser.add_argument('--sim', 40 | action='store_true', 41 | help='simplify onnx model') 42 | parser.add_argument('--input-shape', 43 | nargs='+', 44 | type=int, 45 | default=[1, 3, 640, 640], 46 | help='Model input shape only for api builder') 47 | parser.add_argument('--device', 48 | type=str, 49 | default='cpu', 50 | help='Export ONNX device') 51 | args = parser.parse_args() 52 | assert len(args.input_shape) == 4 53 | PostDetect.conf_thres = args.conf_thres 54 | PostDetect.iou_thres = args.iou_thres 55 | PostDetect.topk = args.topk 56 | return args 57 | 58 | 59 | def main(args): 60 | b = args.input_shape[0] 61 | YOLOv8 = YOLO(args.weights) 62 | model = YOLOv8.model.fuse().eval() 63 | for m in model.modules(): 64 | optim(m) 65 | m.to(args.device) 66 | model.to(args.device) 67 | fake_input = torch.randn(args.input_shape).to(args.device) 68 | for _ in range(2): 69 | model(fake_input) 70 | save_path = args.weights.replace('.pt', '.onnx') 71 | with BytesIO() as f: 72 | torch.onnx.export( 73 | model, 74 | fake_input, 75 | f, 76 | opset_version=args.opset, 77 | input_names=['images'], 78 | output_names=['num_dets', 'bboxes', 'scores', 'labels']) 79 | f.seek(0) 80 | onnx_model = onnx.load(f) 81 | onnx.checker.check_model(onnx_model) 82 | shapes = [b, 1, b, args.topk, 4, b, args.topk, b, args.topk] 83 | for i in onnx_model.graph.output: 84 | for j in i.type.tensor_type.shape.dim: 85 | j.dim_param = str(shapes.pop(0)) 86 | if args.sim: 87 | try: 88 | onnx_model, check = onnxsim.simplify(onnx_model) 89 | assert check, 'assert check failed' 90 | except Exception as e: 91 | print(f'Simplifier failure: {e}') 92 | onnx.save(onnx_model, save_path) 93 | print(f'ONNX export success, saved as {save_path}') 94 | 95 | 96 | if __name__ == '__main__': 97 | main(parse_args()) 98 | -------------------------------------------------------------------------------- /export-seg.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from io import BytesIO 3 | 4 | import onnx 5 | import torch 6 | from ultralytics import YOLO 7 | 8 | from models.common import optim 9 | 10 | try: 11 | import onnxsim 12 | except ImportError: 13 | onnxsim = None 14 | 15 | 16 | def parse_args(): 17 | parser = argparse.ArgumentParser() 18 | parser.add_argument('-w', 19 | '--weights', 20 | type=str, 21 | required=True, 22 | help='PyTorch yolov8 weights') 23 | parser.add_argument('--opset', 24 | type=int, 25 | default=11, 26 | help='ONNX opset version') 27 | parser.add_argument('--sim', 28 | action='store_true', 29 | help='simplify onnx model') 30 | parser.add_argument('--input-shape', 31 | nargs='+', 32 | type=int, 33 | default=[1, 3, 640, 640], 34 | help='Model input shape only for api builder') 35 | parser.add_argument('--device', 36 | type=str, 37 | default='cpu', 38 | help='Export ONNX device') 39 | args = parser.parse_args() 40 | assert len(args.input_shape) == 4 41 | return args 42 | 43 | 44 | def main(args): 45 | YOLOv8 = YOLO(args.weights) 46 | model = YOLOv8.model.fuse().eval() 47 | for m in model.modules(): 48 | optim(m) 49 | m.to(args.device) 50 | model.to(args.device) 51 | fake_input = torch.randn(args.input_shape).to(args.device) 52 | for _ in range(2): 53 | model(fake_input) 54 | save_path = args.weights.replace('.pt', '.onnx') 55 | with BytesIO() as f: 56 | torch.onnx.export(model, 57 | fake_input, 58 | f, 59 | opset_version=args.opset, 60 | input_names=['images'], 61 | output_names=['outputs', 'proto']) 62 | f.seek(0) 63 | onnx_model = onnx.load(f) 64 | onnx.checker.check_model(onnx_model) 65 | if args.sim: 66 | try: 67 | onnx_model, check = onnxsim.simplify(onnx_model) 68 | assert check, 'assert check failed' 69 | except Exception as e: 70 | print(f'Simplifier failure: {e}') 71 | onnx.save(onnx_model, save_path) 72 | print(f'ONNX export success, saved as {save_path}') 73 | 74 | 75 | if __name__ == '__main__': 76 | main(parse_args()) 77 | -------------------------------------------------------------------------------- /gen_pkl.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import pickle 3 | from collections import OrderedDict 4 | 5 | from ultralytics import YOLO 6 | 7 | 8 | def parse_args(): 9 | parser = argparse.ArgumentParser() 10 | parser.add_argument('-w', 11 | '--weights', 12 | type=str, 13 | required=True, 14 | help='YOLOv8 pytorch weights') 15 | parser.add_argument('-o', 16 | '--output', 17 | type=str, 18 | required=True, 19 | help='Output file') 20 | args = parser.parse_args() 21 | return args 22 | 23 | 24 | args = parse_args() 25 | 26 | model = YOLO(args.weights) 27 | model.model.fuse() 28 | YOLOv8 = model.model.model 29 | 30 | strides = YOLOv8[-1].stride.detach().cpu().numpy() 31 | reg_max = YOLOv8[-1].dfl.conv.weight.shape[1] 32 | 33 | state_dict = OrderedDict(GD=model.model.yaml['depth_multiple'], 34 | GW=model.model.yaml['width_multiple'], 35 | strides=strides, 36 | reg_max=reg_max) 37 | 38 | for name, value in YOLOv8.state_dict().items(): 39 | value = value.detach().cpu().numpy() 40 | i = int(name.split('.')[0]) 41 | layer = YOLOv8[i] 42 | module_name = layer.type.split('.')[-1] 43 | stem = module_name + '.' + name 44 | state_dict[stem] = value 45 | 46 | with open(args.output, 'wb') as f: 47 | pickle.dump(state_dict, f) 48 | -------------------------------------------------------------------------------- /infer-cls-without-torch.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from pathlib import Path 3 | 4 | import cv2 5 | import numpy as np 6 | 7 | from config import CLASSES_CLS 8 | from models.utils import blob, path_to_list 9 | 10 | 11 | def main(args: argparse.Namespace) -> None: 12 | if args.method == 'cudart': 13 | from models.cudart_api import TRTEngine 14 | elif args.method == 'pycuda': 15 | from models.pycuda_api import TRTEngine 16 | else: 17 | raise NotImplementedError 18 | 19 | Engine = TRTEngine(args.engine) 20 | H, W = Engine.inp_info[0].shape[-2:] 21 | 22 | images = path_to_list(args.imgs) 23 | save_path = Path(args.out_dir) 24 | 25 | if not args.show and not save_path.exists(): 26 | save_path.mkdir(parents=True, exist_ok=True) 27 | 28 | for image in images: 29 | save_image = save_path / image.name 30 | bgr = cv2.imread(str(image)) 31 | draw = bgr.copy() 32 | bgr = cv2.resize(bgr, (W, H)) 33 | rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) 34 | tensor = blob(rgb, return_seg=False) 35 | tensor = np.ascontiguousarray(tensor) 36 | # inference 37 | data = Engine(tensor) 38 | 39 | data = data[0] 40 | score = data.max().item() 41 | cls_id = data.argmax().item() 42 | cls = CLASSES_CLS[cls_id] 43 | 44 | text = f'{cls}:{score:.3f}' 45 | (_w, _h), _bl = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.8, 1) 46 | _y1 = min(10, draw.shape[0]) 47 | 48 | cv2.rectangle(draw, (10, _y1), (10 + _w, _y1 + _h + _bl), (0, 0, 255), 49 | -1) 50 | cv2.putText(draw, text, (10, _y1 + _h), cv2.FONT_HERSHEY_SIMPLEX, 0.75, 51 | (255, 255, 255), 2) 52 | 53 | if args.show: 54 | cv2.imshow('result', draw) 55 | cv2.waitKey(0) 56 | else: 57 | cv2.imwrite(str(save_image), draw) 58 | 59 | 60 | def parse_args(): 61 | parser = argparse.ArgumentParser() 62 | parser.add_argument('--engine', type=str, help='Engine file') 63 | parser.add_argument('--imgs', type=str, help='Images file') 64 | parser.add_argument('--show', 65 | action='store_true', 66 | help='Show the detection results') 67 | parser.add_argument('--out-dir', 68 | type=str, 69 | default='./output', 70 | help='Path to output file') 71 | parser.add_argument('--method', 72 | type=str, 73 | default='cudart', 74 | help='CUDART pipeline') 75 | args = parser.parse_args() 76 | return args 77 | 78 | 79 | if __name__ == '__main__': 80 | args = parse_args() 81 | main(args) 82 | -------------------------------------------------------------------------------- /infer-cls.py: -------------------------------------------------------------------------------- 1 | from models import TRTModule # isort:skip 2 | import argparse 3 | from pathlib import Path 4 | 5 | import cv2 6 | import torch 7 | 8 | from config import CLASSES_CLS 9 | from models.utils import blob, path_to_list 10 | 11 | 12 | def main(args: argparse.Namespace) -> None: 13 | device = torch.device(args.device) 14 | Engine = TRTModule(args.engine, device) 15 | H, W = Engine.inp_info[0].shape[-2:] 16 | 17 | images = path_to_list(args.imgs) 18 | save_path = Path(args.out_dir) 19 | 20 | if not args.show and not save_path.exists(): 21 | save_path.mkdir(parents=True, exist_ok=True) 22 | 23 | for image in images: 24 | save_image = save_path / image.name 25 | bgr = cv2.imread(str(image)) 26 | draw = bgr.copy() 27 | bgr = cv2.resize(bgr, (W, H)) 28 | rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) 29 | tensor = blob(rgb, return_seg=False) 30 | tensor = torch.asarray(tensor, device=device) 31 | # inference 32 | data = Engine(tensor) 33 | 34 | score, cls_id = data[0].max(0) 35 | score = float(score) 36 | cls_id = int(cls_id) 37 | cls = CLASSES_CLS[cls_id] 38 | 39 | text = f'{cls}:{score:.3f}' 40 | (_w, _h), _bl = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.8, 1) 41 | _y1 = min(10, draw.shape[0]) 42 | 43 | cv2.rectangle(draw, (10, _y1), (10 + _w, _y1 + _h + _bl), (0, 0, 255), 44 | -1) 45 | cv2.putText(draw, text, (10, _y1 + _h), cv2.FONT_HERSHEY_SIMPLEX, 0.75, 46 | (255, 255, 255), 2) 47 | 48 | if args.show: 49 | cv2.imshow('result', draw) 50 | cv2.waitKey(0) 51 | else: 52 | cv2.imwrite(str(save_image), draw) 53 | 54 | 55 | def parse_args() -> argparse.Namespace: 56 | parser = argparse.ArgumentParser() 57 | parser.add_argument('--engine', type=str, help='Engine file') 58 | parser.add_argument('--imgs', type=str, help='Images file') 59 | parser.add_argument('--show', 60 | action='store_true', 61 | help='Show the detection results') 62 | parser.add_argument('--out-dir', 63 | type=str, 64 | default='./output', 65 | help='Path to output file') 66 | parser.add_argument('--device', 67 | type=str, 68 | default='cuda:0', 69 | help='TensorRT infer device') 70 | args = parser.parse_args() 71 | return args 72 | 73 | 74 | if __name__ == '__main__': 75 | args = parse_args() 76 | main(args) 77 | -------------------------------------------------------------------------------- /infer-det-without-torch.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from pathlib import Path 3 | 4 | import cv2 5 | import numpy as np 6 | 7 | from config import CLASSES_DET, COLORS 8 | from models.utils import blob, det_postprocess, letterbox, path_to_list 9 | 10 | 11 | def main(args: argparse.Namespace) -> None: 12 | if args.method == 'cudart': 13 | from models.cudart_api import TRTEngine 14 | elif args.method == 'pycuda': 15 | from models.pycuda_api import TRTEngine 16 | else: 17 | raise NotImplementedError 18 | 19 | Engine = TRTEngine(args.engine) 20 | H, W = Engine.inp_info[0].shape[-2:] 21 | 22 | images = path_to_list(args.imgs) 23 | save_path = Path(args.out_dir) 24 | 25 | if not args.show and not save_path.exists(): 26 | save_path.mkdir(parents=True, exist_ok=True) 27 | 28 | for image in images: 29 | save_image = save_path / image.name 30 | bgr = cv2.imread(str(image)) 31 | draw = bgr.copy() 32 | bgr, ratio, dwdh = letterbox(bgr, (W, H)) 33 | rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) 34 | tensor = blob(rgb, return_seg=False) 35 | dwdh = np.array(dwdh * 2, dtype=np.float32) 36 | tensor = np.ascontiguousarray(tensor) 37 | # inference 38 | data = Engine(tensor) 39 | 40 | bboxes, scores, labels = det_postprocess(data) 41 | if bboxes.size == 0: 42 | # if no bounding box 43 | print(f'{image}: no object!') 44 | continue 45 | bboxes -= dwdh 46 | bboxes /= ratio 47 | 48 | for (bbox, score, label) in zip(bboxes, scores, labels): 49 | bbox = bbox.round().astype(np.int32).tolist() 50 | cls_id = int(label) 51 | cls = CLASSES_DET[cls_id] 52 | color = COLORS[cls] 53 | 54 | text = f'{cls}:{score:.3f}' 55 | x1, y1, x2, y2 = bbox 56 | 57 | (_w, _h), _bl = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 58 | 0.8, 1) 59 | _y1 = min(y1 + 1, draw.shape[0]) 60 | 61 | cv2.rectangle(draw, (x1, y1), (x2, y2), color, 2) 62 | cv2.rectangle(draw, (x1, _y1), (x1 + _w, _y1 + _h + _bl), 63 | (0, 0, 255), -1) 64 | cv2.putText(draw, text, (x1, _y1 + _h), cv2.FONT_HERSHEY_SIMPLEX, 65 | 0.75, (255, 255, 255), 2) 66 | 67 | if args.show: 68 | cv2.imshow('result', draw) 69 | cv2.waitKey(0) 70 | else: 71 | cv2.imwrite(str(save_image), draw) 72 | 73 | 74 | def parse_args(): 75 | parser = argparse.ArgumentParser() 76 | parser.add_argument('--engine', type=str, help='Engine file') 77 | parser.add_argument('--imgs', type=str, help='Images file') 78 | parser.add_argument('--show', 79 | action='store_true', 80 | help='Show the detection results') 81 | parser.add_argument('--out-dir', 82 | type=str, 83 | default='./output', 84 | help='Path to output file') 85 | parser.add_argument('--method', 86 | type=str, 87 | default='cudart', 88 | help='CUDART pipeline') 89 | args = parser.parse_args() 90 | return args 91 | 92 | 93 | if __name__ == '__main__': 94 | args = parse_args() 95 | main(args) 96 | -------------------------------------------------------------------------------- /infer-det.py: -------------------------------------------------------------------------------- 1 | from models import TRTModule # isort:skip 2 | import argparse 3 | from pathlib import Path 4 | 5 | import cv2 6 | import torch 7 | 8 | from config import CLASSES_DET, COLORS 9 | from models.torch_utils import det_postprocess 10 | from models.utils import blob, letterbox, path_to_list 11 | 12 | 13 | def main(args: argparse.Namespace) -> None: 14 | device = torch.device(args.device) 15 | Engine = TRTModule(args.engine, device) 16 | H, W = Engine.inp_info[0].shape[-2:] 17 | 18 | # set desired output names order 19 | Engine.set_desired(['num_dets', 'bboxes', 'scores', 'labels']) 20 | 21 | images = path_to_list(args.imgs) 22 | save_path = Path(args.out_dir) 23 | 24 | if not args.show and not save_path.exists(): 25 | save_path.mkdir(parents=True, exist_ok=True) 26 | 27 | for image in images: 28 | save_image = save_path / image.name 29 | bgr = cv2.imread(str(image)) 30 | draw = bgr.copy() 31 | bgr, ratio, dwdh = letterbox(bgr, (W, H)) 32 | rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) 33 | tensor = blob(rgb, return_seg=False) 34 | dwdh = torch.asarray(dwdh * 2, dtype=torch.float32, device=device) 35 | tensor = torch.asarray(tensor, device=device) 36 | # inference 37 | data = Engine(tensor) 38 | 39 | bboxes, scores, labels = det_postprocess(data) 40 | if bboxes.numel() == 0: 41 | # if no bounding box 42 | print(f'{image}: no object!') 43 | continue 44 | bboxes -= dwdh 45 | bboxes /= ratio 46 | 47 | for (bbox, score, label) in zip(bboxes, scores, labels): 48 | bbox = bbox.round().int().tolist() 49 | cls_id = int(label) 50 | cls = CLASSES_DET[cls_id] 51 | color = COLORS[cls] 52 | 53 | text = f'{cls}:{score:.3f}' 54 | x1, y1, x2, y2 = bbox 55 | 56 | (_w, _h), _bl = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 57 | 0.8, 1) 58 | _y1 = min(y1 + 1, draw.shape[0]) 59 | 60 | cv2.rectangle(draw, (x1, y1), (x2, y2), color, 2) 61 | cv2.rectangle(draw, (x1, _y1), (x1 + _w, _y1 + _h + _bl), 62 | (0, 0, 255), -1) 63 | cv2.putText(draw, text, (x1, _y1 + _h), cv2.FONT_HERSHEY_SIMPLEX, 64 | 0.75, (255, 255, 255), 2) 65 | 66 | if args.show: 67 | cv2.imshow('result', draw) 68 | cv2.waitKey(0) 69 | else: 70 | cv2.imwrite(str(save_image), draw) 71 | 72 | 73 | def parse_args() -> argparse.Namespace: 74 | parser = argparse.ArgumentParser() 75 | parser.add_argument('--engine', type=str, help='Engine file') 76 | parser.add_argument('--imgs', type=str, help='Images file') 77 | parser.add_argument('--show', 78 | action='store_true', 79 | help='Show the detection results') 80 | parser.add_argument('--out-dir', 81 | type=str, 82 | default='./output', 83 | help='Path to output file') 84 | parser.add_argument('--device', 85 | type=str, 86 | default='cuda:0', 87 | help='TensorRT infer device') 88 | args = parser.parse_args() 89 | return args 90 | 91 | 92 | if __name__ == '__main__': 93 | args = parse_args() 94 | main(args) 95 | -------------------------------------------------------------------------------- /infer-obb-without-torch.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from pathlib import Path 3 | 4 | import cv2 5 | import numpy as np 6 | 7 | from config import CLASSES_OBB, COLORS_OBB 8 | from models.utils import blob, letterbox, obb_postprocess, path_to_list 9 | 10 | 11 | def main(args: argparse.Namespace) -> None: 12 | if args.method == 'cudart': 13 | from models.cudart_api import TRTEngine 14 | elif args.method == 'pycuda': 15 | from models.pycuda_api import TRTEngine 16 | else: 17 | raise NotImplementedError 18 | 19 | Engine = TRTEngine(args.engine) 20 | H, W = Engine.inp_info[0].shape[-2:] 21 | 22 | images = path_to_list(args.imgs) 23 | save_path = Path(args.out_dir) 24 | 25 | if not args.show and not save_path.exists(): 26 | save_path.mkdir(parents=True, exist_ok=True) 27 | 28 | for image in images: 29 | save_image = save_path / image.name 30 | bgr = cv2.imread(str(image)) 31 | draw = bgr.copy() 32 | bgr, ratio, dwdh = letterbox(bgr, (W, H)) 33 | rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) 34 | tensor = blob(rgb, return_seg=False) 35 | dwdh = np.array(dwdh, dtype=np.float32) 36 | tensor = np.ascontiguousarray(tensor) 37 | # inference 38 | data = Engine(tensor) 39 | 40 | points, scores, labels = obb_postprocess(data, args.conf_thres, 41 | args.iou_thres) 42 | if points.size == 0: 43 | # if no points 44 | print(f'{image}: no object!') 45 | continue 46 | points -= dwdh 47 | points /= ratio 48 | 49 | for (point, score, label) in zip(points, scores, labels): 50 | point = point.round().astype(np.int32) 51 | label = int(label) 52 | score = float(score) 53 | cls = CLASSES_OBB[label] 54 | color = COLORS_OBB[cls] 55 | cv2.polylines(draw, [point], True, color, 2) 56 | cv2.putText(draw, 57 | f'{cls}:{score:.3f}', (point[0, 0], point[0, 1] - 2), 58 | cv2.FONT_HERSHEY_SIMPLEX, 59 | 0.75, [225, 255, 255], 60 | thickness=2) 61 | 62 | if args.show: 63 | cv2.imshow('result', draw) 64 | cv2.waitKey(0) 65 | else: 66 | cv2.imwrite(str(save_image), draw) 67 | 68 | 69 | def parse_args(): 70 | parser = argparse.ArgumentParser() 71 | parser.add_argument('--engine', type=str, help='Engine file') 72 | parser.add_argument('--imgs', type=str, help='Images file') 73 | parser.add_argument('--show', 74 | action='store_true', 75 | help='Show the detection results') 76 | parser.add_argument('--out-dir', 77 | type=str, 78 | default='./output', 79 | help='Path to output file') 80 | parser.add_argument('--conf-thres', 81 | type=float, 82 | default=0.25, 83 | help='Confidence threshold') 84 | parser.add_argument('--iou-thres', 85 | type=float, 86 | default=0.65, 87 | help='Confidence threshold') 88 | parser.add_argument('--method', 89 | type=str, 90 | default='cudart', 91 | help='CUDART pipeline') 92 | args = parser.parse_args() 93 | return args 94 | 95 | 96 | if __name__ == '__main__': 97 | args = parse_args() 98 | main(args) 99 | -------------------------------------------------------------------------------- /infer-obb.py: -------------------------------------------------------------------------------- 1 | from models import TRTModule # isort:skip 2 | import argparse 3 | from pathlib import Path 4 | 5 | import cv2 6 | import torch 7 | 8 | from config import CLASSES_OBB, COLORS_OBB 9 | from models.torch_utils import obb_postprocess 10 | from models.utils import blob, letterbox, path_to_list 11 | 12 | 13 | def main(args: argparse.Namespace) -> None: 14 | device = torch.device(args.device) 15 | Engine = TRTModule(args.engine, device) 16 | H, W = Engine.inp_info[0].shape[-2:] 17 | 18 | images = path_to_list(args.imgs) 19 | save_path = Path(args.out_dir) 20 | 21 | if not args.show and not save_path.exists(): 22 | save_path.mkdir(parents=True, exist_ok=True) 23 | 24 | for image in images: 25 | save_image = save_path / image.name 26 | bgr = cv2.imread(str(image)) 27 | draw = bgr.copy() 28 | bgr, ratio, dwdh = letterbox(bgr, (W, H)) 29 | rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) 30 | tensor = blob(rgb, return_seg=False) 31 | dwdh = torch.asarray(dwdh, dtype=torch.float32, device=device) 32 | tensor = torch.asarray(tensor, device=device) 33 | # inference 34 | data = Engine(tensor) 35 | 36 | points, scores, labels = obb_postprocess(data, args.conf_thres, 37 | args.iou_thres) 38 | if points.numel() == 0: 39 | # if no points 40 | print(f'{image}: no object!') 41 | continue 42 | points -= dwdh 43 | points /= ratio 44 | 45 | for (point, score, label) in zip(points, scores, labels): 46 | point = point.round().int().cpu().numpy() 47 | label = int(label) 48 | score = float(score) 49 | cls = CLASSES_OBB[label] 50 | color = COLORS_OBB[cls] 51 | cv2.polylines(draw, [point], True, color, 2) 52 | cv2.putText(draw, 53 | f'{cls}:{score:.3f}', (point[0, 0], point[0, 1] - 2), 54 | cv2.FONT_HERSHEY_SIMPLEX, 55 | 0.75, [225, 255, 255], 56 | thickness=2) 57 | 58 | if args.show: 59 | cv2.imshow('result', draw) 60 | cv2.waitKey(0) 61 | else: 62 | cv2.imwrite(str(save_image), draw) 63 | 64 | 65 | def parse_args() -> argparse.Namespace: 66 | parser = argparse.ArgumentParser() 67 | parser.add_argument('--engine', type=str, help='Engine file') 68 | parser.add_argument('--imgs', type=str, help='Images file') 69 | parser.add_argument('--show', 70 | action='store_true', 71 | help='Show the detection results') 72 | parser.add_argument('--out-dir', 73 | type=str, 74 | default='./output', 75 | help='Path to output file') 76 | parser.add_argument('--conf-thres', 77 | type=float, 78 | default=0.25, 79 | help='Confidence threshold') 80 | parser.add_argument('--iou-thres', 81 | type=float, 82 | default=0.65, 83 | help='Confidence threshold') 84 | parser.add_argument('--device', 85 | type=str, 86 | default='cuda:0', 87 | help='TensorRT infer device') 88 | args = parser.parse_args() 89 | return args 90 | 91 | 92 | if __name__ == '__main__': 93 | args = parse_args() 94 | main(args) 95 | -------------------------------------------------------------------------------- /infer-pose-without-torch.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from pathlib import Path 3 | 4 | import cv2 5 | import numpy as np 6 | 7 | from config import COLORS, KPS_COLORS, LIMB_COLORS, SKELETON 8 | from models.utils import blob, letterbox, path_to_list, pose_postprocess 9 | 10 | 11 | def main(args: argparse.Namespace) -> None: 12 | if args.method == 'cudart': 13 | from models.cudart_api import TRTEngine 14 | elif args.method == 'pycuda': 15 | from models.pycuda_api import TRTEngine 16 | else: 17 | raise NotImplementedError 18 | 19 | Engine = TRTEngine(args.engine) 20 | H, W = Engine.inp_info[0].shape[-2:] 21 | 22 | images = path_to_list(args.imgs) 23 | save_path = Path(args.out_dir) 24 | 25 | if not args.show and not save_path.exists(): 26 | save_path.mkdir(parents=True, exist_ok=True) 27 | 28 | for image in images: 29 | save_image = save_path / image.name 30 | bgr = cv2.imread(str(image)) 31 | draw = bgr.copy() 32 | bgr, ratio, dwdh = letterbox(bgr, (W, H)) 33 | dw, dh = int(dwdh[0]), int(dwdh[1]) 34 | rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) 35 | tensor = blob(rgb, return_seg=False) 36 | dwdh = np.array(dwdh * 2, dtype=np.float32) 37 | tensor = np.ascontiguousarray(tensor) 38 | # inference 39 | data = Engine(tensor) 40 | 41 | bboxes, scores, kpts = pose_postprocess(data, args.conf_thres, 42 | args.iou_thres) 43 | if bboxes.size == 0: 44 | # if no bounding box 45 | print(f'{image}: no object!') 46 | continue 47 | bboxes -= dwdh 48 | bboxes /= ratio 49 | 50 | for (bbox, score, kpt) in zip(bboxes, scores, kpts): 51 | bbox = bbox.round().astype(np.int32).tolist() 52 | color = COLORS['person'] 53 | cv2.rectangle(draw, bbox[:2], bbox[2:], color, 2) 54 | cv2.putText(draw, 55 | f'person:{score:.3f}', (bbox[0], bbox[1] - 2), 56 | cv2.FONT_HERSHEY_SIMPLEX, 57 | 0.75, [225, 255, 255], 58 | thickness=2) 59 | for i in range(19): 60 | if i < 17: 61 | px, py, ps = kpt[i] 62 | if ps > 0.5: 63 | kcolor = KPS_COLORS[i] 64 | px = round(float(px - dw) / ratio) 65 | py = round(float(py - dh) / ratio) 66 | cv2.circle(draw, (px, py), 5, kcolor, -1) 67 | xi, yi = SKELETON[i] 68 | pos1_s = kpt[xi - 1][2] 69 | pos2_s = kpt[yi - 1][2] 70 | if pos1_s > 0.5 and pos2_s > 0.5: 71 | limb_color = LIMB_COLORS[i] 72 | pos1_x = round(float(kpt[xi - 1][0] - dw) / ratio) 73 | pos1_y = round(float(kpt[xi - 1][1] - dh) / ratio) 74 | 75 | pos2_x = round(float(kpt[yi - 1][0] - dw) / ratio) 76 | pos2_y = round(float(kpt[yi - 1][1] - dh) / ratio) 77 | 78 | cv2.line(draw, (pos1_x, pos1_y), (pos2_x, pos2_y), 79 | limb_color, 2) 80 | if args.show: 81 | cv2.imshow('result', draw) 82 | cv2.waitKey(0) 83 | else: 84 | cv2.imwrite(str(save_image), draw) 85 | 86 | 87 | def parse_args(): 88 | parser = argparse.ArgumentParser() 89 | parser.add_argument('--engine', type=str, help='Engine file') 90 | parser.add_argument('--imgs', type=str, help='Images file') 91 | parser.add_argument('--show', 92 | action='store_true', 93 | help='Show the detection results') 94 | parser.add_argument('--out-dir', 95 | type=str, 96 | default='./output', 97 | help='Path to output file') 98 | parser.add_argument('--conf-thres', 99 | type=float, 100 | default=0.25, 101 | help='Confidence threshold') 102 | parser.add_argument('--iou-thres', 103 | type=float, 104 | default=0.65, 105 | help='Confidence threshold') 106 | parser.add_argument('--method', 107 | type=str, 108 | default='cudart', 109 | help='CUDART pipeline') 110 | args = parser.parse_args() 111 | return args 112 | 113 | 114 | if __name__ == '__main__': 115 | args = parse_args() 116 | main(args) 117 | -------------------------------------------------------------------------------- /infer-pose.py: -------------------------------------------------------------------------------- 1 | from models import TRTModule # isort:skip 2 | import argparse 3 | from pathlib import Path 4 | 5 | import cv2 6 | import torch 7 | 8 | from config import COLORS, KPS_COLORS, LIMB_COLORS, SKELETON 9 | from models.torch_utils import pose_postprocess 10 | from models.utils import blob, letterbox, path_to_list 11 | 12 | 13 | def main(args: argparse.Namespace) -> None: 14 | device = torch.device(args.device) 15 | Engine = TRTModule(args.engine, device) 16 | H, W = Engine.inp_info[0].shape[-2:] 17 | 18 | images = path_to_list(args.imgs) 19 | save_path = Path(args.out_dir) 20 | 21 | if not args.show and not save_path.exists(): 22 | save_path.mkdir(parents=True, exist_ok=True) 23 | 24 | for image in images: 25 | save_image = save_path / image.name 26 | bgr = cv2.imread(str(image)) 27 | draw = bgr.copy() 28 | bgr, ratio, dwdh = letterbox(bgr, (W, H)) 29 | dw, dh = int(dwdh[0]), int(dwdh[1]) 30 | rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) 31 | tensor = blob(rgb, return_seg=False) 32 | dwdh = torch.asarray(dwdh * 2, dtype=torch.float32, device=device) 33 | tensor = torch.asarray(tensor, device=device) 34 | # inference 35 | data = Engine(tensor) 36 | 37 | bboxes, scores, kpts = pose_postprocess(data, args.conf_thres, 38 | args.iou_thres) 39 | if bboxes.numel() == 0: 40 | # if no bounding box 41 | print(f'{image}: no object!') 42 | continue 43 | bboxes -= dwdh 44 | bboxes /= ratio 45 | 46 | for (bbox, score, kpt) in zip(bboxes, scores, kpts): 47 | bbox = bbox.round().int().tolist() 48 | color = COLORS['person'] 49 | cv2.rectangle(draw, bbox[:2], bbox[2:], color, 2) 50 | cv2.putText(draw, 51 | f'person:{score:.3f}', (bbox[0], bbox[1] - 2), 52 | cv2.FONT_HERSHEY_SIMPLEX, 53 | 0.75, [225, 255, 255], 54 | thickness=2) 55 | for i in range(19): 56 | if i < 17: 57 | px, py, ps = kpt[i] 58 | if ps > 0.5: 59 | kcolor = KPS_COLORS[i] 60 | px = round(float(px - dw) / ratio) 61 | py = round(float(py - dh) / ratio) 62 | cv2.circle(draw, (px, py), 5, kcolor, -1) 63 | xi, yi = SKELETON[i] 64 | pos1_s = kpt[xi - 1][2] 65 | pos2_s = kpt[yi - 1][2] 66 | if pos1_s > 0.5 and pos2_s > 0.5: 67 | limb_color = LIMB_COLORS[i] 68 | pos1_x = round(float(kpt[xi - 1][0] - dw) / ratio) 69 | pos1_y = round(float(kpt[xi - 1][1] - dh) / ratio) 70 | 71 | pos2_x = round(float(kpt[yi - 1][0] - dw) / ratio) 72 | pos2_y = round(float(kpt[yi - 1][1] - dh) / ratio) 73 | 74 | cv2.line(draw, (pos1_x, pos1_y), (pos2_x, pos2_y), 75 | limb_color, 2) 76 | if args.show: 77 | cv2.imshow('result', draw) 78 | cv2.waitKey(0) 79 | else: 80 | cv2.imwrite(str(save_image), draw) 81 | 82 | 83 | def parse_args() -> argparse.Namespace: 84 | parser = argparse.ArgumentParser() 85 | parser.add_argument('--engine', type=str, help='Engine file') 86 | parser.add_argument('--imgs', type=str, help='Images file') 87 | parser.add_argument('--show', 88 | action='store_true', 89 | help='Show the detection results') 90 | parser.add_argument('--out-dir', 91 | type=str, 92 | default='./output', 93 | help='Path to output file') 94 | parser.add_argument('--conf-thres', 95 | type=float, 96 | default=0.25, 97 | help='Confidence threshold') 98 | parser.add_argument('--iou-thres', 99 | type=float, 100 | default=0.65, 101 | help='Confidence threshold') 102 | parser.add_argument('--device', 103 | type=str, 104 | default='cuda:0', 105 | help='TensorRT infer device') 106 | args = parser.parse_args() 107 | return args 108 | 109 | 110 | if __name__ == '__main__': 111 | args = parse_args() 112 | main(args) 113 | -------------------------------------------------------------------------------- /infer-seg-without-torch.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from pathlib import Path 3 | 4 | import cv2 5 | import numpy as np 6 | 7 | from config import ALPHA, CLASSES_SEG, COLORS, MASK_COLORS 8 | from models.utils import blob, letterbox, path_to_list, seg_postprocess 9 | 10 | 11 | def main(args: argparse.Namespace) -> None: 12 | if args.method == 'cudart': 13 | from models.cudart_api import TRTEngine 14 | elif args.method == 'pycuda': 15 | from models.pycuda_api import TRTEngine 16 | else: 17 | raise NotImplementedError 18 | 19 | Engine = TRTEngine(args.engine) 20 | H, W = Engine.inp_info[0].shape[-2:] 21 | 22 | images = path_to_list(args.imgs) 23 | save_path = Path(args.out_dir) 24 | 25 | if not args.show and not save_path.exists(): 26 | save_path.mkdir(parents=True, exist_ok=True) 27 | 28 | for image in images: 29 | save_image = save_path / image.name 30 | bgr = cv2.imread(str(image)) 31 | draw = bgr.copy() 32 | bgr, ratio, dwdh = letterbox(bgr, (W, H)) 33 | dw, dh = int(dwdh[0]), int(dwdh[1]) 34 | rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) 35 | tensor, seg_img = blob(rgb, return_seg=True) 36 | dwdh = np.array(dwdh * 2, dtype=np.float32) 37 | tensor = np.ascontiguousarray(tensor) 38 | # inference 39 | data = Engine(tensor) 40 | 41 | seg_img = seg_img[dh:H - dh, dw:W - dw, [2, 1, 0]] 42 | bboxes, scores, labels, masks = seg_postprocess( 43 | data, bgr.shape[:2], args.conf_thres, args.iou_thres) 44 | if bboxes.size == 0: 45 | # if no bounding box 46 | print(f'{image}: no object!') 47 | continue 48 | masks = masks[:, dh:H - dh, dw:W - dw, :] 49 | mask_colors = MASK_COLORS[labels % len(MASK_COLORS)] 50 | mask_colors = mask_colors.reshape(-1, 1, 1, 3) * ALPHA 51 | mask_colors = masks @ mask_colors 52 | inv_alph_masks = (1 - masks * 0.5).cumprod(0) 53 | mcs = (mask_colors * inv_alph_masks).sum(0) * 2 54 | seg_img = (seg_img * inv_alph_masks[-1] + mcs) * 255 55 | draw = cv2.resize(seg_img.astype(np.uint8), draw.shape[:2][::-1]) 56 | 57 | bboxes -= dwdh 58 | bboxes /= ratio 59 | 60 | for (bbox, score, label) in zip(bboxes, scores, labels): 61 | bbox = bbox.round().astype(np.int32).tolist() 62 | cls_id = int(label) 63 | cls = CLASSES_SEG[cls_id] 64 | color = COLORS[cls] 65 | cv2.rectangle(draw, bbox[:2], bbox[2:], color, 2) 66 | cv2.putText(draw, 67 | f'{cls}:{score:.3f}', (bbox[0], bbox[1] - 2), 68 | cv2.FONT_HERSHEY_SIMPLEX, 69 | 0.75, [225, 255, 255], 70 | thickness=2) 71 | if args.show: 72 | cv2.imshow('result', draw) 73 | cv2.waitKey(0) 74 | else: 75 | cv2.imwrite(str(save_image), draw) 76 | 77 | 78 | def parse_args(): 79 | parser = argparse.ArgumentParser() 80 | parser.add_argument('--engine', type=str, help='Engine file') 81 | parser.add_argument('--imgs', type=str, help='Images file') 82 | parser.add_argument('--show', 83 | action='store_true', 84 | help='Show the detection results') 85 | parser.add_argument('--out-dir', 86 | type=str, 87 | default='./output', 88 | help='Path to output file') 89 | parser.add_argument('--conf-thres', 90 | type=float, 91 | default=0.25, 92 | help='Confidence threshold') 93 | parser.add_argument('--iou-thres', 94 | type=float, 95 | default=0.65, 96 | help='Confidence threshold') 97 | parser.add_argument('--method', 98 | type=str, 99 | default='cudart', 100 | help='CUDART pipeline') 101 | args = parser.parse_args() 102 | return args 103 | 104 | 105 | if __name__ == '__main__': 106 | args = parse_args() 107 | main(args) 108 | -------------------------------------------------------------------------------- /infer-seg.py: -------------------------------------------------------------------------------- 1 | from models import TRTModule # isort:skip 2 | import argparse 3 | from pathlib import Path 4 | 5 | import cv2 6 | import numpy as np 7 | import torch 8 | 9 | from config import ALPHA, CLASSES_SEG, COLORS, MASK_COLORS 10 | from models.torch_utils import seg_postprocess 11 | from models.utils import blob, letterbox, path_to_list 12 | 13 | 14 | def main(args: argparse.Namespace) -> None: 15 | device = torch.device(args.device) 16 | Engine = TRTModule(args.engine, device) 17 | H, W = Engine.inp_info[0].shape[-2:] 18 | 19 | # set desired output names order 20 | Engine.set_desired(['outputs', 'proto']) 21 | 22 | images = path_to_list(args.imgs) 23 | save_path = Path(args.out_dir) 24 | 25 | if not args.show and not save_path.exists(): 26 | save_path.mkdir(parents=True, exist_ok=True) 27 | 28 | for image in images: 29 | save_image = save_path / image.name 30 | bgr = cv2.imread(str(image)) 31 | draw = bgr.copy() 32 | bgr, ratio, dwdh = letterbox(bgr, (W, H)) 33 | dw, dh = int(dwdh[0]), int(dwdh[1]) 34 | rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) 35 | tensor, seg_img = blob(rgb, return_seg=True) 36 | dwdh = torch.asarray(dwdh * 2, dtype=torch.float32, device=device) 37 | tensor = torch.asarray(tensor, device=device) 38 | # inference 39 | data = Engine(tensor) 40 | 41 | seg_img = torch.asarray(seg_img[dh:H - dh, dw:W - dw, [2, 1, 0]], 42 | device=device) 43 | bboxes, scores, labels, masks = seg_postprocess( 44 | data, bgr.shape[:2], args.conf_thres, args.iou_thres) 45 | if bboxes.numel() == 0: 46 | # if no bounding box 47 | print(f'{image}: no object!') 48 | continue 49 | masks = masks[:, dh:H - dh, dw:W - dw, :] 50 | indices = (labels % len(MASK_COLORS)).long() 51 | mask_colors = torch.asarray(MASK_COLORS, device=device)[indices] 52 | mask_colors = mask_colors.view(-1, 1, 1, 3) * ALPHA 53 | mask_colors = masks @ mask_colors 54 | inv_alph_masks = (1 - masks * 0.5).cumprod(0) 55 | mcs = (mask_colors * inv_alph_masks).sum(0) * 2 56 | seg_img = (seg_img * inv_alph_masks[-1] + mcs) * 255 57 | draw = cv2.resize(seg_img.cpu().numpy().astype(np.uint8), 58 | draw.shape[:2][::-1]) 59 | 60 | bboxes -= dwdh 61 | bboxes /= ratio 62 | 63 | for (bbox, score, label) in zip(bboxes, scores, labels): 64 | bbox = bbox.round().int().tolist() 65 | cls_id = int(label) 66 | cls = CLASSES_SEG[cls_id] 67 | color = COLORS[cls] 68 | cv2.rectangle(draw, bbox[:2], bbox[2:], color, 2) 69 | cv2.putText(draw, 70 | f'{cls}:{score:.3f}', (bbox[0], bbox[1] - 2), 71 | cv2.FONT_HERSHEY_SIMPLEX, 72 | 0.75, [225, 255, 255], 73 | thickness=2) 74 | if args.show: 75 | cv2.imshow('result', draw) 76 | cv2.waitKey(0) 77 | else: 78 | cv2.imwrite(str(save_image), draw) 79 | 80 | 81 | def parse_args() -> argparse.Namespace: 82 | parser = argparse.ArgumentParser() 83 | parser.add_argument('--engine', type=str, help='Engine file') 84 | parser.add_argument('--imgs', type=str, help='Images file') 85 | parser.add_argument('--show', 86 | action='store_true', 87 | help='Show the detection results') 88 | parser.add_argument('--out-dir', 89 | type=str, 90 | default='./output', 91 | help='Path to output file') 92 | parser.add_argument('--conf-thres', 93 | type=float, 94 | default=0.25, 95 | help='Confidence threshold') 96 | parser.add_argument('--iou-thres', 97 | type=float, 98 | default=0.65, 99 | help='Confidence threshold') 100 | parser.add_argument('--device', 101 | type=str, 102 | default='cuda:0', 103 | help='TensorRT infer device') 104 | args = parser.parse_args() 105 | return args 106 | 107 | 108 | if __name__ == '__main__': 109 | args = parse_args() 110 | main(args) 111 | -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- 1 | from .engine import EngineBuilder, TRTModule, TRTProfilerV0, TRTProfilerV1 # isort:skip # noqa: E501 2 | import warnings 3 | 4 | import torch 5 | 6 | warnings.filterwarnings(action='ignore', category=torch.jit.TracerWarning) 7 | warnings.filterwarnings(action='ignore', category=torch.jit.ScriptWarning) 8 | warnings.filterwarnings(action='ignore', category=UserWarning) 9 | warnings.filterwarnings(action='ignore', category=FutureWarning) 10 | warnings.filterwarnings(action='ignore', category=DeprecationWarning) 11 | __all__ = ['EngineBuilder', 'TRTModule', 'TRTProfilerV0', 'TRTProfilerV1'] 12 | -------------------------------------------------------------------------------- /models/pycuda_api.py: -------------------------------------------------------------------------------- 1 | import os 2 | import warnings 3 | from dataclasses import dataclass 4 | from pathlib import Path 5 | from typing import List, Optional, Tuple, Union 6 | 7 | import numpy as np 8 | import pycuda.autoinit # noqa F401 9 | import pycuda.driver as cuda 10 | import tensorrt as trt 11 | from numpy import ndarray 12 | 13 | os.environ['CUDA_MODULE_LOADING'] = 'LAZY' 14 | warnings.filterwarnings(action='ignore', category=DeprecationWarning) 15 | 16 | 17 | @dataclass 18 | class Tensor: 19 | name: str 20 | dtype: np.dtype 21 | shape: Tuple 22 | cpu: ndarray 23 | gpu: int 24 | 25 | 26 | class TRTEngine: 27 | 28 | def __init__(self, weight: Union[str, Path]) -> None: 29 | self.weight = Path(weight) if isinstance(weight, str) else weight 30 | self.stream = cuda.Stream(0) 31 | self.__init_engine() 32 | self.__init_bindings() 33 | self.__warm_up() 34 | 35 | def __init_engine(self) -> None: 36 | logger = trt.Logger(trt.Logger.WARNING) 37 | trt.init_libnvinfer_plugins(logger, namespace='') 38 | with trt.Runtime(logger) as runtime: 39 | model = runtime.deserialize_cuda_engine(self.weight.read_bytes()) 40 | 41 | context = model.create_execution_context() 42 | 43 | names = [model.get_binding_name(i) for i in range(model.num_bindings)] 44 | self.num_bindings = model.num_bindings 45 | self.bindings: List[int] = [0] * self.num_bindings 46 | num_inputs, num_outputs = 0, 0 47 | 48 | for i in range(model.num_bindings): 49 | if model.binding_is_input(i): 50 | num_inputs += 1 51 | else: 52 | num_outputs += 1 53 | 54 | self.num_inputs = num_inputs 55 | self.num_outputs = num_outputs 56 | self.model = model 57 | self.context = context 58 | self.input_names = names[:num_inputs] 59 | self.output_names = names[num_inputs:] 60 | 61 | def __init_bindings(self) -> None: 62 | dynamic = False 63 | inp_info = [] 64 | out_info = [] 65 | out_ptrs = [] 66 | for i, name in enumerate(self.input_names): 67 | assert self.model.get_binding_name(i) == name 68 | dtype = trt.nptype(self.model.get_binding_dtype(i)) 69 | shape = tuple(self.model.get_binding_shape(i)) 70 | if -1 in shape: 71 | dynamic |= True 72 | if not dynamic: 73 | cpu = np.empty(shape, dtype) 74 | gpu = cuda.mem_alloc(cpu.nbytes) 75 | cuda.memcpy_htod_async(gpu, cpu, self.stream) 76 | else: 77 | cpu, gpu = np.empty(0), 0 78 | inp_info.append(Tensor(name, dtype, shape, cpu, gpu)) 79 | for i, name in enumerate(self.output_names): 80 | i += self.num_inputs 81 | assert self.model.get_binding_name(i) == name 82 | dtype = trt.nptype(self.model.get_binding_dtype(i)) 83 | shape = tuple(self.model.get_binding_shape(i)) 84 | if not dynamic: 85 | cpu = np.empty(shape, dtype=dtype) 86 | gpu = cuda.mem_alloc(cpu.nbytes) 87 | cuda.memcpy_htod_async(gpu, cpu, self.stream) 88 | out_ptrs.append(gpu) 89 | else: 90 | cpu, gpu = np.empty(0), 0 91 | out_info.append(Tensor(name, dtype, shape, cpu, gpu)) 92 | 93 | self.is_dynamic = dynamic 94 | self.inp_info = inp_info 95 | self.out_info = out_info 96 | self.out_ptrs = out_ptrs 97 | 98 | def __warm_up(self) -> None: 99 | if self.is_dynamic: 100 | print('You engine has dynamic axes, please warm up by yourself !') 101 | return 102 | for _ in range(10): 103 | inputs = [] 104 | for i in self.inp_info: 105 | inputs.append(i.cpu) 106 | self.__call__(inputs) 107 | 108 | def set_profiler(self, profiler: Optional[trt.IProfiler]) -> None: 109 | self.context.profiler = profiler \ 110 | if profiler is not None else trt.Profiler() 111 | 112 | def __call__(self, *inputs) -> Union[Tuple, ndarray]: 113 | 114 | assert len(inputs) == self.num_inputs 115 | contiguous_inputs: List[ndarray] = [ 116 | np.ascontiguousarray(i) for i in inputs 117 | ] 118 | 119 | for i in range(self.num_inputs): 120 | 121 | if self.is_dynamic: 122 | self.context.set_binding_shape( 123 | i, tuple(contiguous_inputs[i].shape)) 124 | self.inp_info[i].gpu = cuda.mem_alloc( 125 | contiguous_inputs[i].nbytes) 126 | 127 | cuda.memcpy_htod_async(self.inp_info[i].gpu, contiguous_inputs[i], 128 | self.stream) 129 | self.bindings[i] = int(self.inp_info[i].gpu) 130 | 131 | output_gpu_ptrs: List[int] = [] 132 | outputs: List[ndarray] = [] 133 | 134 | for i in range(self.num_outputs): 135 | j = i + self.num_inputs 136 | if self.is_dynamic: 137 | shape = tuple(self.context.get_binding_shape(j)) 138 | dtype = self.out_info[i].dtype 139 | cpu = np.empty(shape, dtype=dtype) 140 | gpu = cuda.mem_alloc(cpu.nbytes) 141 | cuda.memcpy_htod_async(gpu, cpu, self.stream) 142 | else: 143 | cpu = self.out_info[i].cpu 144 | gpu = self.out_info[i].gpu 145 | outputs.append(cpu) 146 | output_gpu_ptrs.append(gpu) 147 | self.bindings[j] = int(gpu) 148 | 149 | self.context.execute_async_v2(self.bindings, self.stream.handle) 150 | self.stream.synchronize() 151 | 152 | for i, o in enumerate(output_gpu_ptrs): 153 | cuda.memcpy_dtoh_async(outputs[i], o, self.stream) 154 | 155 | return tuple(outputs) if len(outputs) > 1 else outputs[0] 156 | -------------------------------------------------------------------------------- /models/torch_utils.py: -------------------------------------------------------------------------------- 1 | from typing import List, Tuple, Union 2 | 3 | import torch 4 | import torch.nn.functional as F 5 | from torch import Tensor 6 | from torchvision.ops import batched_nms, nms 7 | 8 | from .utils import obb_postprocess as np_obb_postprocess 9 | 10 | 11 | def seg_postprocess( 12 | data: Tuple[Tensor], 13 | shape: Union[Tuple, List], 14 | conf_thres: float = 0.25, 15 | iou_thres: float = 0.65) \ 16 | -> Tuple[Tensor, Tensor, Tensor, Tensor]: 17 | assert len(data) == 2 18 | h, w = shape[0] // 4, shape[1] // 4 # 4x downsampling 19 | outputs, proto = data[0][0], data[1][0] 20 | bboxes, scores, labels, maskconf = outputs.split([4, 1, 1, 32], 1) 21 | scores, labels = scores.squeeze(), labels.squeeze() 22 | idx = scores > conf_thres 23 | if not idx.any(): # no bounding boxes or seg were created 24 | return bboxes.new_zeros((0, 4)), scores.new_zeros( 25 | (0, )), labels.new_zeros((0, )), bboxes.new_zeros((0, 0, 0, 0)) 26 | bboxes, scores, labels, maskconf = \ 27 | bboxes[idx], scores[idx], labels[idx], maskconf[idx] 28 | idx = batched_nms(bboxes, scores, labels, iou_thres) 29 | bboxes, scores, labels, maskconf = \ 30 | bboxes[idx], scores[idx], labels[idx].int(), maskconf[idx] 31 | masks = (maskconf @ proto).sigmoid().view(-1, h, w) 32 | masks = crop_mask(masks, bboxes / 4.) 33 | masks = F.interpolate(masks[None], 34 | shape, 35 | mode='bilinear', 36 | align_corners=False)[0] 37 | masks = masks.gt_(0.5)[..., None] 38 | return bboxes, scores, labels, masks 39 | 40 | 41 | def pose_postprocess( 42 | data: Union[Tuple, Tensor], 43 | conf_thres: float = 0.25, 44 | iou_thres: float = 0.65) \ 45 | -> Tuple[Tensor, Tensor, Tensor]: 46 | if isinstance(data, tuple): 47 | assert len(data) == 1 48 | data = data[0] 49 | outputs = torch.transpose(data[0], 0, 1).contiguous() 50 | bboxes, scores, kpts = outputs.split([4, 1, 51], 1) 51 | scores, kpts = scores.squeeze(), kpts.squeeze() 52 | idx = scores > conf_thres 53 | if not idx.any(): # no bounding boxes or seg were created 54 | return bboxes.new_zeros((0, 4)), scores.new_zeros( 55 | (0, )), bboxes.new_zeros((0, 0, 0)) 56 | bboxes, scores, kpts = bboxes[idx], scores[idx], kpts[idx] 57 | xycenter, wh = bboxes.chunk(2, -1) 58 | bboxes = torch.cat([xycenter - 0.5 * wh, xycenter + 0.5 * wh], -1) 59 | idx = nms(bboxes, scores, iou_thres) 60 | bboxes, scores, kpts = bboxes[idx], scores[idx], kpts[idx] 61 | return bboxes, scores, kpts.reshape(idx.shape[0], -1, 3) 62 | 63 | 64 | def det_postprocess(data: Tuple[Tensor, Tensor, Tensor, Tensor]): 65 | assert len(data) == 4 66 | iou_thres: float = 0.65 # noqa F841 67 | num_dets, bboxes, scores, labels = data[0][0], data[1][0], data[2][ 68 | 0], data[3][0] 69 | nums = num_dets.item() 70 | if nums == 0: 71 | return bboxes.new_zeros((0, 4)), scores.new_zeros( 72 | (0, )), labels.new_zeros((0, )) 73 | # check score negative 74 | scores[scores < 0] = 1 + scores[scores < 0] 75 | bboxes = bboxes[:nums] 76 | scores = scores[:nums] 77 | labels = labels[:nums] 78 | 79 | return bboxes, scores, labels 80 | 81 | 82 | def obb_postprocess( 83 | data: Union[Tuple, Tensor], 84 | conf_thres: float = 0.25, 85 | iou_thres: float = 0.65) \ 86 | -> Tuple[Tensor, Tensor, Tensor]: 87 | if isinstance(data, tuple): 88 | assert len(data) == 1 89 | data = data[0] 90 | device = data.device 91 | points, scores, labels = np_obb_postprocess(data.cpu().numpy(), conf_thres, 92 | iou_thres) 93 | return torch.from_numpy(points).to(device), torch.from_numpy(scores).to( 94 | device), torch.from_numpy(labels).to(device) 95 | 96 | 97 | def crop_mask(masks: Tensor, bboxes: Tensor) -> Tensor: 98 | n, h, w = masks.shape 99 | x1, y1, x2, y2 = torch.chunk(bboxes[:, :, None], 4, 1) # x1 shape(1,1,n) 100 | r = torch.arange(w, device=masks.device, 101 | dtype=x1.dtype)[None, None, :] # rows shape(1,w,1) 102 | c = torch.arange(h, device=masks.device, 103 | dtype=x1.dtype)[None, :, None] # cols shape(h,1,1) 104 | 105 | return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2)) 106 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | numpy<=1.23.5 2 | onnx 3 | onnxsim 4 | opencv-python 5 | torch 6 | torchvision 7 | ultralytics 8 | # tensorrt 9 | # cuda-python 10 | # pycuda 11 | -------------------------------------------------------------------------------- /trt-profile.py: -------------------------------------------------------------------------------- 1 | from models import TRTModule, TRTProfilerV0 # isort:skip 2 | import argparse 3 | 4 | import torch 5 | 6 | 7 | def profile(args): 8 | device = torch.device(args.device) 9 | Engine = TRTModule(args.engine, device) 10 | profiler = TRTProfilerV0() 11 | Engine.set_profiler(profiler) 12 | random_input = torch.randn(Engine.inp_info[0].shape, device=device) 13 | _ = Engine(random_input) 14 | 15 | 16 | def parse_args(): 17 | parser = argparse.ArgumentParser() 18 | parser.add_argument('--engine', type=str, help='Engine file') 19 | parser.add_argument('--device', 20 | type=str, 21 | default='cuda:0', 22 | help='TensorRT infer device') 23 | args = parser.parse_args() 24 | return args 25 | 26 | 27 | if __name__ == '__main__': 28 | args = parse_args() 29 | profile(args) 30 | --------------------------------------------------------------------------------