├── .gitignore ├── CMakeLists.txt ├── CMakeLists_deepsort-tensorrt_win10.txt ├── CMakeLists_yolov5-deepsort-tensorrt_win10.txt ├── LICENSE ├── README.md ├── assets └── yolosort.gif ├── deepsort ├── include │ ├── datatype.h │ ├── deepsort.h │ ├── deepsortenginegenerator.h │ ├── featuretensor.h │ ├── hungarianoper.h │ ├── kalmanfilter.h │ ├── linear_assignment.h │ ├── logging.h │ ├── matrix.h │ ├── model.hpp │ ├── munkres.h │ ├── nn_matching.h │ ├── track.h │ └── tracker.h └── src │ ├── deepsort.cpp │ ├── deepsortenginegenerator.cpp │ ├── featuretensor.cpp │ ├── hungarianoper.cpp │ ├── kalmanfilter.cpp │ ├── linear_assignment.cpp │ ├── munkres.cpp │ ├── nn_matching.cpp │ ├── track.cpp │ └── tracker.cpp ├── include └── manager.hpp ├── src ├── main.cpp └── manager.cpp └── yolo ├── include ├── calibrator.h ├── common.hpp ├── cuda_utils.h ├── logging.h ├── macros.h ├── utils.h ├── yololayer.cu ├── yololayer.h └── yolov5_lib.h └── src ├── calibrator.cpp └── yolov5_lib.cpp /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | build/ 131 | resources/ 132 | python/ -------------------------------------------------------------------------------- /CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 2.6) 2 | 3 | project(yolosort) 4 | list(APPEND CUDA_NVCC_FLAGS "-std=c++11") 5 | set(CMAKE_CXX_FLAGS "-std=c++0x") 6 | find_package(OpenCV REQUIRED) 7 | add_definitions(-std=c++11) 8 | add_definitions(-DAPI_EXPORTS) 9 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 10 | set(CMAKE_CXX_STANDARD 11) 11 | set(CMAKE_BUILD_TYPE Release) 12 | 13 | find_package(CUDA REQUIRED) 14 | set(CUDA_NVCC_PLAGS ${CUDA_NVCC_PLAGS};-std=c++11;-g;-G;-gencode;arch=compute_53;code=sm_53) 15 | 16 | if(WIN32) 17 | enable_language(CUDA) 18 | endif(WIN32) 19 | 20 | include_directories(${PROJECT_SOURCE_DIR}/deepsort/include) 21 | # include and link dirs of cuda and tensorrt, you need adapt them if yours are different 22 | # cuda 23 | include_directories(/usr/local/cuda/include) 24 | link_directories(/usr/local/cuda/lib64) 25 | # tensorrt 26 | include_directories(/usr/include/aarch64-linux-gnu/) 27 | link_directories(/usr/lib/aarch64-linux-gnu/) 28 | 29 | include_directories( 30 | ${CUDA_INCLUDE_DIRS} 31 | ${OpenCV_INCLUDE_DIRS} 32 | ${PROJECT_SOURCE_DIR}/deepsort/include 33 | ) 34 | aux_source_directory(${PROJECT_SOURCE_DIR}/deepsort/src SRC_DIR) 35 | 36 | # ===== deepsort ===== 37 | add_library(deepsort SHARED ${SRC_DIR}) 38 | target_link_libraries(deepsort 39 | ${CUDA_LIBS} ${OpenCV_LIBS} 40 | cudart nvinfer nvonnxparser 41 | ) 42 | 43 | # ===== yolo ===== 44 | include_directories(${PROJECT_SOURCE_DIR}/yolo/include) 45 | aux_source_directory(${PROJECT_SOURCE_DIR}/yolo/src YOLO_SRC_DIR) 46 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED") 47 | 48 | cuda_add_library(yolov5_trt SHARED ${PROJECT_SOURCE_DIR}/yolo/include/yololayer.cu ${PROJECT_SOURCE_DIR}/yolo/src/yolov5_lib.cpp) 49 | target_link_libraries(yolov5_trt nvinfer cudart deepsort) 50 | 51 | 52 | # ===== main ===== 53 | aux_source_directory(${PROJECT_SOURCE_DIR}/src M_SRC_DIR) 54 | include_directories(${PROJECT_SOURCE_DIR}/include) 55 | 56 | add_executable(yolosort ${M_SRC_DIR}) 57 | 58 | target_link_libraries(yolosort nvinfer cudart yolov5_trt) 59 | 60 | if(UNIX) 61 | add_definitions(-O2 -pthread) 62 | endif(UNIX) 63 | set(CMAKE_CXX_FLAGS "-std=c++0x") 64 | -------------------------------------------------------------------------------- /CMakeLists_deepsort-tensorrt_win10.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 2.6) 2 | project(deepsort) 3 | 4 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11") 5 | set(CMAKE_BUILD_TYPE Release) 6 | set(CUDA_BIN_PATH C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v11.1) #1.修改为自己的cuda路径 7 | set(TRT_DIR "D:\\lbq\\TensorRT-7.2.3.4") #2.修改为自己的tensorRT路径 8 | set(TRT_INCLUDE_DIRS ${TRT_DIR}\\include) #3.修改为自己的tensorRT头文件路径 9 | set(TRT_LIB_DIRS ${TRT_DIR}\\lib) #4.修改为自己的tensorRT库文件路径 10 | set(Eigen3_PATH D:\\lbq\\eigen) #5.修改为自己的Eigen路径 11 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 12 | find_package(OpenCV REQUIRED) 13 | find_package(CUDA REQUIRED) 14 | 15 | 16 | enable_language(CUDA) 17 | include_directories( 18 | ${CUDA_INCLUDE_DIRS} 19 | ${OpenCV_INCLUDE_DIRS} 20 | ${TRT_INCLUDE_DIRS} 21 | ${Eigen3_PATH} 22 | ${PROJECT_SOURCE_DIR}/include 23 | ) 24 | 25 | # tensorRT 26 | link_directories(${TRT_LIB_DIRS}) 27 | link_directories(${OpenCV_LIB_DIRS}) 28 | 29 | aux_source_directory(${PROJECT_SOURCE_DIR}/src SRC_DIR) 30 | 31 | # ===== deepsort ===== 32 | add_library(deepsort STATIC ${SRC_DIR}) 33 | target_link_libraries(deepsort 34 | ${CUDA_LIBRARIES} ${OpenCV_LIBS} 35 | cudart nvinfer nvonnxparser 36 | ) 37 | 38 | # ===== onnx2engine ===== 39 | add_executable(onnx2engine ${PROJECT_SOURCE_DIR}/onnx2engine.cpp) 40 | target_link_libraries(onnx2engine 41 | ${CUDA_LIBRARIES} 42 | cudart nvinfer nvonnxparser deepsort 43 | ) 44 | 45 | # ===== demo ===== 46 | add_executable(demo ${PROJECT_SOURCE_DIR}/demo.cpp) 47 | target_link_libraries(demo 48 | ${CUDA_LIBRARIES} ${OpenCV_LIBS} 49 | cudart nvinfer nvonnxparser deepsort 50 | ) 51 | 52 | 53 | -------------------------------------------------------------------------------- /CMakeLists_yolov5-deepsort-tensorrt_win10.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 2.6) 2 | 3 | project(yolosort) 4 | list(APPEND CUDA_NVCC_FLAGS "-std=c++11") 5 | set(CMAKE_CXX_FLAGS "-std=c++0x") 6 | find_package(OpenCV REQUIRED) 7 | add_definitions(-std=c++11) 8 | add_definitions(-DAPI_EXPORTS) 9 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 10 | set(CMAKE_CXX_STANDARD 11) 11 | set(CMAKE_BUILD_TYPE Release) 12 | 13 | #1.设置CUDA路径(set the path of cuda root ) 14 | set(CUDA_BIN_PATH C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v11.1) 15 | #2.设置tensorrt路径(set the path of tensorrt root) 16 | set(TRT_DIR "D:\\lbq\\TensorRT-7.2.3.4") 17 | #3.设置tensorrt头文件路径(set the path of tensorrt .h format file) 18 | set(TRT_INCLUDE_DIRS ${TRT_DIR}\\include) 19 | #4.设置tensorrt库文件路径(set the path of tensorrt .lib format file) 20 | set(TRT_LIB_DIRS ${TRT_DIR}\\lib) 21 | #5.设置Eigen路径,我用的是Eigen-3.3.9.zip(set the path of Eigen) 22 | set(Eigen3_PATH D:\\lbq\\eigen) 23 | #6.设置dirent路径(set the path of Dirent, download it from https://codeload.github.com/tronkko/dirent/zip/refs/heads/master) 24 | set(Dirent_INCLUDE_DIRS "D:\\lbq\\dirent\\include") #11 25 | # ==编译完成后,在vs2019打开的时候需要在链接器->命令行->其他选项->%(AdditionalOptions) /machine:x64 /FORCE:MULTIPLE 26 | 27 | find_package(CUDA REQUIRED) 28 | set(CUDA_NVCC_PLAGS ${CUDA_NVCC_PLAGS};-std=c++11;-g;-G;-gencode;arch=compute_53;code=sm_53) 29 | 30 | if(WIN32) 31 | enable_language(CUDA) 32 | endif(WIN32) 33 | 34 | include_directories(${PROJECT_SOURCE_DIR}/deepsort/include) 35 | # include and link dirs of cuda and tensorrt, you need adapt them if yours are different 36 | 37 | # tensorrt 38 | link_directories(${TRT_LIB_DIRS}) 39 | link_directories(${OpenCV_LIB_DIRS}) 40 | 41 | include_directories( 42 | ${CUDA_INCLUDE_DIRS} 43 | ${OpenCV_INCLUDE_DIRS} 44 | ${TRT_INCLUDE_DIRS} 45 | ${Eigen3_PATH} 46 | ${Dirent_INCLUDE_DIRS} 47 | ${PROJECT_SOURCE_DIR}/deepsort/include 48 | ) 49 | aux_source_directory(${PROJECT_SOURCE_DIR}/deepsort/src SRC_DIR) 50 | 51 | # ===== deepsort ===== 52 | add_library(deepsort STATIC ${SRC_DIR}) 53 | target_link_libraries(deepsort 54 | ${CUDA_LIBS} ${OpenCV_LIBS} 55 | cudart nvinfer nvonnxparser 56 | ) 57 | 58 | # ===== yolo ===== 59 | include_directories(${PROJECT_SOURCE_DIR}/yolo/include) 60 | aux_source_directory(${PROJECT_SOURCE_DIR}/yolo/src YOLO_SRC_DIR) 61 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -D_MWAITXINTRIN_H_INCLUDED") 62 | cuda_add_library(yolov5_trt STATIC ${PROJECT_SOURCE_DIR}/yolo/include/yololayer.cu ${PROJECT_SOURCE_DIR}/yolo/src/yolov5_lib.cpp) 63 | target_link_libraries(yolov5_trt nvinfer cudart deepsort) 64 | 65 | 66 | # ===== main ===== 67 | aux_source_directory(${PROJECT_SOURCE_DIR}/src M_SRC_DIR) 68 | include_directories(${PROJECT_SOURCE_DIR}/include) 69 | 70 | add_executable(yolosort ${M_SRC_DIR}) 71 | 72 | target_link_libraries(yolosort nvinfer cudart yolov5_trt) 73 | 74 | if(UNIX) 75 | add_definitions(-O2 -pthread) 76 | endif(UNIX) 77 | set(CMAKE_CXX_FLAGS "-std=c++0x") 78 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # A C++ implementation of Yolov5 and Deepsort in Jetson Xavier nx and Jetson nano 2 | [![MIT License](https://img.shields.io/badge/license-MIT-green)](https://opensource.org/licenses/MIT) 3 | [![GitHub stars](https://img.shields.io/github/stars/RichardoMrMu/yolov5-deepsort-tensorrt.svg?style=flat-square&logo=github&label=Stars&logoColor=white)](https://github.com/RichardoMrMu/yolov5-deepsort-tensorrt) 4 | 5 | This repository uses yolov5 and deepsort to follow human heads which can run in Jetson Xavier nx and Jetson nano. 6 | In Jetson Xavier Nx, it can achieve 10 FPS when images contain heads about 70+(you can try python version, when you use python version, you can find it very slow in Jetson Xavier nx , and Deepsort can cost nearly 1s). 7 | 8 | 9 | 10 | Thanks for [B.Q Long](https://github.com/lbq779660843), offer the windows cmakelists.txt. If you want run this rep in windows, you can use [CMakeLists_deepsort-tensorrt_win10.txt](https://github.com/RichardoMrMu/yolov5-deepsort-tensorrt/blob/main/CMakeLists_deepsort-tensorrt_win10.txt) and [CMakeLists_yolov5-deepsort-tensorrt_win10.txt](https://github.com/RichardoMrMu/yolov5-deepsort-tensorrt/blob/main/CMakeLists_yolov5-deepsort-tensorrt_win10.txt). 11 | 12 | You can see video play in [BILIBILI](https://www.bilibili.com/video/BV1nR4y1H7sR/), or [YOUTUBE](https://www.youtube.com/watch?v=29vEi-7mEic) and [YOUTUBE](https://youtu.be/mVq3pWNjb3E). 13 | 14 | ## Requirement 15 | 1. Jetson nano or Jetson Xavier nx 16 | 2. Jetpack 4.5.1 17 | 3. python3 with default(jetson nano or jetson xavier nx has default python3 with tensorrt 7.1.3.0 ) 18 | 4. tensorrt 7.1.3.0 19 | 5. torch 1.8.0 20 | 6. torchvision 0.9.0 21 | 7. torch2trt 0.3.0 22 | 8. onnx 1.4.1 23 | 9. opencv-python 4.5.3.56 24 | 10. protobuf 3.17.3 25 | 11. scipy 1.5.4 26 | 27 | 28 | if you have problem in this project, you can see this [artical](https://blog.csdn.net/weixin_42264234/article/details/120152117). 29 | 30 | ## Comming soon 31 | - [ ] Int8 . 32 | - [ ] IOU Tracking. 33 | - [ ] Faster and use less memory. 34 | ## Speed 35 | 36 | Whole process time from read image to finished deepsort (include every img preprocess and postprocess) 37 | and attention!!! the number of deepsort tracking is 70+, not single or 10-20 persons, is 70+. And all results can get in Jetson Xavier nx. 38 | | Backbone | before TensorRT without tracking |before TensorRT with tracking |TensortRT(detection)| TensorRT(detection + tracking) | FPS(detection + tracking) | 39 | | :-------------- | --------------- | ------------------ |--------------|------------------------------ | ------------------------- | 40 | | Yolov5s_416 | 100ms | 0.9s|10-15ms|100-150ms | 8 ~ 9 | 41 | | Yolov5s-640 | 120ms | 1s|18-20ms|100-150ms | 8 ~ 9 | 42 | 43 | ------ 44 | 45 | ## Build and Run 46 | 47 | ```shell 48 | git clone https://github.com/RichardoMrMu/yolov5-deepsort-tensorrt.git 49 | cd yolov5-deepsort-tensorrt 50 | // before you cmake and make, please change ./src/main.cpp char* yolo_engine = ""; char* sort_engine = ""; to your own path 51 | mkdir build 52 | cmake .. 53 | make 54 | ``` 55 | if you meet some errors in cmake and make, please see this [artical](https://blog.csdn.net/weixin_42264234/article/details/120152117) or see Attention. 56 | 57 | ## DataSet 58 | If you need to train your own model with head detection, you can use this [SCUT-HEAD](https://github.com/HCIILAB/SCUT-HEAD-Dataset-Release), this dataset has bbox with head and can download freely. 59 | 60 | ## Model 61 | You need two model, one is yolov5 model, for detection, generating from [tensorrtx](https://github.com/wang-xinyu/tensorrtx). And the other is deepsort model, for tracking. You should generate the model the same way. 62 | ### Generate yolov5 model 63 | For yolov5 detection model, I choose yolov5s, and choose `yolov5s.pt->yolov5s.wts->yolov5s.engine` 64 | Note that, used models can get from [yolov5](https://github.com/ultralytics/yolov5) and [deepsort](https://github.com/ZQPei/deep_sort_pytorch), and if you need to use your own model, you can follow the `Run Your Custom Model`. 65 | You can also see [tensorrtx official readme](https://github.com/wang-xinyu/tensorrtx/tree/master/yolov5) 66 | The following is deepsort.onnx and deesort.engine files, you can find in baiduyun and [https://github.com/RichardoMrMu/yolov5-deepsort-tensorrt/releases/tag/yolosort](https://github.com/RichardoMrMu/yolov5-deepsort-tensorrt/releases/tag/yolosort) 67 | | Model| Url | 68 | | :-------------- | --------------- | 69 | | 百度云 | [BaiduYun url](https://pan.baidu.com/s/1vpQFsD346lP64O1nhOilkw ) passwd:`z68e`| 70 | 71 | 1. Get yolov5 repository 72 | 73 | 74 | Note that, here uses the official pertained model.And I use yolov5-5, v5.0. So if you train your own model, please be sure your yolov5 code is v5.0. 75 | 76 | ```shell 77 | git clone -b v5.0 https://github.com/ultralytics/yolov5.git 78 | cd yolov5 79 | mkdir weights 80 | cd weights 81 | // download https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5s.pt 82 | wget https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5s.pt 83 | 84 | ``` 85 | 86 | 2. Get tensorrtx. 87 | 88 | For yolov5 v5.0, download .pt from [yolov5 release v5.0](https://github.com/ultralytics/yolov5/releases/tag/v5.0), `git clone -b v5.0 https://github.com/ultralytics/yolov5.git` and `git clone -b yolov5-v5.0 https://github.com/wang-xinyu/tensorrtx.git`, then follow how-to-run in [tensorrtx/yolov5-v5.0](https://github.com/wang-xinyu/tensorrtx/tree/yolov5-v5.0/yolov5). 89 | 90 | 91 | 3. Get xxx.wst model 92 | 93 | ```shell 94 | cp {tensorrtx}/yolov5/gen_wts.py {ultralytics}/yolov5/ 95 | cd yolov5 96 | python3 gen_wts.py -w ./weights/yolov5s.pt -o ./weights/yolov5s.wts 97 | // a file 'yolov5s.wts' will be generated. 98 | ``` 99 | You can get yolov5s.wts model in `yolov5/weights/` 100 | 101 | 4. Build tensorrtx/yolov5 and get tensorrt engine 102 | 103 | ```shell 104 | cd tensorrtx/yolov5 105 | // update CLASS_NUM in yololayer.h if your model is trained on custom dataset 106 | mkdir build 107 | cd build 108 | cp {ultralytics}/yolov5/yolov5s.wts {tensorrtx}/yolov5/build 109 | cmake .. 110 | make 111 | // yolov5s 112 | sudo ./yolov5 -s yolov5s.wts yolov5s.engine s 113 | // test your engine file 114 | sudo ./yolov5 -d yolov5s.engine ../samples 115 | ``` 116 | Then you get the yolov5s.engine, and you can put `yolov5s.engine` in My project. For example 117 | 118 | ```shell 119 | cd {yolov5-deepsort-tensorrt} 120 | mkdir resources 121 | cp {tensorrtx}/yolov5/build/yolov5s.engine {yolov5-deepsort-tensorrt}/resources 122 | ``` 123 | 124 | 5. Get deepsort engine file 125 | You can get deepsort pretrained model in this [drive url](https://drive.google.com/drive/folders/1xhG0kRH1EX5B9_Iz8gQJb7UNnn_riXi6) 126 | and ckpt.t7 is ok. 127 | ```shell 128 | git clone https://github.com/RichardoMrMu/deepsort-tensorrt.git 129 | // 根据github的说明 130 | cp {deepsort-tensorrt}/exportOnnx.py {deep_sort_pytorch}/ 131 | python3 exportOnnx.py 132 | mv {deep_sort_pytorch}/deepsort.onnx {deepsort-tensorrt}/resources 133 | cd {deepsort-tensorrt} 134 | mkdir build 135 | cd build 136 | cmake .. 137 | make 138 | ./onnx2engine ../resources/deepsort.onnx ../resources/deepsort.engine 139 | // test 140 | ./demo ../resource/deepsort.engine ../resources/track.txt 141 | ``` 142 | After all 5 step, you can get the yolov5s.engine and deepsort.engine. 143 | 144 | You may face some problems in getting yolov5s.engine and deepsort.engine, you can upload your issue in github or [csdn artical](https://blog.csdn.net/weixin_42264234/article/details/120152117). 145 |
146 | Different versions of yolov5 147 | 148 | Currently, tensorrt support yolov5 v1.0(yolov5s only), v2.0, v3.0, v3.1, v4.0 and v5.0. 149 | 150 | - For yolov5 v5.0, download .pt from [yolov5 release v5.0](https://github.com/ultralytics/yolov5/releases/tag/v5.0), `git clone -b v5.0 https://github.com/ultralytics/yolov5.git` and `git clone https://github.com/wang-xinyu/tensorrtx.git`, then follow how-to-run in current page. 151 | - For yolov5 v4.0, download .pt from [yolov5 release v4.0](https://github.com/ultralytics/yolov5/releases/tag/v4.0), `git clone -b v4.0 https://github.com/ultralytics/yolov5.git` and `git clone -b yolov5-v4.0 https://github.com/wang-xinyu/tensorrtx.git`, then follow how-to-run in [tensorrtx/yolov5-v4.0](https://github.com/wang-xinyu/tensorrtx/tree/yolov5-v4.0/yolov5). 152 | - For yolov5 v3.1, download .pt from [yolov5 release v3.1](https://github.com/ultralytics/yolov5/releases/tag/v3.1), `git clone -b v3.1 https://github.com/ultralytics/yolov5.git` and `git clone -b yolov5-v3.1 https://github.com/wang-xinyu/tensorrtx.git`, then follow how-to-run in [tensorrtx/yolov5-v3.1](https://github.com/wang-xinyu/tensorrtx/tree/yolov5-v3.1/yolov5). 153 | - For yolov5 v3.0, download .pt from [yolov5 release v3.0](https://github.com/ultralytics/yolov5/releases/tag/v3.0), `git clone -b v3.0 https://github.com/ultralytics/yolov5.git` and `git clone -b yolov5-v3.0 https://github.com/wang-xinyu/tensorrtx.git`, then follow how-to-run in [tensorrtx/yolov5-v3.0](https://github.com/wang-xinyu/tensorrtx/tree/yolov5-v3.0/yolov5). 154 | - For yolov5 v2.0, download .pt from [yolov5 release v2.0](https://github.com/ultralytics/yolov5/releases/tag/v2.0), `git clone -b v2.0 https://github.com/ultralytics/yolov5.git` and `git clone -b yolov5-v2.0 https://github.com/wang-xinyu/tensorrtx.git`, then follow how-to-run in [tensorrtx/yolov5-v2.0](https://github.com/wang-xinyu/tensorrtx/tree/yolov5-v2.0/yolov5). 155 | - For yolov5 v1.0, download .pt from [yolov5 release v1.0](https://github.com/ultralytics/yolov5/releases/tag/v1.0), `git clone -b v1.0 https://github.com/ultralytics/yolov5.git` and `git clone -b yolov5-v1.0 https://github.com/wang-xinyu/tensorrtx.git`, then follow how-to-run in [tensorrtx/yolov5-v1.0](https://github.com/wang-xinyu/tensorrtx/tree/yolov5-v1.0/yolov5). 156 |
157 | 158 |
159 | 160 | Config 161 | 162 | - Choose the model s/m/l/x/s6/m6/l6/x6 from command line arguments. 163 | - Input shape defined in yololayer.h 164 | - Number of classes defined in yololayer.h, **DO NOT FORGET TO ADAPT THIS, If using your own model** 165 | - INT8/FP16/FP32 can be selected by the macro in yolov5.cpp, **INT8 need more steps, pls follow `How to Run` first and then go the `INT8 Quantization` below** 166 | - GPU id can be selected by the macro in yolov5.cpp 167 | - NMS thresh in yolov5.cpp 168 | - BBox confidence thresh in yolov5.cpp 169 | - Batch size in yolov5.cpp 170 |
171 | 172 | ## Run Your Custom Model 173 | You may need train your own model and transfer your trained-model to tensorRT. So you can follow the following steps. 174 | 1. Train Custom Model 175 | You can follow the [official wiki](https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data 176 | ) to train your own model on your dataset. For example, I choose yolov5-s to train my model. 177 | 2. Transfer Custom Model 178 | Just like [tensorRT official guideline](https://github.com/wang-xinyu/tensorrtx/edit/master/yolov5/).When your follow ` Generate yolov5 model` to get yolov5 and tensorrt rep, next step is to transfer your pytorch model to tensorrt. 179 | Before this, you need to change yololayer.h file 20,21 and 22 line(CLASS_NUM，INPUT_H,INPUT_W) to your own parameters. 180 | 181 | ```shell 182 | // before 183 | static constexpr int CLASS_NUM = 80; // 20 184 | static constexpr int INPUT_H = 640; // 21 yolov5's input height and width must be divisible by 32. 185 | static constexpr int INPUT_W = 640; // 22 186 | 187 | // after 188 | // if your model is 2 classfication and image size is 416*416 189 | static constexpr int CLASS_NUM = 2; // 20 190 | static constexpr int INPUT_H = 416; // 21 yolov5's input height and width must be divisible by 32. 191 | static constexpr int INPUT_W = 416; // 22 192 | ``` 193 | 194 | ```shell 195 | cd {tensorrtx}/yolov5/ 196 | // update CLASS_NUM in yololayer.h if your model is trained on custom dataset 197 | 198 | mkdir build 199 | cd build 200 | cp {ultralytics}/yolov5/yolov5s.wts {tensorrtx}/yolov5/build 201 | cmake .. 202 | make 203 | sudo ./yolov5 -s [.wts] [.engine] [s/m/l/x/s6/m6/l6/x6 or c/c6 gd gw] // serialize model to plan file 204 | sudo ./yolov5 -d [.engine] [image folder] // deserialize and run inference, the images in [image folder] will be processed. 205 | // For example yolov5s 206 | sudo ./yolov5 -s yolov5s.wts yolov5s.engine s 207 | sudo ./yolov5 -d yolov5s.engine ../samples 208 | // For example Custom model with depth_multiple=0.17, width_multiple=0.25 in yolov5.yaml 209 | sudo ./yolov5 -s yolov5_custom.wts yolov5.engine c 0.17 0.25 210 | sudo ./yolov5 -d yolov5.engine ../samples 211 | ``` 212 | 213 | In this way, you can get your own tensorrt yolov5 model. Enjoy it! 214 | 215 | # Other Project 216 | - [yolov5-deepsort-tensorrt](https://github.com/RichardoMrMu/yolov5-deepsort-tensorrt) 217 | - [facial-emotion-recognition](https://github.com/RichardoMrMu/facial-emotion-recognition) 218 | - [yolov5-fire-smoke-detect-python](https://github.com/RichardoMrMu/yolov5-fire-smoke-detect-python) 219 | - [yolov5-helmet-detection](https://github.com/RichardoMrMu/yolov5-helmet-detection) 220 | - [yolov5-fire-smoke-detect](https://github.com/RichardoMrMu/yolov5-fire-smoke-detect) 221 | - [gazecapture-web gaze estimation](https://github.com/RichardoMrMu/gazecapture-web) 222 | - [yolov5-helmet-detection-python](https://github.com/RichardoMrMu/yolov5-helmet-detection-python) 223 | - [yolov5-reflective-clothes-detect-python](https://github.com/RichardoMrMu/yolov5-reflective-clothes-detect-python) 224 | - [yolov5-reflective-clothes-detect](https://github.com/RichardoMrMu/yolov5-reflective-clothes-detect) 225 | - [yolov5-smoking-detect](https://github.com/RichardoMrMu/yolov5-smoking-detect) 226 | - [yolov5-mask-detect](https://github.com/RichardoMrMu/yolov5-mask-detect) 227 | - [yolov5-mask-detection-python](https://github.com/RichardoMrMu/yolov5-mask-detection-python) 228 | - [yolov5-smoke-detection-python](https://github.com/RichardoMrMu/yolov5-smoke-detection-python) 229 | - [deepsort-tensorrt](https://github.com/RichardoMrMu/deepsort-tensorrt) 230 | 231 | 232 | 233 | 234 | 235 | -------------------------------------------------------------------------------- /assets/yolosort.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RichardoMrMu/yolov5-deepsort-tensorrt/33cd5cc49371fab325fe6970a4cc73f8f690af70/assets/yolosort.gif -------------------------------------------------------------------------------- /deepsort/include/datatype.h: -------------------------------------------------------------------------------- 1 | #ifndef DATATYPE_H 2 | #define DATATYPE_H 3 | 4 | typedef struct DetectBox { 5 | DetectBox(float x1=0, float y1=0, float x2=0, float y2=0, 6 | float confidence=0, float classID=-1, float trackID=-1) { 7 | this->x1 = x1; 8 | this->y1 = y1; 9 | this->x2 = x2; 10 | this->y2 = y2; 11 | this->confidence = confidence; 12 | this->classID = classID; 13 | this->trackID = trackID; 14 | } 15 | float x1, y1, x2, y2; 16 | float confidence; 17 | float classID; 18 | float trackID; 19 | } DetectBox; 20 | 21 | #endif // DATATYPE_H 22 | 23 | #ifndef DEEPSORTDATATYPE_H 24 | #define DEEPSORTDATATYPE_H 25 | 26 | #include 27 | #include 28 | #include 29 | #include 30 | typedef struct CLSCONF { 31 | CLSCONF() { 32 | this->cls = -1; 33 | this->conf = -1; 34 | } 35 | CLSCONF(int cls, float conf) { 36 | this->cls = cls; 37 | this->conf = conf; 38 | } 39 | int cls; 40 | float conf; 41 | } CLSCONF; 42 | 43 | typedef Eigen::Matrix DETECTBOX; 44 | typedef Eigen::Matrix DETECTBOXSS; 45 | typedef Eigen::Matrix FEATURE; 46 | typedef Eigen::Matrix FEATURESS; 47 | //typedef std::vector FEATURESS; 48 | 49 | //Kalmanfilter 50 | //typedef Eigen::Matrix KAL_FILTER; 51 | typedef Eigen::Matrix KAL_MEAN; 52 | typedef Eigen::Matrix KAL_COVA; 53 | typedef Eigen::Matrix KAL_HMEAN; 54 | typedef Eigen::Matrix KAL_HCOVA; 55 | using KAL_DATA = std::pair; 56 | using KAL_HDATA = std::pair; 57 | 58 | //main 59 | using RESULT_DATA = std::pair; 60 | 61 | //tracker: 62 | using TRACKER_DATA = std::pair; 63 | using MATCH_DATA = std::pair; 64 | typedef struct t{ 65 | std::vector matches; 66 | std::vector unmatched_tracks; 67 | std::vector unmatched_detections; 68 | }TRACHER_MATCHD; 69 | 70 | //linear_assignment: 71 | typedef Eigen::Matrix DYNAMICM; 72 | 73 | #endif //DEEPSORTDATATYPE_H -------------------------------------------------------------------------------- /deepsort/include/deepsort.h: -------------------------------------------------------------------------------- 1 | #ifndef DEEPSORT_H 2 | #define DEEPSORT_H 3 | 4 | #include 5 | #include 6 | #include "featuretensor.h" 7 | #include "tracker.h" 8 | #include "datatype.h" 9 | #include 10 | 11 | using std::vector; 12 | using nvinfer1::ILogger; 13 | 14 | class DeepSort { 15 | public: 16 | DeepSort(std::string modelPath, int batchSize, int featureDim, int gpuID, ILogger* gLogger); 17 | ~DeepSort(); 18 | 19 | public: 20 | void sort(cv::Mat& frame, vector& dets); 21 | 22 | private: 23 | void sort(cv::Mat& frame, DETECTIONS& detections); 24 | void sort(cv::Mat& frame, DETECTIONSV2& detectionsv2); 25 | void sort(vector& dets); 26 | void sort(DETECTIONS& detections); 27 | void init(); 28 | 29 | private: 30 | std::string enginePath; 31 | int batchSize; 32 | int featureDim; 33 | cv::Size imgShape; 34 | float confThres; 35 | float nmsThres; 36 | int maxBudget; 37 | float maxCosineDist; 38 | 39 | private: 40 | vector result; 41 | vector> results; 42 | tracker* objTracker; 43 | FeatureTensor* featureExtractor; 44 | ILogger* gLogger; 45 | int gpuID; 46 | }; 47 | 48 | #endif //deepsort.h 49 | -------------------------------------------------------------------------------- /deepsort/include/deepsortenginegenerator.h: -------------------------------------------------------------------------------- 1 | #ifndef DEEPSORT_ENGINE_GENERATOR_H 2 | #define DEEPSORT_ENGINE_GENERATOR_H 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | using namespace nvinfer1; 9 | 10 | const int IMG_HEIGHT = 128; 11 | const int IMG_WIDTH = 64; 12 | const int MAX_BATCH_SIZE = 128; 13 | const std::string INPUT_NAME("input"); 14 | 15 | class DeepSortEngineGenerator { 16 | public: 17 | DeepSortEngineGenerator(ILogger* gLogger); 18 | ~DeepSortEngineGenerator(); 19 | 20 | public: 21 | void setFP16(bool state); 22 | void createEngine(std::string onnxPath, std::string enginePath); 23 | 24 | private: 25 | std::string modelPath, engingPath; 26 | ILogger* gLogger; 27 | bool useFP16; 28 | }; 29 | 30 | #endif -------------------------------------------------------------------------------- /deepsort/include/featuretensor.h: -------------------------------------------------------------------------------- 1 | #ifndef FEATURETENSOR_H 2 | #define FEATURETENSOR_H 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include "model.hpp" 10 | #include "datatype.h" 11 | #include "cuda_runtime_api.h" 12 | 13 | using std::vector; 14 | using nvinfer1::ILogger; 15 | 16 | class FeatureTensor { 17 | public: 18 | FeatureTensor(const int maxBatchSize, const cv::Size imgShape, const int featureDim, int gpuID, ILogger* gLogger); 19 | ~FeatureTensor(); 20 | 21 | public: 22 | bool getRectsFeature(const cv::Mat& img, DETECTIONS& det); 23 | bool getRectsFeature(DETECTIONS& det); 24 | void loadEngine(std::string enginePath); 25 | void loadOnnx(std::string onnxPath); 26 | int getResult(float*& buffer); 27 | void doInference(vector& imgMats); 28 | 29 | private: 30 | void initResource(); 31 | void doInference(float* inputBuffer, float* outputBuffer); 32 | void mat2stream(vector& imgMats, float* stream); 33 | void stream2det(float* stream, DETECTIONS& det); 34 | 35 | private: 36 | nvinfer1::IRuntime* runtime; 37 | nvinfer1::ICudaEngine* engine; 38 | nvinfer1::IExecutionContext* context; 39 | const int maxBatchSize; 40 | const cv::Size imgShape; 41 | const int featureDim; 42 | 43 | private: 44 | int curBatchSize; 45 | const int inputStreamSize, outputStreamSize; 46 | bool initFlag; 47 | float* const inputBuffer; 48 | float* const outputBuffer; 49 | int inputIndex, outputIndex; 50 | void* buffers[2]; 51 | cudaStream_t cudaStream; 52 | // BGR format 53 | float means[3], std[3]; 54 | const std::string inputName, outputName; 55 | ILogger* gLogger; 56 | }; 57 | 58 | #endif 59 | -------------------------------------------------------------------------------- /deepsort/include/hungarianoper.h: -------------------------------------------------------------------------------- 1 | #ifndef HUNGARIANOPER_H 2 | #define HUNGARIANOPER_H 3 | 4 | #include "munkres.h" 5 | #include "datatype.h" 6 | 7 | 8 | class HungarianOper { 9 | public: 10 | static Eigen::Matrix Solve(const DYNAMICM &cost_matrix); 11 | }; 12 | 13 | #endif // HUNGARIANOPER_H 14 | -------------------------------------------------------------------------------- /deepsort/include/kalmanfilter.h: -------------------------------------------------------------------------------- 1 | #ifndef KALMANFILTER_H 2 | #define KALMANFILTER_H 3 | 4 | #include "datatype.h" 5 | 6 | class KalmanFilter { 7 | public: 8 | static const double chi2inv95[10]; 9 | KalmanFilter(); 10 | KAL_DATA initiate(const DETECTBOX& measurement); 11 | void predict(KAL_MEAN& mean, KAL_COVA& covariance); 12 | KAL_HDATA project(const KAL_MEAN& mean, const KAL_COVA& covariance); 13 | KAL_DATA update(const KAL_MEAN& mean, 14 | const KAL_COVA& covariance, 15 | const DETECTBOX& measurement); 16 | 17 | Eigen::Matrix gating_distance( 18 | const KAL_MEAN& mean, 19 | const KAL_COVA& covariance, 20 | const std::vector& measurements, 21 | bool only_position = false); 22 | 23 | private: 24 | Eigen::Matrix _motion_mat; 25 | Eigen::Matrix _update_mat; 26 | float _std_weight_position; 27 | float _std_weight_velocity; 28 | }; 29 | 30 | #endif // KALMANFILTER_H 31 | -------------------------------------------------------------------------------- /deepsort/include/linear_assignment.h: -------------------------------------------------------------------------------- 1 | #ifndef LINEAR_ASSIGNMENT_H 2 | #define LINEAR_ASSIGNMENT_H 3 | 4 | #include "datatype.h" 5 | #include "tracker.h" 6 | 7 | #define INFTY_COST 1e5 8 | class tracker; 9 | //for matching; 10 | class linear_assignment 11 | { 12 | linear_assignment(); 13 | linear_assignment(const linear_assignment& ); 14 | linear_assignment& operator=(const linear_assignment&); 15 | static linear_assignment* instance; 16 | 17 | public: 18 | static linear_assignment* getInstance(); 19 | TRACHER_MATCHD matching_cascade(tracker* distance_metric, 20 | tracker::GATED_METRIC_FUNC distance_metric_func, 21 | float max_distance, 22 | int cascade_depth, 23 | std::vector& tracks, 24 | const DETECTIONS& detections, 25 | std::vector &track_indices, 26 | std::vector detection_indices = std::vector()); 27 | TRACHER_MATCHD min_cost_matching( 28 | tracker* distance_metric, 29 | tracker::GATED_METRIC_FUNC distance_metric_func, 30 | float max_distance, 31 | std::vector& tracks, 32 | const DETECTIONS& detections, 33 | std::vector& track_indices, 34 | std::vector& detection_indices); 35 | DYNAMICM gate_cost_matrix( 36 | KalmanFilter* kf, 37 | DYNAMICM& cost_matrix, 38 | std::vector& tracks, 39 | const DETECTIONS& detections, 40 | const std::vector& track_indices, 41 | const std::vector& detection_indices, 42 | float gated_cost = INFTY_COST, 43 | bool only_position = false); 44 | }; 45 | 46 | #endif // LINEAR_ASSIGNMENT_H 47 | -------------------------------------------------------------------------------- /deepsort/include/logging.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. 3 | * 4 | * Licensed under the Apache License, Version 2.0 (the "License"); 5 | * you may not use this file except in compliance with the License. 6 | * You may obtain a copy of the License at 7 | * 8 | * http://www.apache.org/licenses/LICENSE-2.0 9 | * 10 | * Unless required by applicable law or agreed to in writing, software 11 | * distributed under the License is distributed on an "AS IS" BASIS, 12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | * See the License for the specific language governing permissions and 14 | * limitations under the License. 15 | */ 16 | 17 | #ifndef TENSORRT_LOGGING_H 18 | #define TENSORRT_LOGGING_H 19 | 20 | #include "NvInferRuntimeCommon.h" 21 | #include 22 | #include 23 | #include 24 | #include 25 | #include 26 | #include 27 | #include 28 | 29 | using Severity = nvinfer1::ILogger::Severity; 30 | 31 | class LogStreamConsumerBuffer : public std::stringbuf 32 | { 33 | public: 34 | LogStreamConsumerBuffer(std::ostream& stream, const std::string& prefix, bool shouldLog) 35 | : mOutput(stream) 36 | , mPrefix(prefix) 37 | , mShouldLog(shouldLog) 38 | { 39 | } 40 | 41 | LogStreamConsumerBuffer(LogStreamConsumerBuffer&& other) 42 | : mOutput(other.mOutput) 43 | { 44 | } 45 | 46 | ~LogStreamConsumerBuffer() 47 | { 48 | // std::streambuf::pbase() gives a pointer to the beginning of the buffered part of the output sequence 49 | // std::streambuf::pptr() gives a pointer to the current position of the output sequence 50 | // if the pointer to the beginning is not equal to the pointer to the current position, 51 | // call putOutput() to log the output to the stream 52 | if (pbase() != pptr()) 53 | { 54 | putOutput(); 55 | } 56 | } 57 | 58 | // synchronizes the stream buffer and returns 0 on success 59 | // synchronizing the stream buffer consists of inserting the buffer contents into the stream, 60 | // resetting the buffer and flushing the stream 61 | virtual int sync() 62 | { 63 | putOutput(); 64 | return 0; 65 | } 66 | 67 | void putOutput() 68 | { 69 | if (mShouldLog) 70 | { 71 | // prepend timestamp 72 | std::time_t timestamp = std::time(nullptr); 73 | tm* tm_local = std::localtime(×tamp); 74 | std::cout << "["; 75 | std::cout << std::setw(2) << std::setfill('0') << 1 + tm_local->tm_mon << "/"; 76 | std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_mday << "/"; 77 | std::cout << std::setw(4) << std::setfill('0') << 1900 + tm_local->tm_year << "-"; 78 | std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_hour << ":"; 79 | std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_min << ":"; 80 | std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_sec << "] "; 81 | // std::stringbuf::str() gets the string contents of the buffer 82 | // insert the buffer contents pre-appended by the appropriate prefix into the stream 83 | mOutput << mPrefix << str(); 84 | // set the buffer to empty 85 | str(""); 86 | // flush the stream 87 | mOutput.flush(); 88 | } 89 | } 90 | 91 | void setShouldLog(bool shouldLog) 92 | { 93 | mShouldLog = shouldLog; 94 | } 95 | 96 | private: 97 | std::ostream& mOutput; 98 | std::string mPrefix; 99 | bool mShouldLog; 100 | }; 101 | 102 | //! 103 | //! \class LogStreamConsumerBase 104 | //! \brief Convenience object used to initialize LogStreamConsumerBuffer before std::ostream in LogStreamConsumer 105 | //! 106 | class LogStreamConsumerBase 107 | { 108 | public: 109 | LogStreamConsumerBase(std::ostream& stream, const std::string& prefix, bool shouldLog) 110 | : mBuffer(stream, prefix, shouldLog) 111 | { 112 | } 113 | 114 | protected: 115 | LogStreamConsumerBuffer mBuffer; 116 | }; 117 | 118 | //! 119 | //! \class LogStreamConsumer 120 | //! \brief Convenience object used to facilitate use of C++ stream syntax when logging messages. 121 | //! Order of base classes is LogStreamConsumerBase and then std::ostream. 122 | //! This is because the LogStreamConsumerBase class is used to initialize the LogStreamConsumerBuffer member field 123 | //! in LogStreamConsumer and then the address of the buffer is passed to std::ostream. 124 | //! This is necessary to prevent the address of an uninitialized buffer from being passed to std::ostream. 125 | //! Please do not change the order of the parent classes. 126 | //! 127 | class LogStreamConsumer : protected LogStreamConsumerBase, public std::ostream 128 | { 129 | public: 130 | //! \brief Creates a LogStreamConsumer which logs messages with level severity. 131 | //! Reportable severity determines if the messages are severe enough to be logged. 132 | LogStreamConsumer(Severity reportableSeverity, Severity severity) 133 | : LogStreamConsumerBase(severityOstream(severity), severityPrefix(severity), severity <= reportableSeverity) 134 | , std::ostream(&mBuffer) // links the stream buffer with the stream 135 | , mShouldLog(severity <= reportableSeverity) 136 | , mSeverity(severity) 137 | { 138 | } 139 | 140 | LogStreamConsumer(LogStreamConsumer&& other) 141 | : LogStreamConsumerBase(severityOstream(other.mSeverity), severityPrefix(other.mSeverity), other.mShouldLog) 142 | , std::ostream(&mBuffer) // links the stream buffer with the stream 143 | , mShouldLog(other.mShouldLog) 144 | , mSeverity(other.mSeverity) 145 | { 146 | } 147 | 148 | void setReportableSeverity(Severity reportableSeverity) 149 | { 150 | mShouldLog = mSeverity <= reportableSeverity; 151 | mBuffer.setShouldLog(mShouldLog); 152 | } 153 | 154 | private: 155 | static std::ostream& severityOstream(Severity severity) 156 | { 157 | return severity >= Severity::kINFO ? std::cout : std::cerr; 158 | } 159 | 160 | static std::string severityPrefix(Severity severity) 161 | { 162 | switch (severity) 163 | { 164 | case Severity::kINTERNAL_ERROR: return "[F] "; 165 | case Severity::kERROR: return "[E] "; 166 | case Severity::kWARNING: return "[W] "; 167 | case Severity::kINFO: return "[I] "; 168 | case Severity::kVERBOSE: return "[V] "; 169 | default: assert(0); return ""; 170 | } 171 | } 172 | 173 | bool mShouldLog; 174 | Severity mSeverity; 175 | }; 176 | 177 | //! \class Logger 178 | //! 179 | //! \brief Class which manages logging of TensorRT tools and samples 180 | //! 181 | //! \details This class provides a common interface for TensorRT tools and samples to log information to the console, 182 | //! and supports logging two types of messages: 183 | //! 184 | //! - Debugging messages with an associated severity (info, warning, error, or internal error/fatal) 185 | //! - Test pass/fail messages 186 | //! 187 | //! The advantage of having all samples use this class for logging as opposed to emitting directly to stdout/stderr is 188 | //! that the logic for controlling the verbosity and formatting of sample output is centralized in one location. 189 | //! 190 | //! In the future, this class could be extended to support dumping test results to a file in some standard format 191 | //! (for example, JUnit XML), and providing additional metadata (e.g. timing the duration of a test run). 192 | //! 193 | //! TODO: For backwards compatibility with existing samples, this class inherits directly from the nvinfer1::ILogger 194 | //! interface, which is problematic since there isn't a clean separation between messages coming from the TensorRT 195 | //! library and messages coming from the sample. 196 | //! 197 | //! In the future (once all samples are updated to use Logger::getTRTLogger() to access the ILogger) we can refactor the 198 | //! class to eliminate the inheritance and instead make the nvinfer1::ILogger implementation a member of the Logger 199 | //! object. 200 | 201 | class Logger : public nvinfer1::ILogger 202 | { 203 | public: 204 | Logger(Severity severity = Severity::kWARNING) 205 | : mReportableSeverity(severity) 206 | { 207 | } 208 | 209 | //! 210 | //! \enum TestResult 211 | //! \brief Represents the state of a given test 212 | //! 213 | enum class TestResult 214 | { 215 | kRUNNING, //!< The test is running 216 | kPASSED, //!< The test passed 217 | kFAILED, //!< The test failed 218 | kWAIVED //!< The test was waived 219 | }; 220 | 221 | //! 222 | //! \brief Forward-compatible method for retrieving the nvinfer::ILogger associated with this Logger 223 | //! \return The nvinfer1::ILogger associated with this Logger 224 | //! 225 | //! TODO Once all samples are updated to use this method to register the logger with TensorRT, 226 | //! we can eliminate the inheritance of Logger from ILogger 227 | //! 228 | nvinfer1::ILogger& getTRTLogger() 229 | { 230 | return *this; 231 | } 232 | 233 | //! 234 | //! \brief Implementation of the nvinfer1::ILogger::log() virtual method 235 | //! 236 | //! Note samples should not be calling this function directly; it will eventually go away once we eliminate the 237 | //! inheritance from nvinfer1::ILogger 238 | //! 239 | void log(Severity severity, const char* msg) override 240 | { 241 | LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl; 242 | } 243 | 244 | //! 245 | //! \brief Method for controlling the verbosity of logging output 246 | //! 247 | //! \param severity The logger will only emit messages that have severity of this level or higher. 248 | //! 249 | void setReportableSeverity(Severity severity) 250 | { 251 | mReportableSeverity = severity; 252 | } 253 | 254 | //! 255 | //! \brief Opaque handle that holds logging information for a particular test 256 | //! 257 | //! This object is an opaque handle to information used by the Logger to print test results. 258 | //! The sample must call Logger::defineTest() in order to obtain a TestAtom that can be used 259 | //! with Logger::reportTest{Start,End}(). 260 | //! 261 | class TestAtom 262 | { 263 | public: 264 | TestAtom(TestAtom&&) = default; 265 | 266 | private: 267 | friend class Logger; 268 | 269 | TestAtom(bool started, const std::string& name, const std::string& cmdline) 270 | : mStarted(started) 271 | , mName(name) 272 | , mCmdline(cmdline) 273 | { 274 | } 275 | 276 | bool mStarted; 277 | std::string mName; 278 | std::string mCmdline; 279 | }; 280 | 281 | //! 282 | //! \brief Define a test for logging 283 | //! 284 | //! \param[in] name The name of the test. This should be a string starting with 285 | //! "TensorRT" and containing dot-separated strings containing 286 | //! the characters [A-Za-z0-9_]. 287 | //! For example, "TensorRT.sample_googlenet" 288 | //! \param[in] cmdline The command line used to reproduce the test 289 | // 290 | //! \return a TestAtom that can be used in Logger::reportTest{Start,End}(). 291 | //! 292 | static TestAtom defineTest(const std::string& name, const std::string& cmdline) 293 | { 294 | return TestAtom(false, name, cmdline); 295 | } 296 | 297 | //! 298 | //! \brief A convenience overloaded version of defineTest() that accepts an array of command-line arguments 299 | //! as input 300 | //! 301 | //! \param[in] name The name of the test 302 | //! \param[in] argc The number of command-line arguments 303 | //! \param[in] argv The array of command-line arguments (given as C strings) 304 | //! 305 | //! \return a TestAtom that can be used in Logger::reportTest{Start,End}(). 306 | static TestAtom defineTest(const std::string& name, int argc, char const* const* argv) 307 | { 308 | auto cmdline = genCmdlineString(argc, argv); 309 | return defineTest(name, cmdline); 310 | } 311 | 312 | //! 313 | //! \brief Report that a test has started. 314 | //! 315 | //! \pre reportTestStart() has not been called yet for the given testAtom 316 | //! 317 | //! \param[in] testAtom The handle to the test that has started 318 | //! 319 | static void reportTestStart(TestAtom& testAtom) 320 | { 321 | reportTestResult(testAtom, TestResult::kRUNNING); 322 | assert(!testAtom.mStarted); 323 | testAtom.mStarted = true; 324 | } 325 | 326 | //! 327 | //! \brief Report that a test has ended. 328 | //! 329 | //! \pre reportTestStart() has been called for the given testAtom 330 | //! 331 | //! \param[in] testAtom The handle to the test that has ended 332 | //! \param[in] result The result of the test. Should be one of TestResult::kPASSED, 333 | //! TestResult::kFAILED, TestResult::kWAIVED 334 | //! 335 | static void reportTestEnd(const TestAtom& testAtom, TestResult result) 336 | { 337 | assert(result != TestResult::kRUNNING); 338 | assert(testAtom.mStarted); 339 | reportTestResult(testAtom, result); 340 | } 341 | 342 | static int reportPass(const TestAtom& testAtom) 343 | { 344 | reportTestEnd(testAtom, TestResult::kPASSED); 345 | return EXIT_SUCCESS; 346 | } 347 | 348 | static int reportFail(const TestAtom& testAtom) 349 | { 350 | reportTestEnd(testAtom, TestResult::kFAILED); 351 | return EXIT_FAILURE; 352 | } 353 | 354 | static int reportWaive(const TestAtom& testAtom) 355 | { 356 | reportTestEnd(testAtom, TestResult::kWAIVED); 357 | return EXIT_SUCCESS; 358 | } 359 | 360 | static int reportTest(const TestAtom& testAtom, bool pass) 361 | { 362 | return pass ? reportPass(testAtom) : reportFail(testAtom); 363 | } 364 | 365 | Severity getReportableSeverity() const 366 | { 367 | return mReportableSeverity; 368 | } 369 | 370 | private: 371 | //! 372 | //! \brief returns an appropriate string for prefixing a log message with the given severity 373 | //! 374 | static const char* severityPrefix(Severity severity) 375 | { 376 | switch (severity) 377 | { 378 | case Severity::kINTERNAL_ERROR: return "[F] "; 379 | case Severity::kERROR: return "[E] "; 380 | case Severity::kWARNING: return "[W] "; 381 | case Severity::kINFO: return "[I] "; 382 | case Severity::kVERBOSE: return "[V] "; 383 | default: assert(0); return ""; 384 | } 385 | } 386 | 387 | //! 388 | //! \brief returns an appropriate string for prefixing a test result message with the given result 389 | //! 390 | static const char* testResultString(TestResult result) 391 | { 392 | switch (result) 393 | { 394 | case TestResult::kRUNNING: return "RUNNING"; 395 | case TestResult::kPASSED: return "PASSED"; 396 | case TestResult::kFAILED: return "FAILED"; 397 | case TestResult::kWAIVED: return "WAIVED"; 398 | default: assert(0); return ""; 399 | } 400 | } 401 | 402 | //! 403 | //! \brief returns an appropriate output stream (cout or cerr) to use with the given severity 404 | //! 405 | static std::ostream& severityOstream(Severity severity) 406 | { 407 | return severity >= Severity::kINFO ? std::cout : std::cerr; 408 | } 409 | 410 | //! 411 | //! \brief method that implements logging test results 412 | //! 413 | static void reportTestResult(const TestAtom& testAtom, TestResult result) 414 | { 415 | severityOstream(Severity::kINFO) << "&&&& " << testResultString(result) << " " << testAtom.mName << " # " 416 | << testAtom.mCmdline << std::endl; 417 | } 418 | 419 | //! 420 | //! \brief generate a command line string from the given (argc, argv) values 421 | //! 422 | static std::string genCmdlineString(int argc, char const* const* argv) 423 | { 424 | std::stringstream ss; 425 | for (int i = 0; i < argc; i++) 426 | { 427 | if (i > 0) 428 | ss << " "; 429 | ss << argv[i]; 430 | } 431 | return ss.str(); 432 | } 433 | 434 | Severity mReportableSeverity; 435 | }; 436 | 437 | namespace 438 | { 439 | 440 | //! 441 | //! \brief produces a LogStreamConsumer object that can be used to log messages of severity kVERBOSE 442 | //! 443 | //! Example usage: 444 | //! 445 | //! LOG_VERBOSE(logger) << "hello world" << std::endl; 446 | //! 447 | inline LogStreamConsumer LOG_VERBOSE(const Logger& logger) 448 | { 449 | return LogStreamConsumer(logger.getReportableSeverity(), Severity::kVERBOSE); 450 | } 451 | 452 | //! 453 | //! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINFO 454 | //! 455 | //! Example usage: 456 | //! 457 | //! LOG_INFO(logger) << "hello world" << std::endl; 458 | //! 459 | inline LogStreamConsumer LOG_INFO(const Logger& logger) 460 | { 461 | return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINFO); 462 | } 463 | 464 | //! 465 | //! \brief produces a LogStreamConsumer object that can be used to log messages of severity kWARNING 466 | //! 467 | //! Example usage: 468 | //! 469 | //! LOG_WARN(logger) << "hello world" << std::endl; 470 | //! 471 | inline LogStreamConsumer LOG_WARN(const Logger& logger) 472 | { 473 | return LogStreamConsumer(logger.getReportableSeverity(), Severity::kWARNING); 474 | } 475 | 476 | //! 477 | //! \brief produces a LogStreamConsumer object that can be used to log messages of severity kERROR 478 | //! 479 | //! Example usage: 480 | //! 481 | //! LOG_ERROR(logger) << "hello world" << std::endl; 482 | //! 483 | inline LogStreamConsumer LOG_ERROR(const Logger& logger) 484 | { 485 | return LogStreamConsumer(logger.getReportableSeverity(), Severity::kERROR); 486 | } 487 | 488 | //! 489 | //! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINTERNAL_ERROR 490 | // ("fatal" severity) 491 | //! 492 | //! Example usage: 493 | //! 494 | //! LOG_FATAL(logger) << "hello world" << std::endl; 495 | //! 496 | inline LogStreamConsumer LOG_FATAL(const Logger& logger) 497 | { 498 | return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINTERNAL_ERROR); 499 | } 500 | 501 | } // anonymous namespace 502 | 503 | #endif // TENSORRT_LOGGING_H 504 | -------------------------------------------------------------------------------- /deepsort/include/matrix.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2007 John Weaver 3 | * 4 | * This program is free software; you can redistribute it and/or modify 5 | * it under the terms of the GNU General Public License as published by 6 | * the Free Software Foundation; either version 2 of the License, or 7 | * (at your option) any later version. 8 | * 9 | * This program is distributed in the hope that it will be useful, 10 | * but WITHOUT ANY WARRANTY; without even the implied warranty of 11 | * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 12 | * GNU General Public License for more details. 13 | * 14 | * You should have received a copy of the GNU General Public License 15 | * along with this program; if not, write to the Free Software 16 | * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA 17 | */ 18 | 19 | #ifndef _MATRIX_H_ 20 | #define _MATRIX_H_ 21 | 22 | #include 23 | #include 24 | #include 25 | #include 26 | #include 27 | #include 28 | 29 | #define XYZMIN(x, y) (x)<(y)?(x):(y) 30 | #define XYZMAX(x, y) (x)>(y)?(x):(y) 31 | 32 | template 33 | class Matrix { 34 | public: 35 | Matrix(){ 36 | m_rows = 0; 37 | m_columns = 0; 38 | m_matrix = nullptr; 39 | } 40 | Matrix(const size_t rows, const size_t columns) { 41 | m_matrix = nullptr; 42 | resize(rows, columns); 43 | } 44 | Matrix(const std::initializer_list> init) { 45 | m_matrix = nullptr; 46 | m_rows = init.size(); 47 | if ( m_rows == 0 ) { 48 | m_columns = 0; 49 | } else { 50 | m_columns = init.begin()->size(); 51 | if ( m_columns > 0 ) { 52 | resize(m_rows, m_columns); 53 | } 54 | } 55 | 56 | size_t i = 0, j; 57 | for ( auto row = init.begin() ; row != init.end() ; ++row, ++i ) { 58 | assert ( row->size() == m_columns && "All rows must have the same number of columns." ); 59 | j = 0; 60 | for ( auto value = row->begin() ; value != row->end() ; ++value, ++j ) { 61 | m_matrix[i][j] = *value; 62 | } 63 | } 64 | } 65 | Matrix(const Matrix &other) { 66 | if ( other.m_matrix != nullptr ) { 67 | // copy arrays 68 | m_matrix = nullptr; 69 | resize(other.m_rows, other.m_columns); 70 | for ( size_t i = 0 ; i < m_rows ; i++ ) { 71 | for ( size_t j = 0 ; j < m_columns ; j++ ) { 72 | m_matrix[i][j] = other.m_matrix[i][j]; 73 | } 74 | } 75 | } else { 76 | m_matrix = nullptr; 77 | m_rows = 0; 78 | m_columns = 0; 79 | } 80 | } 81 | Matrix & operator= (const Matrix &other){ 82 | if ( other.m_matrix != nullptr ) { 83 | // copy arrays 84 | resize(other.m_rows, other.m_columns); 85 | for ( size_t i = 0 ; i < m_rows ; i++ ) { 86 | for ( size_t j = 0 ; j < m_columns ; j++ ) { 87 | m_matrix[i][j] = other.m_matrix[i][j]; 88 | } 89 | } 90 | } else { 91 | // free arrays 92 | for ( size_t i = 0 ; i < m_columns ; i++ ) { 93 | delete [] m_matrix[i]; 94 | } 95 | 96 | delete [] m_matrix; 97 | 98 | m_matrix = nullptr; 99 | m_rows = 0; 100 | m_columns = 0; 101 | } 102 | 103 | return *this; 104 | } 105 | ~Matrix(){ 106 | if ( m_matrix != nullptr ) { 107 | // free arrays 108 | for ( size_t i = 0 ; i < m_rows ; i++ ) { 109 | delete [] m_matrix[i]; 110 | } 111 | 112 | delete [] m_matrix; 113 | } 114 | m_matrix = nullptr; 115 | } 116 | // all operations modify the matrix in-place. 117 | void resize(const size_t rows, const size_t columns, const T default_value = 0) { 118 | assert ( rows > 0 && columns > 0 && "Columns and rows must exist." ); 119 | 120 | if ( m_matrix == nullptr ) { 121 | // alloc arrays 122 | m_matrix = new T*[rows]; // rows 123 | for ( size_t i = 0 ; i < rows ; i++ ) { 124 | m_matrix[i] = new T[columns]; // columns 125 | } 126 | 127 | m_rows = rows; 128 | m_columns = columns; 129 | clear(); 130 | } else { 131 | // save array pointer 132 | T **new_matrix; 133 | // alloc new arrays 134 | new_matrix = new T*[rows]; // rows 135 | for ( size_t i = 0 ; i < rows ; i++ ) { 136 | new_matrix[i] = new T[columns]; // columns 137 | for ( size_t j = 0 ; j < columns ; j++ ) { 138 | new_matrix[i][j] = default_value; 139 | } 140 | } 141 | 142 | // copy data from saved pointer to new arrays 143 | size_t minrows = XYZMIN(rows, m_rows); 144 | size_t mincols = XYZMIN(columns, m_columns); 145 | for ( size_t x = 0 ; x < minrows ; x++ ) { 146 | for ( size_t y = 0 ; y < mincols ; y++ ) { 147 | new_matrix[x][y] = m_matrix[x][y]; 148 | } 149 | } 150 | 151 | // delete old arrays 152 | if ( m_matrix != nullptr ) { 153 | for ( size_t i = 0 ; i < m_rows ; i++ ) { 154 | delete [] m_matrix[i]; 155 | } 156 | 157 | delete [] m_matrix; 158 | } 159 | 160 | m_matrix = new_matrix; 161 | } 162 | 163 | m_rows = rows; 164 | m_columns = columns; 165 | } 166 | void clear() { 167 | assert( m_matrix != nullptr ); 168 | 169 | for ( size_t i = 0 ; i < m_rows ; i++ ) { 170 | for ( size_t j = 0 ; j < m_columns ; j++ ) { 171 | m_matrix[i][j] = 0; 172 | } 173 | } 174 | } 175 | T& operator () (const size_t x, const size_t y) { 176 | assert ( x < m_rows ); 177 | assert ( y < m_columns ); 178 | assert ( m_matrix != nullptr ); 179 | return m_matrix[x][y]; 180 | } 181 | 182 | const T& operator () (const size_t x, const size_t y) const { 183 | assert ( x < m_rows ); 184 | assert ( y < m_columns ); 185 | assert ( m_matrix != nullptr ); 186 | return m_matrix[x][y]; 187 | } 188 | const T mmin() const { 189 | assert( m_matrix != nullptr ); 190 | assert ( m_rows > 0 ); 191 | assert ( m_columns > 0 ); 192 | T min = m_matrix[0][0]; 193 | 194 | for ( size_t i = 0 ; i < m_rows ; i++ ) { 195 | for ( size_t j = 0 ; j < m_columns ; j++ ) { 196 | min = std::min(min, m_matrix[i][j]); 197 | } 198 | } 199 | 200 | return min; 201 | } 202 | 203 | const T mmax() const { 204 | assert( m_matrix != nullptr ); 205 | assert ( m_rows > 0 ); 206 | assert ( m_columns > 0 ); 207 | T max = m_matrix[0][0]; 208 | 209 | for ( size_t i = 0 ; i < m_rows ; i++ ) { 210 | for ( size_t j = 0 ; j < m_columns ; j++ ) { 211 | max = std::max(max, m_matrix[i][j]); 212 | } 213 | } 214 | 215 | return max; 216 | } 217 | inline size_t minsize() { return ((m_rows < m_columns) ? m_rows : m_columns); } 218 | inline size_t columns() const { return m_columns;} 219 | inline size_t rows() const { return m_rows;} 220 | 221 | friend std::ostream& operator<<(std::ostream& os, const Matrix &matrix) 222 | { 223 | os << "Matrix:" << std::endl; 224 | for (size_t row = 0 ; row < matrix.rows() ; row++ ) 225 | { 226 | for (size_t col = 0 ; col < matrix.columns() ; col++ ) 227 | { 228 | os.width(8); 229 | os << matrix(row, col) << ","; 230 | } 231 | os << std::endl; 232 | } 233 | return os; 234 | } 235 | 236 | private: 237 | T **m_matrix; 238 | size_t m_rows; 239 | size_t m_columns; 240 | }; 241 | 242 | //#ifndef USE_EXPORT_KEYWORD 243 | //#include "matrix.cpp" 244 | ////#define export /*export*/ 245 | //#endif 246 | 247 | #endif /* !defined(_MATRIX_H_) */ 248 | 249 | -------------------------------------------------------------------------------- /deepsort/include/model.hpp: -------------------------------------------------------------------------------- 1 | #ifndef MODEL_HPP 2 | #define MODEL_HPP 3 | 4 | #include 5 | #include "datatype.h" 6 | 7 | 8 | // * Each rect's data structure. 9 | // * tlwh: topleft point & (w,h) 10 | // * confidence: detection confidence. 11 | // * feature: the rect's 256d feature. 12 | // */ 13 | 14 | const float kRatio=0.5; 15 | enum DETECTBOX_IDX {IDX_X = 0, IDX_Y, IDX_W, IDX_H }; 16 | 17 | class DETECTION_ROW { 18 | public: 19 | DETECTBOX tlwh; 20 | float confidence; 21 | FEATURE feature; 22 | DETECTBOX to_xyah() const { 23 | //(centerx, centery, ration, h) 24 | DETECTBOX ret = tlwh; 25 | ret(0, IDX_X) += (ret(0, IDX_W)*kRatio); 26 | ret(0, IDX_Y) += (ret(0, IDX_H)*kRatio); 27 | ret(0, IDX_W) /= ret(0, IDX_H); 28 | return ret; 29 | } 30 | DETECTBOX to_tlbr() const { 31 | //(x,y,xx,yy) 32 | DETECTBOX ret = tlwh; 33 | ret(0, IDX_X) += ret(0, IDX_W); 34 | ret(0, IDX_Y) += ret(0, IDX_H); 35 | return ret; 36 | } 37 | }; 38 | 39 | typedef std::vector DETECTIONS; 40 | typedef std::pair, DETECTIONS> DETECTIONSV2; 41 | 42 | #endif // MODEL_HPP 43 | -------------------------------------------------------------------------------- /deepsort/include/munkres.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2007 John Weaver 3 | * Copyright (c) 2015 Miroslav Krajicek 4 | * 5 | * This program is free software; you can redistribute it and/or modify 6 | * it under the terms of the GNU General Public License as published by 7 | * the Free Software Foundation; either version 2 of the License, or 8 | * (at your option) any later version. 9 | * 10 | * This program is distributed in the hope that it will be useful, 11 | * but WITHOUT ANY WARRANTY; without even the implied warranty of 12 | * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 13 | * GNU General Public License for more details. 14 | * 15 | * You should have received a copy of the GNU General Public License 16 | * along with this program; if not, write to the Free Software 17 | * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA 18 | */ 19 | 20 | #if !defined(_MUNKRES_H_) 21 | #define _MUNKRES_H_ 22 | 23 | #include "matrix.h" 24 | 25 | #include 26 | #include 27 | #include 28 | #include 29 | #include 30 | 31 | 32 | template class Munkres 33 | { 34 | static constexpr int NORMAL = 0; 35 | static constexpr int STAR = 1; 36 | static constexpr int PRIME = 2; 37 | public: 38 | 39 | /* 40 | * 41 | * Linear assignment problem solution 42 | * [modifies matrix in-place.] 43 | * matrix(row,col): row major format assumed. 44 | * 45 | * Assignments are remaining 0 values 46 | * (extra 0 values are replaced with -1) 47 | * 48 | */ 49 | void solve(Matrix &m) { 50 | const size_t rows = m.rows(), 51 | columns = m.columns(), 52 | size = XYZMAX(rows, columns); 53 | 54 | #ifdef DEBUG 55 | std::cout << "Munkres input: " << m << std::endl; 56 | #endif 57 | 58 | // Copy input matrix 59 | this->matrix = m; 60 | 61 | if ( rows != columns ) { 62 | // If the input matrix isn't square, make it square 63 | // and fill the empty values with the largest value present 64 | // in the matrix. 65 | matrix.resize(size, size, matrix.mmax()); 66 | } 67 | 68 | 69 | // STAR == 1 == starred, PRIME == 2 == primed 70 | mask_matrix.resize(size, size); 71 | 72 | row_mask = new bool[size]; 73 | col_mask = new bool[size]; 74 | for ( size_t i = 0 ; i < size ; i++ ) { 75 | row_mask[i] = false; 76 | } 77 | 78 | for ( size_t i = 0 ; i < size ; i++ ) { 79 | col_mask[i] = false; 80 | } 81 | 82 | // Prepare the matrix values... 83 | 84 | // If there were any infinities, replace them with a value greater 85 | // than the maximum value in the matrix. 86 | replace_infinites(matrix); 87 | 88 | minimize_along_direction(matrix, rows >= columns); 89 | minimize_along_direction(matrix, rows < columns); 90 | 91 | // Follow the steps 92 | int step = 1; 93 | while ( step ) { 94 | switch ( step ) { 95 | case 1: 96 | step = step1(); 97 | // step is always 2 98 | break; 99 | case 2: 100 | step = step2(); 101 | // step is always either 0 or 3 102 | break; 103 | case 3: 104 | step = step3(); 105 | // step in [3, 4, 5] 106 | break; 107 | case 4: 108 | step = step4(); 109 | // step is always 2 110 | break; 111 | case 5: 112 | step = step5(); 113 | // step is always 3 114 | break; 115 | } 116 | } 117 | 118 | // Store results 119 | for ( size_t row = 0 ; row < size ; row++ ) { 120 | for ( size_t col = 0 ; col < size ; col++ ) { 121 | if ( mask_matrix(row, col) == STAR ) { 122 | matrix(row, col) = 0; 123 | } else { 124 | matrix(row, col) = -1; 125 | } 126 | } 127 | } 128 | 129 | #ifdef DEBUG 130 | std::cout << "Munkres output: " << matrix << std::endl; 131 | #endif 132 | // Remove the excess rows or columns that we added to fit the 133 | // input to a square matrix. 134 | matrix.resize(rows, columns); 135 | 136 | m = matrix; 137 | 138 | delete [] row_mask; 139 | delete [] col_mask; 140 | } 141 | 142 | static void replace_infinites(Matrix &matrix) { 143 | const size_t rows = matrix.rows(), 144 | columns = matrix.columns(); 145 | //assert( rows > 0 && columns > 0 ); 146 | double max = matrix(0, 0); 147 | constexpr auto infinity = std::numeric_limits::infinity(); 148 | 149 | // Find the greatest value in the matrix that isn't infinity. 150 | for ( size_t row = 0 ; row < rows ; row++ ) { 151 | for ( size_t col = 0 ; col < columns ; col++ ) { 152 | if ( matrix(row, col) != infinity ) { 153 | if ( max == infinity ) { 154 | max = matrix(row, col); 155 | } else { 156 | max = XYZMAX(max, matrix(row, col)); 157 | } 158 | } 159 | } 160 | } 161 | 162 | // a value higher than the maximum value present in the matrix. 163 | if ( max == infinity ) { 164 | // This case only occurs when all values are infinite. 165 | max = 0; 166 | } else { 167 | max++; 168 | } 169 | 170 | for ( size_t row = 0 ; row < rows ; row++ ) { 171 | for ( size_t col = 0 ; col < columns ; col++ ) { 172 | if ( matrix(row, col) == infinity ) { 173 | matrix(row, col) = max; 174 | } 175 | } 176 | } 177 | 178 | } 179 | static void minimize_along_direction(Matrix &matrix, const bool over_columns) { 180 | const size_t outer_size = over_columns ? matrix.columns() : matrix.rows(), 181 | inner_size = over_columns ? matrix.rows() : matrix.columns(); 182 | 183 | // Look for a minimum value to subtract from all values along 184 | // the "outer" direction. 185 | for ( size_t i = 0 ; i < outer_size ; i++ ) { 186 | double min = over_columns ? matrix(0, i) : matrix(i, 0); 187 | 188 | // As long as the current minimum is greater than zero, 189 | // keep looking for the minimum. 190 | // Start at one because we already have the 0th value in min. 191 | for ( size_t j = 1 ; j < inner_size && min > 0 ; j++ ) { 192 | min = XYZMIN( 193 | min, 194 | over_columns ? matrix(j, i) : matrix(i, j)); 195 | } 196 | 197 | if ( min > 0 ) { 198 | for ( size_t j = 0 ; j < inner_size ; j++ ) { 199 | if ( over_columns ) { 200 | matrix(j, i) -= min; 201 | } else { 202 | matrix(i, j) -= min; 203 | } 204 | } 205 | } 206 | } 207 | } 208 | 209 | private: 210 | 211 | inline bool find_uncovered_in_matrix(const double item, size_t &row, size_t &col) const { 212 | const size_t rows = matrix.rows(), 213 | columns = matrix.columns(); 214 | 215 | for ( row = 0 ; row < rows ; row++ ) { 216 | if ( !row_mask[row] ) { 217 | for ( col = 0 ; col < columns ; col++ ) { 218 | if ( !col_mask[col] ) { 219 | if ( matrix(row,col) == item ) { 220 | return true; 221 | } 222 | } 223 | } 224 | } 225 | } 226 | 227 | return false; 228 | } 229 | 230 | bool pair_in_list(const std::pair &needle, const std::list > &haystack) { 231 | for ( std::list >::const_iterator i = haystack.begin() ; i != haystack.end() ; i++ ) { 232 | if ( needle == *i ) { 233 | return true; 234 | } 235 | } 236 | 237 | return false; 238 | } 239 | 240 | int step1() { 241 | const size_t rows = matrix.rows(), 242 | columns = matrix.columns(); 243 | 244 | for ( size_t row = 0 ; row < rows ; row++ ) { 245 | for ( size_t col = 0 ; col < columns ; col++ ) { 246 | if ( 0 == matrix(row, col) ) { 247 | for ( size_t nrow = 0 ; nrow < row ; nrow++ ) 248 | if ( STAR == mask_matrix(nrow,col) ) 249 | goto next_column; 250 | 251 | mask_matrix(row,col) = STAR; 252 | goto next_row; 253 | } 254 | next_column:; 255 | } 256 | next_row:; 257 | } 258 | 259 | return 2; 260 | } 261 | 262 | int step2() { 263 | const size_t rows = matrix.rows(), 264 | columns = matrix.columns(); 265 | size_t covercount = 0; 266 | 267 | for ( size_t row = 0 ; row < rows ; row++ ) 268 | for ( size_t col = 0 ; col < columns ; col++ ) 269 | if ( STAR == mask_matrix(row, col) ) { 270 | col_mask[col] = true; 271 | covercount++; 272 | } 273 | 274 | if ( covercount >= matrix.minsize() ) { 275 | #ifdef DEBUG 276 | std::cout << "Final cover count: " << covercount << std::endl; 277 | #endif 278 | return 0; 279 | } 280 | 281 | #ifdef DEBUG 282 | std::cout << "Munkres matrix has " << covercount << " of " << matrix.minsize() << " Columns covered:" << std::endl; 283 | std::cout << matrix << std::endl; 284 | #endif 285 | 286 | 287 | return 3; 288 | } 289 | 290 | int step3() { 291 | /* 292 | Main Zero Search 293 | 294 | 1. Find an uncovered Z in the distance matrix and prime it. If no such zero exists, go to Step 5 295 | 2. If No Z* exists in the row of the Z', go to Step 4. 296 | 3. If a Z* exists, cover this row and uncover the column of the Z*. Return to Step 3.1 to find a new Z 297 | */ 298 | if ( find_uncovered_in_matrix(0, saverow, savecol) ) { 299 | mask_matrix(saverow,savecol) = PRIME; // prime it. 300 | } else { 301 | return 5; 302 | } 303 | 304 | for ( size_t ncol = 0 ; ncol < matrix.columns() ; ncol++ ) { 305 | if ( mask_matrix(saverow,ncol) == STAR ) { 306 | row_mask[saverow] = true; //cover this row and 307 | col_mask[ncol] = false; // uncover the column containing the starred zero 308 | return 3; // repeat 309 | } 310 | } 311 | 312 | return 4; // no starred zero in the row containing this primed zero 313 | } 314 | 315 | int step4() { 316 | const size_t rows = matrix.rows(), 317 | columns = matrix.columns(); 318 | 319 | // seq contains pairs of row/column values where we have found 320 | // either a star or a prime that is part of the ``alternating sequence``. 321 | std::list > seq; 322 | // use saverow, savecol from step 3. 323 | std::pair z0(saverow, savecol); 324 | seq.insert(seq.end(), z0); 325 | 326 | // We have to find these two pairs: 327 | std::pair z1(-1, -1); 328 | std::pair z2n(-1, -1); 329 | 330 | size_t row, col = savecol; 331 | /* 332 | Increment Set of Starred Zeros 333 | 334 | 1. Construct the ``alternating sequence'' of primed and starred zeros: 335 | 336 | Z0 : Unpaired Z' from Step 4.2 337 | Z1 : The Z* in the column of Z0 338 | Z[2N] : The Z' in the row of Z[2N-1], if such a zero exists 339 | Z[2N+1] : The Z* in the column of Z[2N] 340 | 341 | The sequence eventually terminates with an unpaired Z' = Z[2N] for some N. 342 | */ 343 | bool madepair; 344 | do { 345 | madepair = false; 346 | for ( row = 0 ; row < rows ; row++ ) { 347 | if ( mask_matrix(row,col) == STAR ) { 348 | z1.first = row; 349 | z1.second = col; 350 | if ( pair_in_list(z1, seq) ) { 351 | continue; 352 | } 353 | 354 | madepair = true; 355 | seq.insert(seq.end(), z1); 356 | break; 357 | } 358 | } 359 | 360 | if ( !madepair ) 361 | break; 362 | 363 | madepair = false; 364 | 365 | for ( col = 0 ; col < columns ; col++ ) { 366 | if ( mask_matrix(row, col) == PRIME ) { 367 | z2n.first = row; 368 | z2n.second = col; 369 | if ( pair_in_list(z2n, seq) ) { 370 | continue; 371 | } 372 | madepair = true; 373 | seq.insert(seq.end(), z2n); 374 | break; 375 | } 376 | } 377 | } while ( madepair ); 378 | 379 | for ( std::list >::iterator i = seq.begin() ; 380 | i != seq.end() ; 381 | i++ ) { 382 | // 2. Unstar each starred zero of the sequence. 383 | if ( mask_matrix(i->first,i->second) == STAR ) 384 | mask_matrix(i->first,i->second) = NORMAL; 385 | 386 | // 3. Star each primed zero of the sequence, 387 | // thus increasing the number of starred zeros by one. 388 | if ( mask_matrix(i->first,i->second) == PRIME ) 389 | mask_matrix(i->first,i->second) = STAR; 390 | } 391 | 392 | // 4. Erase all primes, uncover all columns and rows, 393 | for ( size_t row = 0 ; row < mask_matrix.rows() ; row++ ) { 394 | for ( size_t col = 0 ; col < mask_matrix.columns() ; col++ ) { 395 | if ( mask_matrix(row,col) == PRIME ) { 396 | mask_matrix(row,col) = NORMAL; 397 | } 398 | } 399 | } 400 | 401 | for ( size_t i = 0 ; i < rows ; i++ ) { 402 | row_mask[i] = false; 403 | } 404 | 405 | for ( size_t i = 0 ; i < columns ; i++ ) { 406 | col_mask[i] = false; 407 | } 408 | 409 | // and return to Step 2. 410 | return 2; 411 | } 412 | 413 | int step5() { 414 | const size_t rows = matrix.rows(), 415 | columns = matrix.columns(); 416 | /* 417 | New Zero Manufactures 418 | 419 | 1. Let h be the smallest uncovered entry in the (modified) distance matrix. 420 | 2. Add h to all covered rows. 421 | 3. Subtract h from all uncovered columns 422 | 4. Return to Step 3, without altering stars, primes, or covers. 423 | */ 424 | double h = 100000;//xyzoylz std::numeric_limits::max(); 425 | for ( size_t row = 0 ; row < rows ; row++ ) { 426 | if ( !row_mask[row] ) { 427 | for ( size_t col = 0 ; col < columns ; col++ ) { 428 | if ( !col_mask[col] ) { 429 | if ( h > matrix(row, col) && matrix(row, col) != 0 ) { 430 | h = matrix(row, col); 431 | } 432 | } 433 | } 434 | } 435 | } 436 | 437 | for ( size_t row = 0 ; row < rows ; row++ ) { 438 | if ( row_mask[row] ) { 439 | for ( size_t col = 0 ; col < columns ; col++ ) { 440 | matrix(row, col) += h; 441 | } 442 | } 443 | } 444 | 445 | for ( size_t col = 0 ; col < columns ; col++ ) { 446 | if ( !col_mask[col] ) { 447 | for ( size_t row = 0 ; row < rows ; row++ ) { 448 | matrix(row, col) -= h; 449 | } 450 | } 451 | } 452 | 453 | return 3; 454 | } 455 | 456 | Matrix mask_matrix; 457 | Matrix matrix; 458 | bool *row_mask; 459 | bool *col_mask; 460 | size_t saverow = 0, savecol = 0; 461 | }; 462 | 463 | 464 | #endif /* !defined(_MUNKRES_H_) */ 465 | -------------------------------------------------------------------------------- /deepsort/include/nn_matching.h: -------------------------------------------------------------------------------- 1 | #ifndef NN_MATCHING_H 2 | #define NN_MATCHING_H 3 | 4 | #include "datatype.h" 5 | 6 | #include 7 | 8 | //A tool to calculate distance; 9 | class NearNeighborDisMetric{ 10 | public: 11 | enum METRIC_TYPE{euclidean=1, cosine}; 12 | NearNeighborDisMetric(METRIC_TYPE metric, 13 | float matching_threshold, 14 | int budget); 15 | DYNAMICM distance(const FEATURESS& features, const std::vector &targets); 16 | // void partial_fit(FEATURESS& features, std::vector targets, std::vector active_targets); 17 | void partial_fit(std::vector& tid_feats, std::vector& active_targets); 18 | float mating_threshold; 19 | 20 | private: 21 | typedef Eigen::VectorXf (NearNeighborDisMetric::*PTRFUN)(const FEATURESS&, const FEATURESS&); 22 | Eigen::VectorXf _nncosine_distance(const FEATURESS& x, const FEATURESS& y); 23 | Eigen::VectorXf _nneuclidean_distance(const FEATURESS& x, const FEATURESS& y); 24 | 25 | Eigen::MatrixXf _pdist(const FEATURESS& x, const FEATURESS& y); 26 | Eigen::MatrixXf _cosine_distance(const FEATURESS & a, const FEATURESS& b, bool data_is_normalized = false); 27 | private: 28 | PTRFUN _metric; 29 | int budget; 30 | std::map samples; 31 | }; 32 | 33 | #endif // NN_MATCHING_H 34 | -------------------------------------------------------------------------------- /deepsort/include/track.h: -------------------------------------------------------------------------------- 1 | #ifndef TRACK_H 2 | #define TRACK_H 3 | 4 | 5 | #include "kalmanfilter.h" 6 | #include "datatype.h" 7 | #include "model.hpp" 8 | 9 | class Track 10 | { 11 | /*""" 12 | A single target track with state space `(x, y, a, h)` and associated 13 | velocities, where `(x, y)` is the center of the bounding box, `a` is the 14 | aspect ratio and `h` is the height. 15 | 16 | Parameters 17 | ---------- 18 | mean : ndarray 19 | Mean vector of the initial state distribution. 20 | covariance : ndarray 21 | Covariance matrix of the initial state distribution. 22 | track_id : int 23 | A unique track identifier. 24 | n_init : int 25 | Number of consecutive detections before the track is confirmed. The 26 | track state is set to `Deleted` if a miss occurs within the first 27 | `n_init` frames. 28 | max_age : int 29 | The maximum number of consecutive misses before the track state is 30 | set to `Deleted`. 31 | feature : Optional[ndarray] 32 | Feature vector of the detection this track originates from. If not None, 33 | this feature is added to the `features` cache. 34 | 35 | Attributes 36 | ---------- 37 | mean : ndarray 38 | Mean vector of the initial state distribution. 39 | covariance : ndarray 40 | Covariance matrix of the initial state distribution. 41 | track_id : int 42 | A unique track identifier. 43 | hits : int 44 | Total number of measurement updates. 45 | age : int 46 | Total number of frames since first occurance. 47 | time_since_update : int 48 | Total number of frames since last measurement update. 49 | state : TrackState 50 | The current track state. 51 | features : List[ndarray] 52 | A cache of features. On each measurement update, the associated feature 53 | vector is added to this list. 54 | 55 | """*/ 56 | enum TrackState {Tentative = 1, Confirmed, Deleted}; 57 | 58 | public: 59 | Track(KAL_MEAN& mean, KAL_COVA& covariance, int track_id, 60 | int n_init, int max_age, const FEATURE& feature); 61 | Track(KAL_MEAN& mean, KAL_COVA& covariance, int track_id, 62 | int n_init, int max_age, const FEATURE& feature, int cls, float conf); 63 | void predit(KalmanFilter* kf); 64 | void update(KalmanFilter* const kf, const DETECTION_ROW &detection); 65 | void update(KalmanFilter* const kf, const DETECTION_ROW & detection, CLSCONF pair_det); 66 | void mark_missed(); 67 | bool is_confirmed(); 68 | bool is_deleted(); 69 | bool is_tentative(); 70 | DETECTBOX to_tlwh(); 71 | int time_since_update; 72 | int track_id; 73 | FEATURESS features; 74 | KAL_MEAN mean; 75 | KAL_COVA covariance; 76 | 77 | int hits; 78 | int age; 79 | int _n_init; 80 | int _max_age; 81 | TrackState state; 82 | 83 | int cls; 84 | float conf; 85 | private: 86 | void featuresAppendOne(const FEATURE& f); 87 | }; 88 | 89 | #endif // TRACK_H 90 | -------------------------------------------------------------------------------- /deepsort/include/tracker.h: -------------------------------------------------------------------------------- 1 | #ifndef TRACKER_H 2 | #define TRACKER_H 3 | 4 | 5 | #include 6 | 7 | #include "kalmanfilter.h" 8 | #include "track.h" 9 | #include "model.hpp" 10 | 11 | using namespace std; 12 | 13 | class NearNeighborDisMetric; 14 | 15 | class tracker 16 | { 17 | public: 18 | NearNeighborDisMetric* metric; 19 | float max_iou_distance; 20 | int max_age; 21 | int n_init; 22 | 23 | KalmanFilter* kf; 24 | 25 | int _next_idx; 26 | public: 27 | std::vector tracks; 28 | tracker(/*NearNeighborDisMetric* metric,*/ 29 | float max_cosine_distance, int nn_budget, 30 | float max_iou_distance = 0.7, 31 | int max_age = 200, int n_init=20); 32 | void predict(); 33 | void update(const DETECTIONS& detections); 34 | void update(const DETECTIONSV2& detectionsv2); 35 | typedef DYNAMICM (tracker::* GATED_METRIC_FUNC)( 36 | std::vector& tracks, 37 | const DETECTIONS& dets, 38 | const std::vector& track_indices, 39 | const std::vector& detection_indices); 40 | private: 41 | void _match(const DETECTIONS& detections, TRACHER_MATCHD& res); 42 | void _initiate_track(const DETECTION_ROW& detection); 43 | void _initiate_track(const DETECTION_ROW& detection, CLSCONF clsConf); 44 | public: 45 | DYNAMICM gated_matric( 46 | std::vector& tracks, 47 | const DETECTIONS& dets, 48 | const std::vector& track_indices, 49 | const std::vector& detection_indices); 50 | DYNAMICM iou_cost( 51 | std::vector& tracks, 52 | const DETECTIONS& dets, 53 | const std::vector& track_indices, 54 | const std::vector& detection_indices); 55 | Eigen::VectorXf iou(DETECTBOX& bbox, 56 | DETECTBOXSS &candidates); 57 | }; 58 | 59 | #endif // TRACKER_H 60 | -------------------------------------------------------------------------------- /deepsort/src/deepsort.cpp: -------------------------------------------------------------------------------- 1 | #include "deepsort.h" 2 | 3 | DeepSort::DeepSort(std::string modelPath, int batchSize, int featureDim, int gpuID, ILogger* gLogger) { 4 | this->gpuID = gpuID; 5 | this->enginePath = modelPath; 6 | this->batchSize = batchSize; 7 | this->featureDim = featureDim; 8 | this->imgShape = cv::Size(64, 128); 9 | this->maxBudget = 100; 10 | this->maxCosineDist = 0.2; 11 | this->gLogger = gLogger; 12 | init(); 13 | } 14 | 15 | void DeepSort::init() { 16 | objTracker = new tracker(maxCosineDist, maxBudget); 17 | featureExtractor = new FeatureTensor(batchSize, imgShape, featureDim, gpuID, gLogger); 18 | int ret = enginePath.find(".onnx"); 19 | if (ret != -1) 20 | featureExtractor->loadOnnx(enginePath); 21 | else 22 | featureExtractor->loadEngine(enginePath); 23 | } 24 | 25 | DeepSort::~DeepSort() { 26 | delete objTracker; 27 | } 28 | 29 | void DeepSort::sort(cv::Mat& frame, vector& dets) { 30 | // preprocess Mat -> DETECTION 31 | DETECTIONS detections; 32 | vector clsConf; 33 | 34 | for (DetectBox i : dets) { 35 | DETECTBOX box(i.x1, i.y1, i.x2-i.x1, i.y2-i.y1); 36 | DETECTION_ROW d; 37 | d.tlwh = box; 38 | d.confidence = i.confidence; 39 | detections.push_back(d); 40 | clsConf.push_back(CLSCONF((int)i.classID, i.confidence)); 41 | } 42 | result.clear(); 43 | results.clear(); 44 | if (detections.size() > 0) { 45 | DETECTIONSV2 detectionsv2 = make_pair(clsConf, detections); 46 | sort(frame, detectionsv2); 47 | } 48 | // postprocess DETECTION -> Mat 49 | dets.clear(); 50 | for (auto r : result) { 51 | DETECTBOX i = r.second; 52 | DetectBox b(i(0), i(1), i(2)+i(0), i(3)+i(1), 1.); 53 | b.trackID = (float)r.first; 54 | dets.push_back(b); 55 | } 56 | for (int i = 0; i < results.size(); ++i) { 57 | CLSCONF c = results[i].first; 58 | dets[i].classID = c.cls; 59 | dets[i].confidence = c.conf; 60 | } 61 | } 62 | 63 | 64 | void DeepSort::sort(cv::Mat& frame, DETECTIONS& detections) { 65 | bool flag = featureExtractor->getRectsFeature(frame, detections); 66 | if (flag) { 67 | objTracker->predict(); 68 | objTracker->update(detections); 69 | //result.clear(); 70 | for (Track& track : objTracker->tracks) { 71 | if (!track.is_confirmed() || track.time_since_update > 1) 72 | continue; 73 | result.push_back(make_pair(track.track_id, track.to_tlwh())); 74 | } 75 | } 76 | } 77 | 78 | void DeepSort::sort(cv::Mat& frame, DETECTIONSV2& detectionsv2) { 79 | std::vector& clsConf = detectionsv2.first; 80 | DETECTIONS& detections = detectionsv2.second; 81 | bool flag = featureExtractor->getRectsFeature(frame, detections); 82 | if (flag) { 83 | objTracker->predict(); 84 | objTracker->update(detectionsv2); 85 | result.clear(); 86 | results.clear(); 87 | for (Track& track : objTracker->tracks) { 88 | if (!track.is_confirmed() || track.time_since_update > 1) 89 | continue; 90 | result.push_back(make_pair(track.track_id, track.to_tlwh())); 91 | results.push_back(make_pair(CLSCONF(track.cls, track.conf) ,track.to_tlwh())); 92 | } 93 | } 94 | } 95 | 96 | void DeepSort::sort(vector& dets) { 97 | DETECTIONS detections; 98 | for (DetectBox i : dets) { 99 | DETECTBOX box(i.x1, i.y1, i.x2-i.x1, i.y2-i.y1); 100 | DETECTION_ROW d; 101 | d.tlwh = box; 102 | d.confidence = i.confidence; 103 | detections.push_back(d); 104 | } 105 | if (detections.size() > 0) 106 | sort(detections); 107 | dets.clear(); 108 | for (auto r : result) { 109 | DETECTBOX i = r.second; 110 | DetectBox b(i(0), i(1), i(2), i(3), 1.); 111 | b.trackID = r.first; 112 | dets.push_back(b); 113 | } 114 | } 115 | 116 | void DeepSort::sort(DETECTIONS& detections) { 117 | bool flag = featureExtractor->getRectsFeature(detections); 118 | if (flag) { 119 | objTracker->predict(); 120 | objTracker->update(detections); 121 | result.clear(); 122 | for (Track& track : objTracker->tracks) { 123 | if (!track.is_confirmed() || track.time_since_update > 1) 124 | continue; 125 | result.push_back(make_pair(track.track_id, track.to_tlwh())); 126 | } 127 | } 128 | } 129 | -------------------------------------------------------------------------------- /deepsort/src/deepsortenginegenerator.cpp: -------------------------------------------------------------------------------- 1 | #include "deepsortenginegenerator.h" 2 | #include "assert.h" 3 | #include 4 | #include 5 | 6 | DeepSortEngineGenerator::DeepSortEngineGenerator(ILogger* gLogger) { 7 | this->gLogger = gLogger; 8 | } 9 | 10 | DeepSortEngineGenerator::~DeepSortEngineGenerator() { 11 | 12 | } 13 | 14 | void DeepSortEngineGenerator::setFP16(bool state) { 15 | this->useFP16 = state; 16 | } 17 | 18 | void DeepSortEngineGenerator::createEngine(std::string onnxPath, std::string enginePath) { 19 | // Load onnx model 20 | auto builder = createInferBuilder(*gLogger); 21 | assert(builder != nullptr); 22 | const auto explicitBatch = 1U << static_cast(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); 23 | auto network = builder->createNetworkV2(explicitBatch); 24 | assert(network != nullptr); 25 | auto config = builder->createBuilderConfig(); 26 | assert(config != nullptr); 27 | 28 | auto profile = builder->createOptimizationProfile(); 29 | Dims dims = Dims4{1, 3, IMG_HEIGHT, IMG_WIDTH}; 30 | profile->setDimensions(INPUT_NAME.c_str(), 31 | OptProfileSelector::kMIN, Dims4{1, dims.d[1], dims.d[2], dims.d[3]}); 32 | profile->setDimensions(INPUT_NAME.c_str(), 33 | OptProfileSelector::kOPT, Dims4{MAX_BATCH_SIZE, dims.d[1], dims.d[2], dims.d[3]}); 34 | profile->setDimensions(INPUT_NAME.c_str(), 35 | OptProfileSelector::kMAX, Dims4{MAX_BATCH_SIZE, dims.d[1], dims.d[2], dims.d[3]}); 36 | config->addOptimizationProfile(profile); 37 | 38 | nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, *gLogger); 39 | assert(parser != nullptr); 40 | auto parsed = parser->parseFromFile(onnxPath.c_str(), static_cast(ILogger::Severity::kWARNING)); 41 | assert(parsed); 42 | if (useFP16) config->setFlag(BuilderFlag::kFP16); 43 | config->setMaxWorkspaceSize(1 << 20); 44 | ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); 45 | 46 | // Serialize model and save engine 47 | IHostMemory* modelStream = engine->serialize(); 48 | std::string serializeStr; 49 | std::ofstream serializeOutputStream; 50 | serializeStr.resize(modelStream->size()); 51 | memcpy((void*)serializeStr.data(), modelStream->data(), modelStream->size()); 52 | serializeOutputStream.open(enginePath); 53 | serializeOutputStream << serializeStr; 54 | serializeOutputStream.close(); 55 | } -------------------------------------------------------------------------------- /deepsort/src/featuretensor.cpp: -------------------------------------------------------------------------------- 1 | #include "featuretensor.h" 2 | #include 3 | 4 | using namespace nvinfer1; 5 | 6 | #define INPUTSTREAM_SIZE (maxBatchSize*3*imgShape.area()) 7 | #define OUTPUTSTREAM_SIZE (maxBatchSize*featureDim) 8 | 9 | FeatureTensor::FeatureTensor(const int maxBatchSize, const cv::Size imgShape, const int featureDim, int gpuID, ILogger* gLogger) 10 | : maxBatchSize(maxBatchSize), imgShape(imgShape), featureDim(featureDim), 11 | inputStreamSize(INPUTSTREAM_SIZE), outputStreamSize(OUTPUTSTREAM_SIZE), 12 | inputBuffer(new float[inputStreamSize]), outputBuffer(new float[outputStreamSize]), 13 | inputName("input"), outputName("output") { 14 | cudaSetDevice(gpuID); 15 | this->gLogger = gLogger; 16 | runtime = nullptr; 17 | engine = nullptr; 18 | context = nullptr; 19 | 20 | means[0] = 0.485, means[1] = 0.456, means[2] = 0.406; 21 | std[0] = 0.229, std[1] = 0.224, std[2] = 0.225; 22 | 23 | initFlag = false; 24 | } 25 | 26 | FeatureTensor::~FeatureTensor() { 27 | delete [] inputBuffer; 28 | delete [] outputBuffer; 29 | if (initFlag) { 30 | // cudaStreamSynchronize(cudaStream); 31 | cudaStreamDestroy(cudaStream); 32 | cudaFree(buffers[inputIndex]); 33 | cudaFree(buffers[outputIndex]); 34 | } 35 | } 36 | 37 | bool FeatureTensor::getRectsFeature(const cv::Mat& img, DETECTIONS& det) { 38 | std::vector mats; 39 | for (auto& dbox : det) { 40 | cv::Rect rect = cv::Rect(int(dbox.tlwh(0)), int(dbox.tlwh(1)), 41 | int(dbox.tlwh(2)), int(dbox.tlwh(3))); 42 | rect.x -= (rect.height * 0.5 - rect.width) * 0.5; 43 | rect.width = rect.height * 0.5; 44 | rect.x = (rect.x >= 0 ? rect.x : 0); 45 | rect.y = (rect.y >= 0 ? rect.y : 0); 46 | rect.width = (rect.x + rect.width <= img.cols ? rect.width : (img.cols - rect.x)); 47 | rect.height = (rect.y + rect.height <= img.rows ? rect.height : (img.rows - rect.y)); 48 | cv::Mat tempMat = img(rect).clone(); 49 | cv::resize(tempMat, tempMat, imgShape); 50 | mats.push_back(tempMat); 51 | } 52 | doInference(mats); 53 | // decode output to det 54 | stream2det(outputBuffer, det); 55 | return true; 56 | } 57 | 58 | bool FeatureTensor::getRectsFeature(DETECTIONS& det) { 59 | return true; 60 | } 61 | 62 | void FeatureTensor::loadEngine(std::string enginePath) { 63 | // Deserialize model 64 | runtime = createInferRuntime(*gLogger); 65 | assert(runtime != nullptr); 66 | std::ifstream engineStream(enginePath); 67 | std::string engineCache(""); 68 | while (engineStream.peek() != EOF) { 69 | std::stringstream buffer; 70 | buffer << engineStream.rdbuf(); 71 | engineCache.append(buffer.str()); 72 | } 73 | engineStream.close(); 74 | engine = runtime->deserializeCudaEngine(engineCache.data(), engineCache.size(), nullptr); 75 | assert(engine != nullptr); 76 | context = engine->createExecutionContext(); 77 | assert(context != nullptr); 78 | initResource(); 79 | } 80 | 81 | void FeatureTensor::loadOnnx(std::string onnxPath) { 82 | auto builder = createInferBuilder(*gLogger); 83 | assert(builder != nullptr); 84 | const auto explicitBatch = 1U << static_cast(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); 85 | auto network = builder->createNetworkV2(explicitBatch); 86 | assert(network != nullptr); 87 | auto config = builder->createBuilderConfig(); 88 | assert(config != nullptr); 89 | 90 | auto profile = builder->createOptimizationProfile(); 91 | Dims dims = Dims4{1, 3, imgShape.height, imgShape.width}; 92 | profile->setDimensions(inputName.c_str(), 93 | OptProfileSelector::kMIN, Dims4{1, dims.d[1], dims.d[2], dims.d[3]}); 94 | profile->setDimensions(inputName.c_str(), 95 | OptProfileSelector::kOPT, Dims4{maxBatchSize, dims.d[1], dims.d[2], dims.d[3]}); 96 | profile->setDimensions(inputName.c_str(), 97 | OptProfileSelector::kMAX, Dims4{maxBatchSize, dims.d[1], dims.d[2], dims.d[3]}); 98 | config->addOptimizationProfile(profile); 99 | 100 | nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, *gLogger); 101 | assert(parser != nullptr); 102 | auto parsed = parser->parseFromFile(onnxPath.c_str(), static_cast(ILogger::Severity::kWARNING)); 103 | assert(parsed); 104 | config->setMaxWorkspaceSize(1 << 20); 105 | engine = builder->buildEngineWithConfig(*network, *config); 106 | assert(engine != nullptr); 107 | context = engine->createExecutionContext(); 108 | assert(context != nullptr); 109 | initResource(); 110 | } 111 | 112 | int FeatureTensor::getResult(float*& buffer) { 113 | if (buffer != nullptr) 114 | delete buffer; 115 | int curStreamSize = curBatchSize*featureDim; 116 | buffer = new float[curStreamSize]; 117 | for (int i = 0; i < curStreamSize; ++i) { 118 | buffer[i] = outputBuffer[i]; 119 | } 120 | return curStreamSize; 121 | } 122 | 123 | void FeatureTensor::doInference(vector& imgMats) { 124 | mat2stream(imgMats, inputBuffer); 125 | doInference(inputBuffer, outputBuffer); 126 | } 127 | 128 | void FeatureTensor::initResource() { 129 | inputIndex = engine->getBindingIndex(inputName.c_str()); 130 | outputIndex = engine->getBindingIndex(outputName.c_str()); 131 | // Create CUDA stream 132 | cudaStreamCreate(&cudaStream); 133 | buffers[inputIndex] = inputBuffer; 134 | buffers[outputIndex] = outputBuffer; 135 | 136 | // Malloc CUDA memory 137 | cudaMalloc(&buffers[inputIndex], inputStreamSize * sizeof(float)); 138 | cudaMalloc(&buffers[outputIndex], outputStreamSize * sizeof(float)); 139 | 140 | initFlag = true; 141 | } 142 | 143 | void FeatureTensor::doInference(float* inputBuffer, float* outputBuffer) { 144 | cudaMemcpyAsync(buffers[inputIndex], inputBuffer, inputStreamSize * sizeof(float), cudaMemcpyHostToDevice, cudaStream); 145 | Dims4 inputDims{curBatchSize, 3, imgShape.height, imgShape.width}; 146 | context->setBindingDimensions(0, inputDims); 147 | 148 | context->enqueueV2(buffers, cudaStream, nullptr); 149 | cudaMemcpyAsync(outputBuffer, buffers[outputIndex], outputStreamSize * sizeof(float), cudaMemcpyDeviceToHost, cudaStream); 150 | // cudaStreamSynchronize(cudaStream); 151 | } 152 | 153 | void FeatureTensor::mat2stream(vector& imgMats, float* stream) { 154 | int imgArea = imgShape.area(); 155 | curBatchSize = imgMats.size(); 156 | if (curBatchSize > maxBatchSize) { 157 | std::cout << "[WARNING]::Batch size overflow, input will be truncated!" << std::endl; 158 | curBatchSize = maxBatchSize; 159 | } 160 | for (int batch = 0; batch < curBatchSize; ++batch) { 161 | cv::Mat tempMat = imgMats[batch]; 162 | int i = 0; 163 | for (int row = 0; row < imgShape.height; ++row) { 164 | uchar* uc_pixel = tempMat.data + row * tempMat.step; 165 | for (int col = 0; col < imgShape.width; ++col) { 166 | stream[batch * 3 * imgArea + i] = ((float)uc_pixel[0] / 255.0 - means[0]) / std[0]; 167 | stream[batch * 3 * imgArea + i + imgArea] = ((float)uc_pixel[1] / 255.0 - means[1]) / std[1]; 168 | stream[batch * 3 * imgArea + i + 2 * imgArea] = ((float)uc_pixel[2] / 255.0 - means[2]) / std[2]; 169 | uc_pixel += 3; 170 | ++i; 171 | } 172 | } 173 | } 174 | } 175 | 176 | void FeatureTensor::stream2det(float* stream, DETECTIONS& det) { 177 | int i = 0; 178 | for (DETECTION_ROW& dbox : det) { 179 | for (int j = 0; j < featureDim; ++j) 180 | dbox.feature[j] = stream[i * featureDim + j]; 181 | // dbox.feature[j] = (float)1.0; 182 | ++i; 183 | } 184 | } 185 | -------------------------------------------------------------------------------- /deepsort/src/hungarianoper.cpp: -------------------------------------------------------------------------------- 1 | #include "hungarianoper.h" 2 | 3 | Eigen::Matrix HungarianOper::Solve(const DYNAMICM &cost_matrix) { 4 | int rows = cost_matrix.rows(); 5 | int cols = cost_matrix.cols(); 6 | Matrix matrix(rows, cols); 7 | for (int row = 0; row < rows; row++) { 8 | for (int col = 0; col < cols; col++) { 9 | matrix(row, col) = cost_matrix(row, col); 10 | } 11 | } 12 | //Munkres get matrix; 13 | Munkres m; 14 | m.solve(matrix); 15 | 16 | // 17 | std::vector> pairs; 18 | for (int row = 0; row < rows; row++) { 19 | for (int col = 0; col < cols; col++) { 20 | int tmp = (int)matrix(row, col); 21 | if (tmp == 0) pairs.push_back(std::make_pair(row, col)); 22 | } 23 | } 24 | // 25 | int count = pairs.size(); 26 | Eigen::Matrix re(count, 2); 27 | for (int i = 0; i < count; i++) { 28 | re(i, 0) = pairs[i].first; 29 | re(i, 1) = pairs[i].second; 30 | } 31 | return re; 32 | }//end Solve; 33 | -------------------------------------------------------------------------------- /deepsort/src/kalmanfilter.cpp: -------------------------------------------------------------------------------- 1 | #include "kalmanfilter.h" 2 | #include 3 | #include 4 | 5 | const double KalmanFilter::chi2inv95[10] = { 6 | 0, 7 | 3.8415, 8 | 5.9915, 9 | 7.8147, 10 | 9.4877, 11 | 11.070, 12 | 12.592, 13 | 14.067, 14 | 15.507, 15 | 16.919 16 | }; 17 | KalmanFilter::KalmanFilter() { 18 | int ndim = 4; 19 | double dt = 1.; 20 | 21 | _motion_mat = Eigen::MatrixXf::Identity(8, 8); 22 | for(int i = 0; i < ndim; i++) { 23 | _motion_mat(i, ndim+i) = dt; 24 | } 25 | _update_mat = Eigen::MatrixXf::Identity(4, 8); 26 | 27 | this->_std_weight_position = 1. / 20; 28 | this->_std_weight_velocity = 1. / 160; 29 | } 30 | 31 | KAL_DATA KalmanFilter::initiate(const DETECTBOX& measurement) { 32 | DETECTBOX mean_pos = measurement; 33 | DETECTBOX mean_vel; 34 | for(int i = 0; i < 4; i++) mean_vel(i) = 0; 35 | 36 | KAL_MEAN mean; 37 | for(int i = 0; i < 8; i++){ 38 | if(i < 4) mean(i) = mean_pos(i); 39 | else mean(i) = mean_vel(i - 4); 40 | } 41 | 42 | KAL_MEAN std; 43 | std(0) = 2 * _std_weight_position * measurement[3]; 44 | std(1) = 2 * _std_weight_position * measurement[3]; 45 | std(2) = 1e-2; 46 | std(3) = 2 * _std_weight_position * measurement[3]; 47 | std(4) = 10 * _std_weight_velocity * measurement[3]; 48 | std(5) = 10 * _std_weight_velocity * measurement[3]; 49 | std(6) = 1e-5; 50 | std(7) = 10 * _std_weight_velocity * measurement[3]; 51 | 52 | KAL_MEAN tmp = std.array().square(); 53 | KAL_COVA var = tmp.asDiagonal(); 54 | return std::make_pair(mean, var); 55 | } 56 | 57 | void KalmanFilter::predict(KAL_MEAN &mean, KAL_COVA &covariance) { 58 | //revise the data; 59 | DETECTBOX std_pos; 60 | std_pos << _std_weight_position * mean(3), 61 | _std_weight_position * mean(3), 62 | 1e-2, 63 | _std_weight_position * mean(3); 64 | DETECTBOX std_vel; 65 | std_vel << _std_weight_velocity * mean(3), 66 | _std_weight_velocity * mean(3), 67 | 1e-5, 68 | _std_weight_velocity * mean(3); 69 | KAL_MEAN tmp; 70 | tmp.block<1,4>(0,0) = std_pos; 71 | tmp.block<1,4>(0,4) = std_vel; 72 | tmp = tmp.array().square(); 73 | KAL_COVA motion_cov = tmp.asDiagonal(); 74 | KAL_MEAN mean1 = this->_motion_mat * mean.transpose(); 75 | KAL_COVA covariance1 = this->_motion_mat * covariance *(_motion_mat.transpose()); 76 | covariance1 += motion_cov; 77 | 78 | mean = mean1; 79 | covariance = covariance1; 80 | } 81 | 82 | KAL_HDATA KalmanFilter::project(const KAL_MEAN &mean, const KAL_COVA &covariance) { 83 | DETECTBOX std; 84 | std << _std_weight_position * mean(3), _std_weight_position * mean(3), 85 | 1e-1, _std_weight_position * mean(3); 86 | KAL_HMEAN mean1 = _update_mat * mean.transpose(); 87 | KAL_HCOVA covariance1 = _update_mat * covariance * (_update_mat.transpose()); 88 | Eigen::Matrix diag = std.asDiagonal(); 89 | diag = diag.array().square().matrix(); 90 | covariance1 += diag; 91 | // covariance1.diagonal() << diag; 92 | return std::make_pair(mean1, covariance1); 93 | } 94 | 95 | KAL_DATA 96 | KalmanFilter::update( 97 | const KAL_MEAN &mean, 98 | const KAL_COVA &covariance, 99 | const DETECTBOX &measurement) { 100 | KAL_HDATA pa = project(mean, covariance); 101 | KAL_HMEAN projected_mean = pa.first; 102 | KAL_HCOVA projected_cov = pa.second; 103 | 104 | //chol_factor, lower = 105 | //scipy.linalg.cho_factor(projected_cov, lower=True, check_finite=False) 106 | //kalmain_gain = 107 | //scipy.linalg.cho_solve((cho_factor, lower), 108 | //np.dot(covariance, self._upadte_mat.T).T, 109 | //check_finite=False).T 110 | Eigen::Matrix B = (covariance * (_update_mat.transpose())).transpose(); 111 | Eigen::Matrix kalman_gain = (projected_cov.llt().solve(B)).transpose(); // eg.8x4 112 | Eigen::Matrix innovation = measurement - projected_mean; //eg.1x4 113 | auto tmp = innovation*(kalman_gain.transpose()); 114 | KAL_MEAN new_mean = (mean.array() + tmp.array()).matrix(); 115 | KAL_COVA new_covariance = covariance - kalman_gain*projected_cov*(kalman_gain.transpose()); 116 | return std::make_pair(new_mean, new_covariance); 117 | } 118 | 119 | Eigen::Matrix 120 | KalmanFilter::gating_distance( 121 | const KAL_MEAN &mean, 122 | const KAL_COVA &covariance, 123 | const std::vector &measurements, 124 | bool only_position) { 125 | KAL_HDATA pa = this->project(mean, covariance); 126 | if(only_position) { 127 | printf("not implement!"); 128 | exit(0); 129 | } 130 | KAL_HMEAN mean1 = pa.first; 131 | KAL_HCOVA covariance1 = pa.second; 132 | 133 | // Eigen::Matrix d(size, 4); 134 | DETECTBOXSS d(measurements.size(), 4); 135 | int pos = 0; 136 | for(DETECTBOX box:measurements) { 137 | d.row(pos++) = box - mean1; 138 | } 139 | Eigen::Matrix factor = covariance1.llt().matrixL(); 140 | Eigen::Matrix z = factor.triangularView().solve(d).transpose(); 141 | auto zz = ((z.array())*(z.array())).matrix(); 142 | auto square_maha = zz.colwise().sum(); 143 | return square_maha; 144 | } 145 | 146 | -------------------------------------------------------------------------------- /deepsort/src/linear_assignment.cpp: -------------------------------------------------------------------------------- 1 | #include "linear_assignment.h" 2 | #include "hungarianoper.h" 3 | #include 4 | 5 | linear_assignment *linear_assignment::instance = NULL; 6 | linear_assignment::linear_assignment() 7 | { 8 | } 9 | 10 | linear_assignment *linear_assignment::getInstance() 11 | { 12 | if(instance == NULL) instance = new linear_assignment(); 13 | return instance; 14 | } 15 | 16 | TRACHER_MATCHD 17 | linear_assignment::matching_cascade( 18 | tracker *distance_metric, 19 | tracker::GATED_METRIC_FUNC distance_metric_func, 20 | float max_distance, 21 | int cascade_depth, 22 | std::vector &tracks, 23 | const DETECTIONS &detections, 24 | std::vector& track_indices, 25 | std::vector detection_indices) 26 | { 27 | TRACHER_MATCHD res; 28 | // std::cout << "distance_metric" << distance_metric << std::endl; 29 | // std::cout << "max_distance" << max_distance << std::endl; 30 | // std::cout << "cascade_depth" << cascade_depth << std::endl; 31 | // std::cout << "tracks [" << std::endl; 32 | // for (auto i : tracks) 33 | // std::cout << i.hits << ", "; 34 | // std::cout << "]" << endl; 35 | // std::cout << "detections [" << std::endl; 36 | // for (auto i : detections) 37 | // std::cout << i.confidence << ", "; 38 | // std::cout << "]" << endl; 39 | // std::cout << "track_indices [" << std::endl; 40 | // for (auto i : track_indices) 41 | // std::cout << i << ", "; 42 | // std::cout << "]" << endl; 43 | // std::cout << "detection_indices [" << std::endl; 44 | // for (auto i : detection_indices) 45 | // std::cout << i << ", "; 46 | // std::cout << "]" << endl; 47 | // !!!python diff: track_indices will never be None. 48 | // if(track_indices.empty() == true) { 49 | // for(size_t i = 0; i < tracks.size(); i++) { 50 | // track_indices.push_back(i); 51 | // } 52 | // } 53 | 54 | //!!!python diff: detection_indices will always be None. 55 | for(size_t i = 0; i < detections.size(); i++) { 56 | detection_indices.push_back(int(i)); 57 | } 58 | 59 | std::vector unmatched_detections; 60 | unmatched_detections.assign(detection_indices.begin(), detection_indices.end()); 61 | res.matches.clear(); 62 | std::vector track_indices_l; 63 | 64 | std::map matches_trackid; 65 | for(int level = 0; level < cascade_depth; level++) { 66 | if(unmatched_detections.size() == 0) break; //No detections left; 67 | 68 | track_indices_l.clear(); 69 | for(int k:track_indices) { 70 | if(tracks[k].time_since_update == 1+level) 71 | track_indices_l.push_back(k); 72 | } 73 | if(track_indices_l.size() == 0) continue; //Nothing to match at this level. 74 | 75 | TRACHER_MATCHD tmp = min_cost_matching( 76 | distance_metric, distance_metric_func, 77 | max_distance, tracks, detections, track_indices_l, 78 | unmatched_detections); 79 | unmatched_detections.assign(tmp.unmatched_detections.begin(), tmp.unmatched_detections.end()); 80 | for(size_t i = 0; i < tmp.matches.size(); i++) { 81 | MATCH_DATA pa = tmp.matches[i]; 82 | res.matches.push_back(pa); 83 | matches_trackid.insert(pa); 84 | } 85 | } 86 | res.unmatched_detections.assign(unmatched_detections.begin(), unmatched_detections.end()); 87 | for(size_t i = 0; i < track_indices.size(); i++) { 88 | int tid = track_indices[i]; 89 | if(matches_trackid.find(tid) == matches_trackid.end()) 90 | res.unmatched_tracks.push_back(tid); 91 | } 92 | return res; 93 | } 94 | 95 | TRACHER_MATCHD 96 | linear_assignment::min_cost_matching(tracker *distance_metric, 97 | tracker::GATED_METRIC_FUNC distance_metric_func, 98 | float max_distance, 99 | std::vector &tracks, 100 | const DETECTIONS &detections, 101 | std::vector &track_indices, 102 | std::vector &detection_indices) 103 | { 104 | TRACHER_MATCHD res; 105 | //!!!python diff: track_indices && detection_indices will never be None. 106 | // if(track_indices.empty() == true) { 107 | // for(size_t i = 0; i < tracks.size(); i++) { 108 | // track_indices.push_back(i); 109 | // } 110 | // } 111 | // if(detection_indices.empty() == true) { 112 | // for(size_t i = 0; i < detections.size(); i++) { 113 | // detection_indices.push_back(int(i)); 114 | // } 115 | // } 116 | if((detection_indices.size() == 0) || (track_indices.size() == 0)) { 117 | res.matches.clear(); 118 | res.unmatched_tracks.assign(track_indices.begin(), track_indices.end()); 119 | res.unmatched_detections.assign(detection_indices.begin(), detection_indices.end()); 120 | return res; 121 | } 122 | DYNAMICM cost_matrix = (distance_metric->*(distance_metric_func))( 123 | tracks, detections, track_indices, detection_indices); 124 | for(int i = 0; i < cost_matrix.rows(); i++) { 125 | for(int j = 0; j < cost_matrix.cols(); j++) { 126 | float tmp = cost_matrix(i,j); 127 | if(tmp > max_distance) cost_matrix(i,j) = max_distance + 1e-5; 128 | } 129 | } 130 | Eigen::Matrix indices = HungarianOper::Solve(cost_matrix); 131 | res.matches.clear(); 132 | res.unmatched_tracks.clear(); 133 | res.unmatched_detections.clear(); 134 | for(size_t col = 0; col < detection_indices.size(); col++) { 135 | bool flag = false; 136 | for(int i = 0; i < indices.rows(); i++) 137 | if(indices(i, 1) == col) { flag = true; break;} 138 | if(flag == false)res.unmatched_detections.push_back(detection_indices[col]); 139 | } 140 | for(size_t row = 0; row < track_indices.size(); row++) { 141 | bool flag = false; 142 | for(int i = 0; i < indices.rows(); i++) 143 | if(indices(i, 0) == row) { flag = true; break; } 144 | if(flag == false) res.unmatched_tracks.push_back(track_indices[row]); 145 | } 146 | for(int i = 0; i < indices.rows(); i++) { 147 | int row = indices(i, 0); 148 | int col = indices(i, 1); 149 | 150 | int track_idx = track_indices[row]; 151 | int detection_idx = detection_indices[col]; 152 | if(cost_matrix(row, col) > max_distance) { 153 | res.unmatched_tracks.push_back(track_idx); 154 | res.unmatched_detections.push_back(detection_idx); 155 | } else res.matches.push_back(std::make_pair(track_idx, detection_idx)); 156 | } 157 | return res; 158 | } 159 | 160 | DYNAMICM 161 | linear_assignment::gate_cost_matrix( 162 | KalmanFilter *kf, 163 | DYNAMICM &cost_matrix, 164 | std::vector &tracks, 165 | const DETECTIONS &detections, 166 | const std::vector &track_indices, 167 | const std::vector &detection_indices, 168 | float gated_cost, bool only_position) 169 | { 170 | // std::cout << "input cost matric" << cost_matrix << std::endl; 171 | int gating_dim = (only_position == true?2:4); 172 | double gating_threshold = KalmanFilter::chi2inv95[gating_dim]; 173 | std::vector measurements; 174 | for(int i:detection_indices) { 175 | DETECTION_ROW t = detections[i]; 176 | measurements.push_back(t.to_xyah()); 177 | } 178 | for(size_t i = 0; i < track_indices.size(); i++) { 179 | Track& track = tracks[track_indices[i]]; 180 | Eigen::Matrix gating_distance = kf->gating_distance( 181 | track.mean, track.covariance, measurements, only_position); 182 | for (int j = 0; j < gating_distance.cols(); j++) { 183 | if (gating_distance(0, j) > gating_threshold) cost_matrix(i, j) = gated_cost; 184 | } 185 | } 186 | // std::cout << "out cost matrix" << cost_matrix << std::endl; 187 | return cost_matrix; 188 | } 189 | 190 | -------------------------------------------------------------------------------- /deepsort/src/munkres.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2007 John Weaver 3 | * Copyright (c) 2015 Miroslav Krajicek 4 | * 5 | * This program is free software; you can redistribute it and/or modify 6 | * it under the terms of the GNU General Public License as published by 7 | * the Free Software Foundation; either version 2 of the License, or 8 | * (at your option) any later version. 9 | * 10 | * This program is distributed in the hope that it will be useful, 11 | * but WITHOUT ANY WARRANTY; without even the implied warranty of 12 | * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 13 | * GNU General Public License for more details. 14 | * 15 | * You should have received a copy of the GNU General Public License 16 | * along with this program; if not, write to the Free Software 17 | * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA 18 | */ 19 | 20 | #include "munkres.h" 21 | 22 | template class Munkres; 23 | template class Munkres; 24 | template class Munkres; 25 | 26 | -------------------------------------------------------------------------------- /deepsort/src/nn_matching.cpp: -------------------------------------------------------------------------------- 1 | #include "nn_matching.h" 2 | #include 3 | 4 | using namespace Eigen; 5 | 6 | NearNeighborDisMetric::NearNeighborDisMetric( 7 | NearNeighborDisMetric::METRIC_TYPE metric, 8 | float matching_threshold, int budget) 9 | { 10 | if (metric == euclidean) { 11 | _metric = 12 | &NearNeighborDisMetric::_nneuclidean_distance; 13 | } else if (metric == cosine) { 14 | _metric = 15 | &NearNeighborDisMetric::_nncosine_distance; 16 | } 17 | 18 | this->mating_threshold = matching_threshold; 19 | this->budget = budget; 20 | this->samples.clear(); 21 | } 22 | 23 | DYNAMICM 24 | NearNeighborDisMetric::distance( 25 | const FEATURESS & features, 26 | const std::vector < int >&targets) 27 | { 28 | DYNAMICM cost_matrix = Eigen::MatrixXf::Zero(targets.size(), features.rows()); 29 | int idx = 0; 30 | for (int target:targets) { 31 | cost_matrix.row(idx) = (this->*_metric) (this->samples[target], features); 32 | idx++; 33 | } 34 | return cost_matrix; 35 | } 36 | 37 | void 38 | NearNeighborDisMetric::partial_fit( 39 | std::vector &tid_feats, 40 | std::vector < int >&active_targets) 41 | { 42 | /*python code: 43 | * let feature(target_id) append to samples; 44 | * && delete not comfirmed target_id from samples. 45 | * update samples; 46 | */ 47 | for (TRACKER_DATA & data:tid_feats) { 48 | int track_id = data.first; 49 | FEATURESS newFeatOne = data.second; 50 | 51 | if (samples.find(track_id) != samples.end()) { //append 52 | int oldSize = samples[track_id].rows(); 53 | int addSize = newFeatOne.rows(); 54 | int newSize = oldSize + addSize; 55 | 56 | if (newSize <= this->budget) { 57 | FEATURESS newSampleFeatures(newSize, 256); 58 | newSampleFeatures.block(0, 0, oldSize, 256) = samples[track_id]; 59 | newSampleFeatures.block(oldSize, 0, addSize, 256) = newFeatOne; 60 | samples[track_id] = newSampleFeatures; 61 | } else { 62 | if (oldSize < this->budget) { //original space is not enough; 63 | FEATURESS newSampleFeatures(this->budget, 256); 64 | if (addSize >= this->budget) { 65 | newSampleFeatures = newFeatOne.block(0, 0, this->budget, 256); 66 | } else { 67 | newSampleFeatures.block(0, 0, this->budget - addSize, 256) = 68 | samples[track_id].block(addSize - 1, 0, this->budget - addSize, 256).eval(); 69 | newSampleFeatures.block(this->budget - addSize, 0, addSize,256) = newFeatOne; 70 | } 71 | samples[track_id] = newSampleFeatures; 72 | } else { //original space is ok; 73 | if (addSize >= this->budget) { 74 | samples[track_id] = newFeatOne.block(0, 0, this->budget, 256); 75 | } else { 76 | samples[track_id].block(0, 0, this->budget - addSize, 256) = 77 | samples[track_id].block(addSize - 1, 0, this->budget - addSize, 256).eval(); 78 | samples[track_id].block(this->budget - addSize, 0, addSize, 256) = newFeatOne; 79 | } 80 | } 81 | } 82 | } else { //not exit, create new one; 83 | samples[track_id] = newFeatOne; 84 | } 85 | } //add features; 86 | 87 | //erase the samples which not in active_targets; 88 | for (std::map < int, FEATURESS >::iterator i = samples.begin(); i != samples.end();) { 89 | bool flag = false; 90 | for (int j:active_targets) if (j == i->first) { flag = true; break; } 91 | if (flag == false)samples.erase(i++); 92 | else i++; 93 | } 94 | } 95 | 96 | Eigen::VectorXf 97 | NearNeighborDisMetric::_nncosine_distance( 98 | const FEATURESS & x, const FEATURESS & y) 99 | { 100 | MatrixXf distances = _cosine_distance(x, y); 101 | VectorXf res = distances.colwise().minCoeff().transpose(); 102 | return res; 103 | } 104 | 105 | Eigen::VectorXf 106 | NearNeighborDisMetric::_nneuclidean_distance( 107 | const FEATURESS & x, const FEATURESS & y) 108 | { 109 | MatrixXf distances = _pdist(x, y); 110 | VectorXf res = distances.colwise().maxCoeff().transpose(); 111 | res = res.array().max(VectorXf::Zero(res.rows()).array()); 112 | return res; 113 | } 114 | 115 | Eigen::MatrixXf 116 | NearNeighborDisMetric::_pdist(const FEATURESS & x, const FEATURESS & y) 117 | { 118 | int len1 = x.rows(), len2 = y.rows(); 119 | if (len1 == 0 || len2 == 0) { 120 | return Eigen::MatrixXf::Zero(len1, len2); 121 | } 122 | MatrixXf res = -2.0 * x * y.transpose(); 123 | res = res.colwise() + x.rowwise().squaredNorm(); 124 | res = res.rowwise() + y.rowwise().squaredNorm().transpose(); 125 | res = res.array().max(MatrixXf::Zero(res.rows(), res.cols()).array()); 126 | return res; 127 | } 128 | 129 | Eigen::MatrixXf 130 | NearNeighborDisMetric::_cosine_distance( 131 | const FEATURESS & a, const FEATURESS & b, bool data_is_normalized) 132 | { 133 | FEATURESS aa = a; 134 | FEATURESS bb = b; 135 | if (!data_is_normalized) { 136 | //undo: 137 | for (int i = 0; i < a.rows(); ++i) { 138 | aa.row(i) = a.row(i) / sqrt(a.row(i).squaredNorm()); 139 | } 140 | for (int i = 0; i < b.rows(); ++i) { 141 | bb.row(i) = b.row(i) / sqrt(b.row(i).squaredNorm()); 142 | } 143 | } 144 | MatrixXf res = 1. - (aa * bb.transpose()).array(); 145 | return res; 146 | } 147 | -------------------------------------------------------------------------------- /deepsort/src/track.cpp: -------------------------------------------------------------------------------- 1 | #include "track.h" 2 | #include 3 | 4 | Track::Track(KAL_MEAN & mean, KAL_COVA & covariance, int track_id, int n_init, int max_age, const FEATURE & feature) 5 | { 6 | this->mean = mean; 7 | this->covariance = covariance; 8 | this->track_id = track_id; 9 | this->hits = 1; 10 | this->age = 1; 11 | this->time_since_update = 0; 12 | this->state = TrackState::Tentative; 13 | features = FEATURESS(1, 256); 14 | features.row(0) = feature; //features.rows() must = 0; 15 | 16 | this->_n_init = n_init; 17 | this->_max_age = max_age; 18 | } 19 | 20 | Track::Track(KAL_MEAN & mean, KAL_COVA & covariance, int track_id, int n_init, int max_age, const FEATURE & feature, int cls, float conf) 21 | { 22 | this->mean = mean; 23 | this->covariance = covariance; 24 | this->track_id = track_id; 25 | this->hits = 1; 26 | this->age = 1; 27 | this->time_since_update = 0; 28 | this->state = TrackState::Tentative; 29 | features = FEATURESS(1, 256); 30 | features.row(0) = feature; //features.rows() must = 0; 31 | 32 | this->_n_init = n_init; 33 | this->_max_age = max_age; 34 | 35 | this->cls = cls; 36 | this->conf = conf; 37 | } 38 | 39 | void Track::predit(KalmanFilter * kf) 40 | { 41 | /*Propagate the state distribution to the current time step using a 42 | Kalman filter prediction step. 43 | 44 | Parameters 45 | ---------- 46 | kf : kalman_filter.KalmanFilterd 47 | The Kalman filter. 48 | */ 49 | 50 | kf->predict(this->mean, this->covariance); 51 | 52 | 53 | this->age += 1; 54 | this->time_since_update += 1; 55 | } 56 | 57 | void Track::update(KalmanFilter * const kf, const DETECTION_ROW & detection) 58 | { 59 | KAL_DATA pa = kf->update(this->mean, this->covariance, detection.to_xyah()); 60 | this->mean = pa.first; 61 | this->covariance = pa.second; 62 | 63 | featuresAppendOne(detection.feature); 64 | // this->features.row(features.rows()) = detection.feature; 65 | this->hits += 1; 66 | this->time_since_update = 0; 67 | if (this->state == TrackState::Tentative && this->hits >= this->_n_init) { 68 | this->state = TrackState::Confirmed; 69 | } 70 | } 71 | 72 | void Track::update(KalmanFilter * const kf, const DETECTION_ROW & detection, CLSCONF pair_det) 73 | { 74 | KAL_DATA pa = kf->update(this->mean, this->covariance, detection.to_xyah()); 75 | this->mean = pa.first; 76 | this->covariance = pa.second; 77 | 78 | featuresAppendOne(detection.feature); 79 | // this->features.row(features.rows()) = detection.feature; 80 | this->hits += 1; 81 | this->time_since_update = 0; 82 | if (this->state == TrackState::Tentative && this->hits >= this->_n_init) { 83 | this->state = TrackState::Confirmed; 84 | } 85 | this->cls = pair_det.cls; 86 | this->conf = pair_det.conf; 87 | } 88 | 89 | void Track::mark_missed() 90 | { 91 | if (this->state == TrackState::Tentative) { 92 | this->state = TrackState::Deleted; 93 | } else if (this->time_since_update > this->_max_age) { 94 | this->state = TrackState::Deleted; 95 | } 96 | } 97 | 98 | bool Track::is_confirmed() 99 | { 100 | return this->state == TrackState::Confirmed; 101 | } 102 | 103 | bool Track::is_deleted() 104 | { 105 | return this->state == TrackState::Deleted; 106 | } 107 | 108 | bool Track::is_tentative() 109 | { 110 | return this->state == TrackState::Tentative; 111 | } 112 | 113 | DETECTBOX Track::to_tlwh() 114 | { 115 | DETECTBOX ret = mean.leftCols(4); 116 | ret(2) *= ret(3); 117 | ret.leftCols(2) -= (ret.rightCols(2) / 2); 118 | return ret; 119 | } 120 | 121 | void Track::featuresAppendOne(const FEATURE & f) 122 | { 123 | int size = this->features.rows(); 124 | FEATURESS newfeatures = FEATURESS(size + 1, 256); 125 | newfeatures.block(0, 0, size, 256) = this->features; 126 | newfeatures.row(size) = f; 127 | features = newfeatures; 128 | } 129 | -------------------------------------------------------------------------------- /deepsort/src/tracker.cpp: -------------------------------------------------------------------------------- 1 | #include "tracker.h" 2 | #include "nn_matching.h" 3 | #include "linear_assignment.h" 4 | using namespace std; 5 | 6 | #define MY_inner_DEBUG 7 | #ifdef MY_inner_DEBUG 8 | #include 9 | #include 10 | #endif 11 | using namespace std; 12 | tracker::tracker( /*NearNeighborDisMetric *metric, */ 13 | float max_cosine_distance, int nn_budget, 14 | float max_iou_distance, int max_age, int n_init) 15 | { 16 | this->metric = new NearNeighborDisMetric( 17 | NearNeighborDisMetric::METRIC_TYPE::cosine, 18 | max_cosine_distance, nn_budget); 19 | this->max_iou_distance = max_iou_distance; 20 | this->max_age = max_age; 21 | this->n_init = n_init; 22 | 23 | this->kf = new KalmanFilter(); 24 | this->tracks.clear(); 25 | this->_next_idx = 1; 26 | } 27 | 28 | void tracker::predict() 29 | { 30 | for (Track & track:tracks) { 31 | track.predit(kf); 32 | } 33 | } 34 | 35 | void tracker::update(const DETECTIONS & detections) 36 | { 37 | TRACHER_MATCHD res; 38 | _match(detections, res); 39 | 40 | vector < MATCH_DATA > &matches = res.matches; 41 | for (MATCH_DATA & data:matches) { 42 | int track_idx = data.first; 43 | int detection_idx = data.second; 44 | tracks[track_idx].update(this->kf, detections[detection_idx]); 45 | } 46 | vector < int >&unmatched_tracks = res.unmatched_tracks; 47 | for (int &track_idx:unmatched_tracks) { 48 | this->tracks[track_idx].mark_missed(); 49 | } 50 | vector < int >&unmatched_detections = res.unmatched_detections; 51 | for (int &detection_idx:unmatched_detections) { 52 | this->_initiate_track(detections[detection_idx]); 53 | } 54 | vector < Track >::iterator it; 55 | for (it = tracks.begin(); it != tracks.end();) { 56 | if ((*it).is_deleted()) it = tracks.erase(it); 57 | else ++it; 58 | } 59 | vector < int >active_targets; 60 | vector < TRACKER_DATA > tid_features; 61 | for (Track & track:tracks) { 62 | if (track.is_confirmed() == false) continue; 63 | active_targets.push_back(track.track_id); 64 | tid_features.push_back(std::make_pair(track. track_id, track.features)); 65 | FEATURESS t = FEATURESS(0, 256); 66 | track.features = t; 67 | } 68 | this->metric->partial_fit(tid_features, active_targets); 69 | } 70 | 71 | void tracker::update(const DETECTIONSV2 & detectionsv2) 72 | { 73 | const vector& clsConf = detectionsv2.first; 74 | const DETECTIONS& detections = detectionsv2.second; 75 | TRACHER_MATCHD res; 76 | _match(detections, res); 77 | 78 | vector < MATCH_DATA > &matches = res.matches; 79 | for (MATCH_DATA & data:matches) { 80 | int track_idx = data.first; 81 | int detection_idx = data.second; 82 | tracks[track_idx].update(this->kf, detections[detection_idx], clsConf[detection_idx]); 83 | } 84 | vector < int >&unmatched_tracks = res.unmatched_tracks; 85 | for (int &track_idx:unmatched_tracks) { 86 | this->tracks[track_idx].mark_missed(); 87 | } 88 | vector < int >&unmatched_detections = res.unmatched_detections; 89 | for (int &detection_idx:unmatched_detections) { 90 | this->_initiate_track(detections[detection_idx], clsConf[detection_idx]); 91 | } 92 | vector < Track >::iterator it; 93 | for (it = tracks.begin(); it != tracks.end();) { 94 | if ((*it).is_deleted()) it = tracks.erase(it); 95 | else ++it; 96 | } 97 | vector < int >active_targets; 98 | vector < TRACKER_DATA > tid_features; 99 | for (Track & track:tracks) { 100 | if (track.is_confirmed() == false) continue; 101 | active_targets.push_back(track.track_id); 102 | tid_features.push_back(std::make_pair(track. track_id, track.features)); 103 | FEATURESS t = FEATURESS(0, 256); 104 | track.features = t; 105 | } 106 | this->metric->partial_fit(tid_features, active_targets); 107 | } 108 | 109 | void tracker::_match(const DETECTIONS & detections, TRACHER_MATCHD & res) 110 | { 111 | vector < int >confirmed_tracks; 112 | vector < int >unconfirmed_tracks; 113 | int idx = 0; 114 | for (Track & t:tracks) { 115 | if (t.is_confirmed()) confirmed_tracks.push_back(idx); 116 | else unconfirmed_tracks.push_back(idx); 117 | idx++; 118 | } 119 | 120 | TRACHER_MATCHD matcha = linear_assignment::getInstance()-> matching_cascade( 121 | this, &tracker::gated_matric, 122 | this->metric->mating_threshold, 123 | this->max_age, 124 | this->tracks, 125 | detections, 126 | confirmed_tracks); 127 | vector < int >iou_track_candidates; 128 | iou_track_candidates.assign(unconfirmed_tracks.begin(), unconfirmed_tracks.end()); 129 | vector < int >::iterator it; 130 | for (it = matcha.unmatched_tracks.begin(); it != matcha.unmatched_tracks.end();) { 131 | int idx = *it; 132 | if (tracks[idx].time_since_update == 1) { //push into unconfirmed 133 | iou_track_candidates.push_back(idx); 134 | it = matcha.unmatched_tracks.erase(it); 135 | continue; 136 | } 137 | ++it; 138 | } 139 | TRACHER_MATCHD matchb = linear_assignment::getInstance()->min_cost_matching( 140 | this, &tracker::iou_cost, 141 | this->max_iou_distance, 142 | this->tracks, 143 | detections, 144 | iou_track_candidates, 145 | matcha.unmatched_detections); 146 | //get result: 147 | res.matches.assign(matcha.matches.begin(), matcha.matches.end()); 148 | res.matches.insert(res.matches.end(), matchb.matches.begin(), matchb.matches.end()); 149 | //unmatched_tracks; 150 | res.unmatched_tracks.assign( 151 | matcha.unmatched_tracks.begin(), 152 | matcha.unmatched_tracks.end()); 153 | res.unmatched_tracks.insert( 154 | res.unmatched_tracks.end(), 155 | matchb.unmatched_tracks.begin(), 156 | matchb.unmatched_tracks.end()); 157 | res.unmatched_detections.assign( 158 | matchb.unmatched_detections.begin(), 159 | matchb.unmatched_detections.end()); 160 | } 161 | 162 | void tracker::_initiate_track(const DETECTION_ROW & detection) 163 | { 164 | KAL_DATA data = kf->initiate(detection.to_xyah()); 165 | KAL_MEAN mean = data.first; 166 | KAL_COVA covariance = data.second; 167 | 168 | this->tracks.push_back(Track(mean, covariance, this->_next_idx, this->n_init, 169 | this->max_age, detection.feature)); 170 | _next_idx += 1; 171 | } 172 | void tracker::_initiate_track(const DETECTION_ROW& detection, CLSCONF clsConf) 173 | { 174 | KAL_DATA data = kf->initiate(detection.to_xyah()); 175 | KAL_MEAN mean = data.first; 176 | KAL_COVA covariance = data.second; 177 | 178 | this->tracks.push_back(Track(mean, covariance, this->_next_idx, this->n_init, 179 | this->max_age, detection.feature, clsConf.cls, clsConf.conf)); 180 | _next_idx += 1; 181 | } 182 | 183 | DYNAMICM tracker::gated_matric( 184 | std::vector < Track > &tracks, 185 | const DETECTIONS & dets, 186 | const std::vector < int >&track_indices, 187 | const std::vector < int >&detection_indices) 188 | { 189 | FEATURESS features(detection_indices.size(), 256); 190 | int pos = 0; 191 | for (int i:detection_indices) { 192 | features.row(pos++) = dets[i].feature; 193 | } 194 | vector < int >targets; 195 | for (int i:track_indices) { 196 | targets.push_back(tracks[i].track_id); 197 | } 198 | DYNAMICM cost_matrix = this->metric->distance(features, targets); 199 | DYNAMICM res = linear_assignment::getInstance()->gate_cost_matrix( 200 | this->kf, cost_matrix, tracks, dets, track_indices, 201 | detection_indices); 202 | return res; 203 | } 204 | 205 | DYNAMICM 206 | tracker::iou_cost( 207 | std::vector < Track > &tracks, 208 | const DETECTIONS & dets, 209 | const std::vector < int >&track_indices, 210 | const std::vector < int >&detection_indices) 211 | { 212 | //!!!python diff: track_indices && detection_indices will never be None. 213 | // if(track_indices.empty() == true) { 214 | // for(size_t i = 0; i < tracks.size(); i++) { 215 | // track_indices.push_back(i); 216 | // } 217 | // } 218 | // if(detection_indices.empty() == true) { 219 | // for(size_t i = 0; i < dets.size(); i++) { 220 | // detection_indices.push_back(i); 221 | // } 222 | // } 223 | int rows = track_indices.size(); 224 | int cols = detection_indices.size(); 225 | DYNAMICM cost_matrix = Eigen::MatrixXf::Zero(rows, cols); 226 | for (int i = 0; i < rows; i++) { 227 | int track_idx = track_indices[i]; 228 | if (tracks[track_idx].time_since_update > 1) { 229 | cost_matrix.row(i) = Eigen::RowVectorXf::Constant(cols, INFTY_COST); 230 | continue; 231 | } 232 | DETECTBOX bbox = tracks[track_idx].to_tlwh(); 233 | int csize = detection_indices.size(); 234 | DETECTBOXSS candidates(csize, 4); 235 | for (int k = 0; k < csize; k++) candidates.row(k) = dets[detection_indices[k]].tlwh; 236 | Eigen::RowVectorXf rowV = (1. - iou(bbox, candidates).array()).matrix().transpose(); 237 | cost_matrix.row(i) = rowV; 238 | } 239 | return cost_matrix; 240 | } 241 | 242 | Eigen::VectorXf 243 | tracker::iou(DETECTBOX & bbox, DETECTBOXSS & candidates) 244 | { 245 | float bbox_tl_1 = bbox[0]; 246 | float bbox_tl_2 = bbox[1]; 247 | float bbox_br_1 = bbox[0] + bbox[2]; 248 | float bbox_br_2 = bbox[1] + bbox[3]; 249 | float area_bbox = bbox[2] * bbox[3]; 250 | 251 | Eigen::Matrix < float, -1, 2 > candidates_tl; 252 | Eigen::Matrix < float, -1, 2 > candidates_br; 253 | candidates_tl = candidates.leftCols(2); 254 | candidates_br = candidates.rightCols(2) + candidates_tl; 255 | 256 | int size = int (candidates.rows()); 257 | // Eigen::VectorXf area_intersection(size); 258 | // Eigen::VectorXf area_candidates(size); 259 | Eigen::VectorXf res(size); 260 | for (int i = 0; i < size; i++) { 261 | float tl_1 = std::max(bbox_tl_1, candidates_tl(i, 0)); 262 | float tl_2 = std::max(bbox_tl_2, candidates_tl(i, 1)); 263 | float br_1 = std::min(bbox_br_1, candidates_br(i, 0)); 264 | float br_2 = std::min(bbox_br_2, candidates_br(i, 1)); 265 | 266 | float w = br_1 - tl_1; w = (w < 0 ? 0 : w); 267 | float h = br_2 - tl_2; h = (h < 0 ? 0 : h); 268 | float area_intersection = w * h; 269 | float area_candidates = candidates(i, 2) * candidates(i, 3); 270 | res[i] = area_intersection / (area_bbox + area_candidates - area_intersection); 271 | } 272 | return res; 273 | } 274 | -------------------------------------------------------------------------------- /include/manager.hpp: -------------------------------------------------------------------------------- 1 | #ifndef _MANAGER_H 2 | #define _MANAGER_H 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include "deepsort.h" 9 | #include "logging.h" 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include "time.h" 15 | 16 | #include 17 | #include 18 | #include "yolov5_lib.h" 19 | #include "deepsort.h" 20 | 21 | using std::vector; 22 | using namespace cv; 23 | //static Logger gLogger; 24 | 25 | class Trtyolosort{ 26 | public: 27 | // init 28 | Trtyolosort(char *yolo_engine_path,char *sort_engine_path); 29 | // detect and show 30 | int TrtDetect(cv::Mat &frame,float &conf_thresh,std::vector &det); 31 | void showDetection(cv::Mat& img, std::vector& boxes); 32 | 33 | private: 34 | char* yolo_engine_path_ = NULL; 35 | char* sort_engine_path_ = NULL; 36 | void *trt_engine = NULL; 37 | // deepsort parms 38 | DeepSort* DS; 39 | std::vector t; 40 | //std::vector det; 41 | }; 42 | #endif // _MANAGER_H 43 | 44 | -------------------------------------------------------------------------------- /src/main.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include "manager.hpp" 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | using namespace cv; 12 | 13 | 14 | 15 | 16 | int main(){ 17 | // calculate every person's (id,(up_num,down_num,average_x,average_y)) 18 | map> personstate; 19 | map classidmap; 20 | bool is_first = true; 21 | char* yolo_engine = ""; 22 | char* sort_engine = ""; 23 | float conf_thre = 0.4; 24 | Trtyolosort yosort(yolo_engine,sort_engine); 25 | VideoCapture capture; 26 | cv::Mat frame; 27 | frame = capture.open(""); 28 | if (!capture.isOpened()){ 29 | std::cout<<"can not open"< det; 34 | auto start_draw_time = std::chrono::system_clock::now(); 35 | 36 | clock_t start_draw,end_draw; 37 | start_draw = clock(); 38 | int i = 0; 39 | while(capture.read(frame)){ 40 | if (i%3==0){ 41 | //std::cout<<"origin img size:"<(end - start).count(); 46 | std::cout << "delay_infer:" << delay_infer << "ms" << std::endl; 47 | } 48 | i++; 49 | } 50 | capture.release(); 51 | return 0; 52 | 53 | } 54 | -------------------------------------------------------------------------------- /src/manager.cpp: -------------------------------------------------------------------------------- 1 | #include "manager.hpp" 2 | using std::vector; 3 | using namespace cv; 4 | static Logger gLogger; 5 | 6 | Trtyolosort::Trtyolosort(char *yolo_engine_path,char *sort_engine_path){ 7 | sort_engine_path_ = sort_engine_path; 8 | yolo_engine_path_ = yolo_engine_path; 9 | trt_engine = yolov5_trt_create(yolo_engine_path_); 10 | printf("create yolov5-trt , instance = %p\n", trt_engine); 11 | DS = new DeepSort(sort_engine_path_, 128, 256, 0, &gLogger); 12 | 13 | } 14 | void Trtyolosort::showDetection(cv::Mat& img, std::vector& boxes) { 15 | cv::Mat temp = img.clone(); 16 | for (auto box : boxes) { 17 | cv::Point lt(box.x1, box.y1); 18 | cv::Point br(box.x2, box.y2); 19 | cv::rectangle(temp, lt, br, cv::Scalar(255, 0, 0), 1); 20 | //std::string lbl = cv::format("ID:%d_C:%d_CONF:%.2f", (int)box.trackID, (int)box.classID, box.confidence); 21 | //std::string lbl = cv::format("ID:%d_C:%d", (int)box.trackID, (int)box.classID); 22 | std::string lbl = cv::format("ID:%d_x:%f_y:%f",(int)box.trackID,(box.x1+box.x2)/2,(box.y1+box.y2)/2); 23 | cv::putText(temp, lbl, lt, cv::FONT_HERSHEY_COMPLEX, 0.5, cv::Scalar(0,255,0)); 24 | } 25 | cv::imshow("img", temp); 26 | cv::waitKey(1); 27 | } 28 | int Trtyolosort::TrtDetect(cv::Mat &frame,float &conf_thresh,std::vector &det){ 29 | // yolo detect 30 | auto ret = yolov5_trt_detect(trt_engine, frame, conf_thresh,det); 31 | DS->sort(frame,det); 32 | //showDetection(frame,det); 33 | return 1 ; 34 | 35 | } 36 | -------------------------------------------------------------------------------- /yolo/include/calibrator.h: -------------------------------------------------------------------------------- 1 | #ifndef ENTROPY_CALIBRATOR_H 2 | #define ENTROPY_CALIBRATOR_H 3 | 4 | #include 5 | #include 6 | #include 7 | #include "macros.h" 8 | 9 | //! \class Int8EntropyCalibrator2 10 | //! 11 | //! \brief Implements Entropy calibrator 2. 12 | //! CalibrationAlgoType is kENTROPY_CALIBRATION_2. 13 | //! 14 | class Int8EntropyCalibrator2 : public nvinfer1::IInt8EntropyCalibrator2 15 | { 16 | public: 17 | Int8EntropyCalibrator2(int batchsize, int input_w, int input_h, const char* img_dir, const char* calib_table_name, const char* input_blob_name, bool read_cache = true); 18 | 19 | virtual ~Int8EntropyCalibrator2(); 20 | int getBatchSize() const TRT_NOEXCEPT override; 21 | bool getBatch(void* bindings[], const char* names[], int nbBindings) TRT_NOEXCEPT override; 22 | const void* readCalibrationCache(size_t& length) TRT_NOEXCEPT override; 23 | void writeCalibrationCache(const void* cache, size_t length) TRT_NOEXCEPT override; 24 | 25 | private: 26 | int batchsize_; 27 | int input_w_; 28 | int input_h_; 29 | int img_idx_; 30 | std::string img_dir_; 31 | std::vector img_files_; 32 | size_t input_count_; 33 | std::string calib_table_name_; 34 | const char* input_blob_name_; 35 | bool read_cache_; 36 | void* device_input_; 37 | std::vector calib_cache_; 38 | }; 39 | 40 | #endif // ENTROPY_CALIBRATOR_H 41 | -------------------------------------------------------------------------------- /yolo/include/common.hpp: -------------------------------------------------------------------------------- 1 | #ifndef YOLOV5_COMMON_H_ 2 | #define YOLOV5_COMMON_H_ 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include "NvInfer.h" 10 | #include "yololayer.h" 11 | 12 | using namespace nvinfer1; 13 | 14 | cv::Rect get_rect(cv::Mat& img, float bbox[4]) { 15 | int l, r, t, b; 16 | float r_w = Yolo::INPUT_W / (img.cols * 1.0); 17 | float r_h = Yolo::INPUT_H / (img.rows * 1.0); 18 | if (r_h > r_w) { 19 | l = bbox[0] - bbox[2] / 2.f; 20 | r = bbox[0] + bbox[2] / 2.f; 21 | t = bbox[1] - bbox[3] / 2.f - (Yolo::INPUT_H - r_w * img.rows) / 2; 22 | b = bbox[1] + bbox[3] / 2.f - (Yolo::INPUT_H - r_w * img.rows) / 2; 23 | l = l / r_w; 24 | r = r / r_w; 25 | t = t / r_w; 26 | b = b / r_w; 27 | } else { 28 | l = bbox[0] - bbox[2] / 2.f - (Yolo::INPUT_W - r_h * img.cols) / 2; 29 | r = bbox[0] + bbox[2] / 2.f - (Yolo::INPUT_W - r_h * img.cols) / 2; 30 | t = bbox[1] - bbox[3] / 2.f; 31 | b = bbox[1] + bbox[3] / 2.f; 32 | l = l / r_h; 33 | r = r / r_h; 34 | t = t / r_h; 35 | b = b / r_h; 36 | } 37 | return cv::Rect(l, t, r - l, b - t); 38 | } 39 | 40 | float iou(float lbox[4], float rbox[4]) { 41 | float interBox[] = { 42 | (std::max)(lbox[0] - lbox[2] / 2.f , rbox[0] - rbox[2] / 2.f), //left 43 | (std::min)(lbox[0] + lbox[2] / 2.f , rbox[0] + rbox[2] / 2.f), //right 44 | (std::max)(lbox[1] - lbox[3] / 2.f , rbox[1] - rbox[3] / 2.f), //top 45 | (std::min)(lbox[1] + lbox[3] / 2.f , rbox[1] + rbox[3] / 2.f), //bottom 46 | }; 47 | 48 | if (interBox[2] > interBox[3] || interBox[0] > interBox[1]) 49 | return 0.0f; 50 | 51 | float interBoxS = (interBox[1] - interBox[0])*(interBox[3] - interBox[2]); 52 | return interBoxS / (lbox[2] * lbox[3] + rbox[2] * rbox[3] - interBoxS); 53 | } 54 | 55 | bool cmp(const Yolo::Detection& a, const Yolo::Detection& b) { 56 | return a.conf > b.conf; 57 | } 58 | 59 | void nms(std::vector& res, float *output, float conf_thresh, float nms_thresh = 0.5) { 60 | int det_size = sizeof(Yolo::Detection) / sizeof(float); 61 | std::map> m; 62 | for (int i = 0; i < output[0] && i < Yolo::MAX_OUTPUT_BBOX_COUNT; i++) { 63 | if (output[1 + det_size * i + 4] <= conf_thresh) continue; 64 | Yolo::Detection det; 65 | memcpy(&det, &output[1 + det_size * i], det_size * sizeof(float)); 66 | if (m.count(det.class_id) == 0) m.emplace(det.class_id, std::vector()); 67 | m[det.class_id].push_back(det); 68 | } 69 | for (auto it = m.begin(); it != m.end(); it++) { 70 | //std::cout << it->second[0].class_id << " --- " << std::endl; 71 | auto& dets = it->second; 72 | std::sort(dets.begin(), dets.end(), cmp); 73 | for (size_t m = 0; m < dets.size(); ++m) { 74 | auto& item = dets[m]; 75 | res.push_back(item); 76 | for (size_t n = m + 1; n < dets.size(); ++n) { 77 | if (iou(item.bbox, dets[n].bbox) > nms_thresh) { 78 | dets.erase(dets.begin() + n); 79 | --n; 80 | } 81 | } 82 | } 83 | } 84 | } 85 | 86 | // TensorRT weight files have a simple space delimited format: 87 | // [type] [size] 88 | std::map loadWeights(const std::string file) { 89 | std::cout << "Loading weights: " << file << std::endl; 90 | std::map weightMap; 91 | 92 | // Open weights file 93 | std::ifstream input(file); 94 | assert(input.is_open() && "Unable to load weight file. please check if the .wts file path is right!!!!!!"); 95 | 96 | // Read number of weight blobs 97 | int32_t count; 98 | input >> count; 99 | assert(count > 0 && "Invalid weight map file."); 100 | 101 | while (count--) 102 | { 103 | Weights wt{ DataType::kFLOAT, nullptr, 0 }; 104 | uint32_t size; 105 | 106 | // Read name and type of blob 107 | std::string name; 108 | input >> name >> std::dec >> size; 109 | wt.type = DataType::kFLOAT; 110 | 111 | // Load blob 112 | uint32_t* val = reinterpret_cast(malloc(sizeof(val) * size)); 113 | for (uint32_t x = 0, y = size; x < y; ++x) 114 | { 115 | input >> std::hex >> val[x]; 116 | } 117 | wt.values = val; 118 | 119 | wt.count = size; 120 | weightMap[name] = wt; 121 | } 122 | 123 | return weightMap; 124 | } 125 | 126 | IScaleLayer* addBatchNorm2d(INetworkDefinition *network, std::map& weightMap, ITensor& input, std::string lname, float eps) { 127 | float *gamma = (float*)weightMap[lname + ".weight"].values; 128 | float *beta = (float*)weightMap[lname + ".bias"].values; 129 | float *mean = (float*)weightMap[lname + ".running_mean"].values; 130 | float *var = (float*)weightMap[lname + ".running_var"].values; 131 | int len = weightMap[lname + ".running_var"].count; 132 | 133 | float *scval = reinterpret_cast(malloc(sizeof(float) * len)); 134 | for (int i = 0; i < len; i++) { 135 | scval[i] = gamma[i] / sqrt(var[i] + eps); 136 | } 137 | Weights scale{ DataType::kFLOAT, scval, len }; 138 | 139 | float *shval = reinterpret_cast(malloc(sizeof(float) * len)); 140 | for (int i = 0; i < len; i++) { 141 | shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps); 142 | } 143 | Weights shift{ DataType::kFLOAT, shval, len }; 144 | 145 | float *pval = reinterpret_cast(malloc(sizeof(float) * len)); 146 | for (int i = 0; i < len; i++) { 147 | pval[i] = 1.0; 148 | } 149 | Weights power{ DataType::kFLOAT, pval, len }; 150 | 151 | weightMap[lname + ".scale"] = scale; 152 | weightMap[lname + ".shift"] = shift; 153 | weightMap[lname + ".power"] = power; 154 | IScaleLayer* scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power); 155 | assert(scale_1); 156 | return scale_1; 157 | } 158 | 159 | ILayer* convBlock(INetworkDefinition *network, std::map& weightMap, ITensor& input, int outch, int ksize, int s, int g, std::string lname) { 160 | Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; 161 | int p = ksize / 2; 162 | IConvolutionLayer* conv1 = network->addConvolutionNd(input, outch, DimsHW{ ksize, ksize }, weightMap[lname + ".conv.weight"], emptywts); 163 | assert(conv1); 164 | conv1->setStrideNd(DimsHW{ s, s }); 165 | conv1->setPaddingNd(DimsHW{ p, p }); 166 | conv1->setNbGroups(g); 167 | IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + ".bn", 1e-3); 168 | 169 | // silu = x * sigmoid 170 | auto sig = network->addActivation(*bn1->getOutput(0), ActivationType::kSIGMOID); 171 | assert(sig); 172 | auto ew = network->addElementWise(*bn1->getOutput(0), *sig->getOutput(0), ElementWiseOperation::kPROD); 173 | assert(ew); 174 | return ew; 175 | } 176 | 177 | ILayer* focus(INetworkDefinition *network, std::map& weightMap, ITensor& input, int inch, int outch, int ksize, std::string lname) { 178 | ISliceLayer *s1 = network->addSlice(input, Dims3{ 0, 0, 0 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 }); 179 | ISliceLayer *s2 = network->addSlice(input, Dims3{ 0, 1, 0 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 }); 180 | ISliceLayer *s3 = network->addSlice(input, Dims3{ 0, 0, 1 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 }); 181 | ISliceLayer *s4 = network->addSlice(input, Dims3{ 0, 1, 1 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 }); 182 | ITensor* inputTensors[] = { s1->getOutput(0), s2->getOutput(0), s3->getOutput(0), s4->getOutput(0) }; 183 | auto cat = network->addConcatenation(inputTensors, 4); 184 | auto conv = convBlock(network, weightMap, *cat->getOutput(0), outch, ksize, 1, 1, lname + ".conv"); 185 | return conv; 186 | } 187 | 188 | ILayer* bottleneck(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, bool shortcut, int g, float e, std::string lname) { 189 | auto cv1 = convBlock(network, weightMap, input, (int)((float)c2 * e), 1, 1, 1, lname + ".cv1"); 190 | auto cv2 = convBlock(network, weightMap, *cv1->getOutput(0), c2, 3, 1, g, lname + ".cv2"); 191 | if (shortcut && c1 == c2) { 192 | auto ew = network->addElementWise(input, *cv2->getOutput(0), ElementWiseOperation::kSUM); 193 | return ew; 194 | } 195 | return cv2; 196 | } 197 | 198 | ILayer* bottleneckCSP(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, int n, bool shortcut, int g, float e, std::string lname) { 199 | Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; 200 | int c_ = (int)((float)c2 * e); 201 | auto cv1 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv1"); 202 | auto cv2 = network->addConvolutionNd(input, c_, DimsHW{ 1, 1 }, weightMap[lname + ".cv2.weight"], emptywts); 203 | ITensor *y1 = cv1->getOutput(0); 204 | for (int i = 0; i < n; i++) { 205 | auto b = bottleneck(network, weightMap, *y1, c_, c_, shortcut, g, 1.0, lname + ".m." + std::to_string(i)); 206 | y1 = b->getOutput(0); 207 | } 208 | auto cv3 = network->addConvolutionNd(*y1, c_, DimsHW{ 1, 1 }, weightMap[lname + ".cv3.weight"], emptywts); 209 | 210 | ITensor* inputTensors[] = { cv3->getOutput(0), cv2->getOutput(0) }; 211 | auto cat = network->addConcatenation(inputTensors, 2); 212 | 213 | IScaleLayer* bn = addBatchNorm2d(network, weightMap, *cat->getOutput(0), lname + ".bn", 1e-4); 214 | auto lr = network->addActivation(*bn->getOutput(0), ActivationType::kLEAKY_RELU); 215 | lr->setAlpha(0.1); 216 | 217 | auto cv4 = convBlock(network, weightMap, *lr->getOutput(0), c2, 1, 1, 1, lname + ".cv4"); 218 | return cv4; 219 | } 220 | 221 | ILayer* C3(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, int n, bool shortcut, int g, float e, std::string lname) { 222 | int c_ = (int)((float)c2 * e); 223 | auto cv1 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv1"); 224 | auto cv2 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv2"); 225 | ITensor *y1 = cv1->getOutput(0); 226 | for (int i = 0; i < n; i++) { 227 | auto b = bottleneck(network, weightMap, *y1, c_, c_, shortcut, g, 1.0, lname + ".m." + std::to_string(i)); 228 | y1 = b->getOutput(0); 229 | } 230 | 231 | ITensor* inputTensors[] = { y1, cv2->getOutput(0) }; 232 | auto cat = network->addConcatenation(inputTensors, 2); 233 | 234 | auto cv3 = convBlock(network, weightMap, *cat->getOutput(0), c2, 1, 1, 1, lname + ".cv3"); 235 | return cv3; 236 | } 237 | 238 | ILayer* SPP(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, int k1, int k2, int k3, std::string lname) { 239 | int c_ = c1 / 2; 240 | auto cv1 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv1"); 241 | 242 | auto pool1 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX, DimsHW{ k1, k1 }); 243 | pool1->setPaddingNd(DimsHW{ k1 / 2, k1 / 2 }); 244 | pool1->setStrideNd(DimsHW{ 1, 1 }); 245 | auto pool2 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX, DimsHW{ k2, k2 }); 246 | pool2->setPaddingNd(DimsHW{ k2 / 2, k2 / 2 }); 247 | pool2->setStrideNd(DimsHW{ 1, 1 }); 248 | auto pool3 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX, DimsHW{ k3, k3 }); 249 | pool3->setPaddingNd(DimsHW{ k3 / 2, k3 / 2 }); 250 | pool3->setStrideNd(DimsHW{ 1, 1 }); 251 | 252 | ITensor* inputTensors[] = { cv1->getOutput(0), pool1->getOutput(0), pool2->getOutput(0), pool3->getOutput(0) }; 253 | auto cat = network->addConcatenation(inputTensors, 4); 254 | 255 | auto cv2 = convBlock(network, weightMap, *cat->getOutput(0), c2, 1, 1, 1, lname + ".cv2"); 256 | return cv2; 257 | } 258 | 259 | std::vector> getAnchors(std::map& weightMap, std::string lname) { 260 | std::vector> anchors; 261 | Weights wts = weightMap[lname + ".anchor_grid"]; 262 | int anchor_len = Yolo::CHECK_COUNT * 2; 263 | for (int i = 0; i < wts.count / anchor_len; i++) { 264 | auto *p = (const float*)wts.values + i * anchor_len; 265 | std::vector anchor(p, p + anchor_len); 266 | anchors.push_back(anchor); 267 | } 268 | return anchors; 269 | } 270 | 271 | IPluginV2Layer* addYoLoLayer(INetworkDefinition *network, std::map& weightMap, std::string lname, std::vector dets) { 272 | auto creator = getPluginRegistry()->getPluginCreator("YoloLayer_TRT", "1"); 273 | auto anchors = getAnchors(weightMap, lname); 274 | PluginField plugin_fields[2]; 275 | int netinfo[4] = {Yolo::CLASS_NUM, Yolo::INPUT_W, Yolo::INPUT_H, Yolo::MAX_OUTPUT_BBOX_COUNT}; 276 | plugin_fields[0].data = netinfo; 277 | plugin_fields[0].length = 4; 278 | plugin_fields[0].name = "netinfo"; 279 | plugin_fields[0].type = PluginFieldType::kFLOAT32; 280 | int scale = 8; 281 | std::vector kernels; 282 | for (size_t i = 0; i < anchors.size(); i++) { 283 | Yolo::YoloKernel kernel; 284 | kernel.width = Yolo::INPUT_W / scale; 285 | kernel.height = Yolo::INPUT_H / scale; 286 | memcpy(kernel.anchors, &anchors[i][0], anchors[i].size() * sizeof(float)); 287 | kernels.push_back(kernel); 288 | scale *= 2; 289 | } 290 | plugin_fields[1].data = &kernels[0]; 291 | plugin_fields[1].length = kernels.size(); 292 | plugin_fields[1].name = "kernels"; 293 | plugin_fields[1].type = PluginFieldType::kFLOAT32; 294 | PluginFieldCollection plugin_data; 295 | plugin_data.nbFields = 2; 296 | plugin_data.fields = plugin_fields; 297 | IPluginV2 *plugin_obj = creator->createPlugin("yololayer", &plugin_data); 298 | std::vector input_tensors; 299 | for (auto det: dets) { 300 | input_tensors.push_back(det->getOutput(0)); 301 | } 302 | auto yolo = network->addPluginV2(&input_tensors[0], input_tensors.size(), *plugin_obj); 303 | return yolo; 304 | } 305 | #endif 306 | 307 | -------------------------------------------------------------------------------- /yolo/include/cuda_utils.h: -------------------------------------------------------------------------------- 1 | #ifndef TRTX_CUDA_UTILS_H_ 2 | #define TRTX_CUDA_UTILS_H_ 3 | 4 | #include 5 | 6 | #ifndef CUDA_CHECK 7 | #define CUDA_CHECK(callstr)\ 8 | {\ 9 | cudaError_t error_code = callstr;\ 10 | if (error_code != cudaSuccess) {\ 11 | std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__;\ 12 | assert(0);\ 13 | }\ 14 | } 15 | #endif // CUDA_CHECK 16 | 17 | #endif // TRTX_CUDA_UTILS_H_ 18 | 19 | -------------------------------------------------------------------------------- /yolo/include/logging.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. 3 | * 4 | * Licensed under the Apache License, Version 2.0 (the "License"); 5 | * you may not use this file except in compliance with the License. 6 | * You may obtain a copy of the License at 7 | * 8 | * http://www.apache.org/licenses/LICENSE-2.0 9 | * 10 | * Unless required by applicable law or agreed to in writing, software 11 | * distributed under the License is distributed on an "AS IS" BASIS, 12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | * See the License for the specific language governing permissions and 14 | * limitations under the License. 15 | */ 16 | 17 | #ifndef TENSORRT_LOGGING_H 18 | #define TENSORRT_LOGGING_H 19 | 20 | #include "NvInferRuntimeCommon.h" 21 | #include 22 | #include 23 | #include 24 | #include 25 | #include 26 | #include 27 | #include 28 | #include "macros.h" 29 | 30 | using Severity = nvinfer1::ILogger::Severity; 31 | 32 | class LogStreamConsumerBuffer : public std::stringbuf 33 | { 34 | public: 35 | LogStreamConsumerBuffer(std::ostream& stream, const std::string& prefix, bool shouldLog) 36 | : mOutput(stream) 37 | , mPrefix(prefix) 38 | , mShouldLog(shouldLog) 39 | { 40 | } 41 | 42 | LogStreamConsumerBuffer(LogStreamConsumerBuffer&& other) 43 | : mOutput(other.mOutput) 44 | { 45 | } 46 | 47 | ~LogStreamConsumerBuffer() 48 | { 49 | // std::streambuf::pbase() gives a pointer to the beginning of the buffered part of the output sequence 50 | // std::streambuf::pptr() gives a pointer to the current position of the output sequence 51 | // if the pointer to the beginning is not equal to the pointer to the current position, 52 | // call putOutput() to log the output to the stream 53 | if (pbase() != pptr()) 54 | { 55 | putOutput(); 56 | } 57 | } 58 | 59 | // synchronizes the stream buffer and returns 0 on success 60 | // synchronizing the stream buffer consists of inserting the buffer contents into the stream, 61 | // resetting the buffer and flushing the stream 62 | virtual int sync() 63 | { 64 | putOutput(); 65 | return 0; 66 | } 67 | 68 | void putOutput() 69 | { 70 | if (mShouldLog) 71 | { 72 | // prepend timestamp 73 | std::time_t timestamp = std::time(nullptr); 74 | tm* tm_local = std::localtime(×tamp); 75 | std::cout << "["; 76 | std::cout << std::setw(2) << std::setfill('0') << 1 + tm_local->tm_mon << "/"; 77 | std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_mday << "/"; 78 | std::cout << std::setw(4) << std::setfill('0') << 1900 + tm_local->tm_year << "-"; 79 | std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_hour << ":"; 80 | std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_min << ":"; 81 | std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_sec << "] "; 82 | // std::stringbuf::str() gets the string contents of the buffer 83 | // insert the buffer contents pre-appended by the appropriate prefix into the stream 84 | mOutput << mPrefix << str(); 85 | // set the buffer to empty 86 | str(""); 87 | // flush the stream 88 | mOutput.flush(); 89 | } 90 | } 91 | 92 | void setShouldLog(bool shouldLog) 93 | { 94 | mShouldLog = shouldLog; 95 | } 96 | 97 | private: 98 | std::ostream& mOutput; 99 | std::string mPrefix; 100 | bool mShouldLog; 101 | }; 102 | 103 | //! 104 | //! \class LogStreamConsumerBase 105 | //! \brief Convenience object used to initialize LogStreamConsumerBuffer before std::ostream in LogStreamConsumer 106 | //! 107 | class LogStreamConsumerBase 108 | { 109 | public: 110 | LogStreamConsumerBase(std::ostream& stream, const std::string& prefix, bool shouldLog) 111 | : mBuffer(stream, prefix, shouldLog) 112 | { 113 | } 114 | 115 | protected: 116 | LogStreamConsumerBuffer mBuffer; 117 | }; 118 | 119 | //! 120 | //! \class LogStreamConsumer 121 | //! \brief Convenience object used to facilitate use of C++ stream syntax when logging messages. 122 | //! Order of base classes is LogStreamConsumerBase and then std::ostream. 123 | //! This is because the LogStreamConsumerBase class is used to initialize the LogStreamConsumerBuffer member field 124 | //! in LogStreamConsumer and then the address of the buffer is passed to std::ostream. 125 | //! This is necessary to prevent the address of an uninitialized buffer from being passed to std::ostream. 126 | //! Please do not change the order of the parent classes. 127 | //! 128 | class LogStreamConsumer : protected LogStreamConsumerBase, public std::ostream 129 | { 130 | public: 131 | //! \brief Creates a LogStreamConsumer which logs messages with level severity. 132 | //! Reportable severity determines if the messages are severe enough to be logged. 133 | LogStreamConsumer(Severity reportableSeverity, Severity severity) 134 | : LogStreamConsumerBase(severityOstream(severity), severityPrefix(severity), severity <= reportableSeverity) 135 | , std::ostream(&mBuffer) // links the stream buffer with the stream 136 | , mShouldLog(severity <= reportableSeverity) 137 | , mSeverity(severity) 138 | { 139 | } 140 | 141 | LogStreamConsumer(LogStreamConsumer&& other) 142 | : LogStreamConsumerBase(severityOstream(other.mSeverity), severityPrefix(other.mSeverity), other.mShouldLog) 143 | , std::ostream(&mBuffer) // links the stream buffer with the stream 144 | , mShouldLog(other.mShouldLog) 145 | , mSeverity(other.mSeverity) 146 | { 147 | } 148 | 149 | void setReportableSeverity(Severity reportableSeverity) 150 | { 151 | mShouldLog = mSeverity <= reportableSeverity; 152 | mBuffer.setShouldLog(mShouldLog); 153 | } 154 | 155 | private: 156 | static std::ostream& severityOstream(Severity severity) 157 | { 158 | return severity >= Severity::kINFO ? std::cout : std::cerr; 159 | } 160 | 161 | static std::string severityPrefix(Severity severity) 162 | { 163 | switch (severity) 164 | { 165 | case Severity::kINTERNAL_ERROR: return "[F] "; 166 | case Severity::kERROR: return "[E] "; 167 | case Severity::kWARNING: return "[W] "; 168 | case Severity::kINFO: return "[I] "; 169 | case Severity::kVERBOSE: return "[V] "; 170 | default: assert(0); return ""; 171 | } 172 | } 173 | 174 | bool mShouldLog; 175 | Severity mSeverity; 176 | }; 177 | 178 | //! \class Logger 179 | //! 180 | //! \brief Class which manages logging of TensorRT tools and samples 181 | //! 182 | //! \details This class provides a common interface for TensorRT tools and samples to log information to the console, 183 | //! and supports logging two types of messages: 184 | //! 185 | //! - Debugging messages with an associated severity (info, warning, error, or internal error/fatal) 186 | //! - Test pass/fail messages 187 | //! 188 | //! The advantage of having all samples use this class for logging as opposed to emitting directly to stdout/stderr is 189 | //! that the logic for controlling the verbosity and formatting of sample output is centralized in one location. 190 | //! 191 | //! In the future, this class could be extended to support dumping test results to a file in some standard format 192 | //! (for example, JUnit XML), and providing additional metadata (e.g. timing the duration of a test run). 193 | //! 194 | //! TODO: For backwards compatibility with existing samples, this class inherits directly from the nvinfer1::ILogger 195 | //! interface, which is problematic since there isn't a clean separation between messages coming from the TensorRT 196 | //! library and messages coming from the sample. 197 | //! 198 | //! In the future (once all samples are updated to use Logger::getTRTLogger() to access the ILogger) we can refactor the 199 | //! class to eliminate the inheritance and instead make the nvinfer1::ILogger implementation a member of the Logger 200 | //! object. 201 | 202 | class Logger : public nvinfer1::ILogger 203 | { 204 | public: 205 | Logger(Severity severity = Severity::kWARNING) 206 | : mReportableSeverity(severity) 207 | { 208 | } 209 | 210 | //! 211 | //! \enum TestResult 212 | //! \brief Represents the state of a given test 213 | //! 214 | enum class TestResult 215 | { 216 | kRUNNING, //!< The test is running 217 | kPASSED, //!< The test passed 218 | kFAILED, //!< The test failed 219 | kWAIVED //!< The test was waived 220 | }; 221 | 222 | //! 223 | //! \brief Forward-compatible method for retrieving the nvinfer::ILogger associated with this Logger 224 | //! \return The nvinfer1::ILogger associated with this Logger 225 | //! 226 | //! TODO Once all samples are updated to use this method to register the logger with TensorRT, 227 | //! we can eliminate the inheritance of Logger from ILogger 228 | //! 229 | nvinfer1::ILogger& getTRTLogger() 230 | { 231 | return *this; 232 | } 233 | 234 | //! 235 | //! \brief Implementation of the nvinfer1::ILogger::log() virtual method 236 | //! 237 | //! Note samples should not be calling this function directly; it will eventually go away once we eliminate the 238 | //! inheritance from nvinfer1::ILogger 239 | //! 240 | void log(Severity severity, const char* msg) TRT_NOEXCEPT override 241 | { 242 | LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl; 243 | } 244 | 245 | //! 246 | //! \brief Method for controlling the verbosity of logging output 247 | //! 248 | //! \param severity The logger will only emit messages that have severity of this level or higher. 249 | //! 250 | void setReportableSeverity(Severity severity) 251 | { 252 | mReportableSeverity = severity; 253 | } 254 | 255 | //! 256 | //! \brief Opaque handle that holds logging information for a particular test 257 | //! 258 | //! This object is an opaque handle to information used by the Logger to print test results. 259 | //! The sample must call Logger::defineTest() in order to obtain a TestAtom that can be used 260 | //! with Logger::reportTest{Start,End}(). 261 | //! 262 | class TestAtom 263 | { 264 | public: 265 | TestAtom(TestAtom&&) = default; 266 | 267 | private: 268 | friend class Logger; 269 | 270 | TestAtom(bool started, const std::string& name, const std::string& cmdline) 271 | : mStarted(started) 272 | , mName(name) 273 | , mCmdline(cmdline) 274 | { 275 | } 276 | 277 | bool mStarted; 278 | std::string mName; 279 | std::string mCmdline; 280 | }; 281 | 282 | //! 283 | //! \brief Define a test for logging 284 | //! 285 | //! \param[in] name The name of the test. This should be a string starting with 286 | //! "TensorRT" and containing dot-separated strings containing 287 | //! the characters [A-Za-z0-9_]. 288 | //! For example, "TensorRT.sample_googlenet" 289 | //! \param[in] cmdline The command line used to reproduce the test 290 | // 291 | //! \return a TestAtom that can be used in Logger::reportTest{Start,End}(). 292 | //! 293 | static TestAtom defineTest(const std::string& name, const std::string& cmdline) 294 | { 295 | return TestAtom(false, name, cmdline); 296 | } 297 | 298 | //! 299 | //! \brief A convenience overloaded version of defineTest() that accepts an array of command-line arguments 300 | //! as input 301 | //! 302 | //! \param[in] name The name of the test 303 | //! \param[in] argc The number of command-line arguments 304 | //! \param[in] argv The array of command-line arguments (given as C strings) 305 | //! 306 | //! \return a TestAtom that can be used in Logger::reportTest{Start,End}(). 307 | static TestAtom defineTest(const std::string& name, int argc, char const* const* argv) 308 | { 309 | auto cmdline = genCmdlineString(argc, argv); 310 | return defineTest(name, cmdline); 311 | } 312 | 313 | //! 314 | //! \brief Report that a test has started. 315 | //! 316 | //! \pre reportTestStart() has not been called yet for the given testAtom 317 | //! 318 | //! \param[in] testAtom The handle to the test that has started 319 | //! 320 | static void reportTestStart(TestAtom& testAtom) 321 | { 322 | reportTestResult(testAtom, TestResult::kRUNNING); 323 | assert(!testAtom.mStarted); 324 | testAtom.mStarted = true; 325 | } 326 | 327 | //! 328 | //! \brief Report that a test has ended. 329 | //! 330 | //! \pre reportTestStart() has been called for the given testAtom 331 | //! 332 | //! \param[in] testAtom The handle to the test that has ended 333 | //! \param[in] result The result of the test. Should be one of TestResult::kPASSED, 334 | //! TestResult::kFAILED, TestResult::kWAIVED 335 | //! 336 | static void reportTestEnd(const TestAtom& testAtom, TestResult result) 337 | { 338 | assert(result != TestResult::kRUNNING); 339 | assert(testAtom.mStarted); 340 | reportTestResult(testAtom, result); 341 | } 342 | 343 | static int reportPass(const TestAtom& testAtom) 344 | { 345 | reportTestEnd(testAtom, TestResult::kPASSED); 346 | return EXIT_SUCCESS; 347 | } 348 | 349 | static int reportFail(const TestAtom& testAtom) 350 | { 351 | reportTestEnd(testAtom, TestResult::kFAILED); 352 | return EXIT_FAILURE; 353 | } 354 | 355 | static int reportWaive(const TestAtom& testAtom) 356 | { 357 | reportTestEnd(testAtom, TestResult::kWAIVED); 358 | return EXIT_SUCCESS; 359 | } 360 | 361 | static int reportTest(const TestAtom& testAtom, bool pass) 362 | { 363 | return pass ? reportPass(testAtom) : reportFail(testAtom); 364 | } 365 | 366 | Severity getReportableSeverity() const 367 | { 368 | return mReportableSeverity; 369 | } 370 | 371 | private: 372 | //! 373 | //! \brief returns an appropriate string for prefixing a log message with the given severity 374 | //! 375 | static const char* severityPrefix(Severity severity) 376 | { 377 | switch (severity) 378 | { 379 | case Severity::kINTERNAL_ERROR: return "[F] "; 380 | case Severity::kERROR: return "[E] "; 381 | case Severity::kWARNING: return "[W] "; 382 | case Severity::kINFO: return "[I] "; 383 | case Severity::kVERBOSE: return "[V] "; 384 | default: assert(0); return ""; 385 | } 386 | } 387 | 388 | //! 389 | //! \brief returns an appropriate string for prefixing a test result message with the given result 390 | //! 391 | static const char* testResultString(TestResult result) 392 | { 393 | switch (result) 394 | { 395 | case TestResult::kRUNNING: return "RUNNING"; 396 | case TestResult::kPASSED: return "PASSED"; 397 | case TestResult::kFAILED: return "FAILED"; 398 | case TestResult::kWAIVED: return "WAIVED"; 399 | default: assert(0); return ""; 400 | } 401 | } 402 | 403 | //! 404 | //! \brief returns an appropriate output stream (cout or cerr) to use with the given severity 405 | //! 406 | static std::ostream& severityOstream(Severity severity) 407 | { 408 | return severity >= Severity::kINFO ? std::cout : std::cerr; 409 | } 410 | 411 | //! 412 | //! \brief method that implements logging test results 413 | //! 414 | static void reportTestResult(const TestAtom& testAtom, TestResult result) 415 | { 416 | severityOstream(Severity::kINFO) << "&&&& " << testResultString(result) << " " << testAtom.mName << " # " 417 | << testAtom.mCmdline << std::endl; 418 | } 419 | 420 | //! 421 | //! \brief generate a command line string from the given (argc, argv) values 422 | //! 423 | static std::string genCmdlineString(int argc, char const* const* argv) 424 | { 425 | std::stringstream ss; 426 | for (int i = 0; i < argc; i++) 427 | { 428 | if (i > 0) 429 | ss << " "; 430 | ss << argv[i]; 431 | } 432 | return ss.str(); 433 | } 434 | 435 | Severity mReportableSeverity; 436 | }; 437 | 438 | namespace 439 | { 440 | 441 | //! 442 | //! \brief produces a LogStreamConsumer object that can be used to log messages of severity kVERBOSE 443 | //! 444 | //! Example usage: 445 | //! 446 | //! LOG_VERBOSE(logger) << "hello world" << std::endl; 447 | //! 448 | inline LogStreamConsumer LOG_VERBOSE(const Logger& logger) 449 | { 450 | return LogStreamConsumer(logger.getReportableSeverity(), Severity::kVERBOSE); 451 | } 452 | 453 | //! 454 | //! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINFO 455 | //! 456 | //! Example usage: 457 | //! 458 | //! LOG_INFO(logger) << "hello world" << std::endl; 459 | //! 460 | inline LogStreamConsumer LOG_INFO(const Logger& logger) 461 | { 462 | return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINFO); 463 | } 464 | 465 | //! 466 | //! \brief produces a LogStreamConsumer object that can be used to log messages of severity kWARNING 467 | //! 468 | //! Example usage: 469 | //! 470 | //! LOG_WARN(logger) << "hello world" << std::endl; 471 | //! 472 | inline LogStreamConsumer LOG_WARN(const Logger& logger) 473 | { 474 | return LogStreamConsumer(logger.getReportableSeverity(), Severity::kWARNING); 475 | } 476 | 477 | //! 478 | //! \brief produces a LogStreamConsumer object that can be used to log messages of severity kERROR 479 | //! 480 | //! Example usage: 481 | //! 482 | //! LOG_ERROR(logger) << "hello world" << std::endl; 483 | //! 484 | inline LogStreamConsumer LOG_ERROR(const Logger& logger) 485 | { 486 | return LogStreamConsumer(logger.getReportableSeverity(), Severity::kERROR); 487 | } 488 | 489 | //! 490 | //! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINTERNAL_ERROR 491 | // ("fatal" severity) 492 | //! 493 | //! Example usage: 494 | //! 495 | //! LOG_FATAL(logger) << "hello world" << std::endl; 496 | //! 497 | inline LogStreamConsumer LOG_FATAL(const Logger& logger) 498 | { 499 | return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINTERNAL_ERROR); 500 | } 501 | 502 | } // anonymous namespace 503 | 504 | #endif // TENSORRT_LOGGING_H 505 | -------------------------------------------------------------------------------- /yolo/include/macros.h: -------------------------------------------------------------------------------- 1 | #ifndef __MACROS_H 2 | #define __MACROS_H 3 | 4 | #ifdef API_EXPORTS 5 | #if defined(_MSC_VER) 6 | #define API __declspec(dllexport) 7 | #else 8 | #define API __attribute__((visibility("default"))) 9 | #endif 10 | #else 11 | 12 | #if defined(_MSC_VER) 13 | #define API __declspec(dllimport) 14 | #else 15 | #define API 16 | #endif 17 | #endif // API_EXPORTS 18 | 19 | #if NV_TENSORRT_MAJOR >= 8 20 | #define TRT_NOEXCEPT noexcept 21 | #define TRT_CONST_ENQUEUE const 22 | #else 23 | #define TRT_NOEXCEPT 24 | #define TRT_CONST_ENQUEUE 25 | #endif 26 | 27 | #endif // __MACROS_H 28 | -------------------------------------------------------------------------------- /yolo/include/utils.h: -------------------------------------------------------------------------------- 1 | #ifndef TRTX_YOLOV5_UTILS_H_ 2 | #define TRTX_YOLOV5_UTILS_H_ 3 | 4 | #include 5 | #include 6 | 7 | static inline cv::Mat preprocess_img(cv::Mat& img, int input_w, int input_h) { 8 | int w, h, x, y; 9 | float r_w = input_w / (img.cols*1.0); 10 | float r_h = input_h / (img.rows*1.0); 11 | if (r_h > r_w) { 12 | w = input_w; 13 | h = r_w * img.rows; 14 | x = 0; 15 | y = (input_h - h) / 2; 16 | } else { 17 | w = r_h * img.cols; 18 | h = input_h; 19 | x = (input_w - w) / 2; 20 | y = 0; 21 | } 22 | cv::Mat re(h, w, CV_8UC3); 23 | cv::resize(img, re, re.size(), 0, 0, cv::INTER_LINEAR); 24 | cv::Mat out(input_h, input_w, CV_8UC3, cv::Scalar(128, 128, 128)); 25 | re.copyTo(out(cv::Rect(x, y, re.cols, re.rows))); 26 | return out; 27 | } 28 | 29 | static inline int read_files_in_dir(const char *p_dir_name, std::vector &file_names) { 30 | DIR *p_dir = opendir(p_dir_name); 31 | if (p_dir == nullptr) { 32 | return -1; 33 | } 34 | 35 | struct dirent* p_file = nullptr; 36 | while ((p_file = readdir(p_dir)) != nullptr) { 37 | if (strcmp(p_file->d_name, ".") != 0 && 38 | strcmp(p_file->d_name, "..") != 0) { 39 | //std::string cur_file_name(p_dir_name); 40 | //cur_file_name += "/"; 41 | //cur_file_name += p_file->d_name; 42 | std::string cur_file_name(p_file->d_name); 43 | file_names.push_back(cur_file_name); 44 | } 45 | } 46 | 47 | closedir(p_dir); 48 | return 0; 49 | } 50 | 51 | #endif // TRTX_YOLOV5_UTILS_H_ 52 | 53 | -------------------------------------------------------------------------------- /yolo/include/yololayer.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include "yololayer.h" 5 | #include "cuda_utils.h" 6 | 7 | namespace Tn 8 | { 9 | template 10 | void write(char*& buffer, const T& val) 11 | { 12 | *reinterpret_cast(buffer) = val; 13 | buffer += sizeof(T); 14 | } 15 | 16 | template 17 | void read(const char*& buffer, T& val) 18 | { 19 | val = *reinterpret_cast(buffer); 20 | buffer += sizeof(T); 21 | } 22 | } 23 | 24 | using namespace Yolo; 25 | 26 | namespace nvinfer1 27 | { 28 | YoloLayerPlugin::YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut, const std::vector& vYoloKernel) 29 | { 30 | mClassCount = classCount; 31 | mYoloV5NetWidth = netWidth; 32 | mYoloV5NetHeight = netHeight; 33 | mMaxOutObject = maxOut; 34 | mYoloKernel = vYoloKernel; 35 | mKernelCount = vYoloKernel.size(); 36 | 37 | CUDA_CHECK(cudaMallocHost(&mAnchor, mKernelCount * sizeof(void*))); 38 | size_t AnchorLen = sizeof(float)* CHECK_COUNT * 2; 39 | for (int ii = 0; ii < mKernelCount; ii++) 40 | { 41 | CUDA_CHECK(cudaMalloc(&mAnchor[ii], AnchorLen)); 42 | const auto& yolo = mYoloKernel[ii]; 43 | CUDA_CHECK(cudaMemcpy(mAnchor[ii], yolo.anchors, AnchorLen, cudaMemcpyHostToDevice)); 44 | } 45 | } 46 | YoloLayerPlugin::~YoloLayerPlugin() 47 | { 48 | for (int ii = 0; ii < mKernelCount; ii++) 49 | { 50 | CUDA_CHECK(cudaFree(mAnchor[ii])); 51 | } 52 | CUDA_CHECK(cudaFreeHost(mAnchor)); 53 | } 54 | 55 | // create the plugin at runtime from a byte stream 56 | YoloLayerPlugin::YoloLayerPlugin(const void* data, size_t length) 57 | { 58 | using namespace Tn; 59 | const char *d = reinterpret_cast(data), *a = d; 60 | read(d, mClassCount); 61 | read(d, mThreadCount); 62 | read(d, mKernelCount); 63 | read(d, mYoloV5NetWidth); 64 | read(d, mYoloV5NetHeight); 65 | read(d, mMaxOutObject); 66 | mYoloKernel.resize(mKernelCount); 67 | auto kernelSize = mKernelCount * sizeof(YoloKernel); 68 | memcpy(mYoloKernel.data(), d, kernelSize); 69 | d += kernelSize; 70 | CUDA_CHECK(cudaMallocHost(&mAnchor, mKernelCount * sizeof(void*))); 71 | size_t AnchorLen = sizeof(float)* CHECK_COUNT * 2; 72 | for (int ii = 0; ii < mKernelCount; ii++) 73 | { 74 | CUDA_CHECK(cudaMalloc(&mAnchor[ii], AnchorLen)); 75 | const auto& yolo = mYoloKernel[ii]; 76 | CUDA_CHECK(cudaMemcpy(mAnchor[ii], yolo.anchors, AnchorLen, cudaMemcpyHostToDevice)); 77 | } 78 | assert(d == a + length); 79 | } 80 | 81 | void YoloLayerPlugin::serialize(void* buffer) const TRT_NOEXCEPT 82 | { 83 | using namespace Tn; 84 | char* d = static_cast(buffer), *a = d; 85 | write(d, mClassCount); 86 | write(d, mThreadCount); 87 | write(d, mKernelCount); 88 | write(d, mYoloV5NetWidth); 89 | write(d, mYoloV5NetHeight); 90 | write(d, mMaxOutObject); 91 | auto kernelSize = mKernelCount * sizeof(YoloKernel); 92 | memcpy(d, mYoloKernel.data(), kernelSize); 93 | d += kernelSize; 94 | 95 | assert(d == a + getSerializationSize()); 96 | } 97 | 98 | size_t YoloLayerPlugin::getSerializationSize() const TRT_NOEXCEPT 99 | { 100 | return sizeof(mClassCount) + sizeof(mThreadCount) + sizeof(mKernelCount) + sizeof(Yolo::YoloKernel) * mYoloKernel.size() + sizeof(mYoloV5NetWidth) + sizeof(mYoloV5NetHeight) + sizeof(mMaxOutObject); 101 | } 102 | 103 | int YoloLayerPlugin::initialize() TRT_NOEXCEPT 104 | { 105 | return 0; 106 | } 107 | 108 | Dims YoloLayerPlugin::getOutputDimensions(int index, const Dims* inputs, int nbInputDims) TRT_NOEXCEPT 109 | { 110 | //output the result to channel 111 | int totalsize = mMaxOutObject * sizeof(Detection) / sizeof(float); 112 | 113 | return Dims3(totalsize + 1, 1, 1); 114 | } 115 | 116 | // Set plugin namespace 117 | void YoloLayerPlugin::setPluginNamespace(const char* pluginNamespace) TRT_NOEXCEPT 118 | { 119 | mPluginNamespace = pluginNamespace; 120 | } 121 | 122 | const char* YoloLayerPlugin::getPluginNamespace() const TRT_NOEXCEPT 123 | { 124 | return mPluginNamespace; 125 | } 126 | 127 | // Return the DataType of the plugin output at the requested index 128 | DataType YoloLayerPlugin::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const TRT_NOEXCEPT 129 | { 130 | return DataType::kFLOAT; 131 | } 132 | 133 | // Return true if output tensor is broadcast across a batch. 134 | bool YoloLayerPlugin::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const TRT_NOEXCEPT 135 | { 136 | return false; 137 | } 138 | 139 | // Return true if plugin can use input that is broadcast across batch without replication. 140 | bool YoloLayerPlugin::canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT 141 | { 142 | return false; 143 | } 144 | 145 | void YoloLayerPlugin::configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) TRT_NOEXCEPT 146 | { 147 | } 148 | 149 | // Attach the plugin object to an execution context and grant the plugin the access to some context resource. 150 | void YoloLayerPlugin::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) TRT_NOEXCEPT 151 | { 152 | } 153 | 154 | // Detach the plugin object from its execution context. 155 | void YoloLayerPlugin::detachFromContext() TRT_NOEXCEPT {} 156 | 157 | const char* YoloLayerPlugin::getPluginType() const TRT_NOEXCEPT 158 | { 159 | return "YoloLayer_TRT"; 160 | } 161 | 162 | const char* YoloLayerPlugin::getPluginVersion() const TRT_NOEXCEPT 163 | { 164 | return "1"; 165 | } 166 | 167 | void YoloLayerPlugin::destroy() TRT_NOEXCEPT 168 | { 169 | delete this; 170 | } 171 | 172 | // Clone the plugin 173 | IPluginV2IOExt* YoloLayerPlugin::clone() const TRT_NOEXCEPT 174 | { 175 | YoloLayerPlugin* p = new YoloLayerPlugin(mClassCount, mYoloV5NetWidth, mYoloV5NetHeight, mMaxOutObject, mYoloKernel); 176 | p->setPluginNamespace(mPluginNamespace); 177 | return p; 178 | } 179 | 180 | __device__ float Logist(float data) { return 1.0f / (1.0f + expf(-data)); }; 181 | 182 | __global__ void CalDetection(const float *input, float *output, int noElements, 183 | const int netwidth, const int netheight, int maxoutobject, int yoloWidth, int yoloHeight, const float anchors[CHECK_COUNT * 2], int classes, int outputElem) 184 | { 185 | 186 | int idx = threadIdx.x + blockDim.x * blockIdx.x; 187 | if (idx >= noElements) return; 188 | 189 | int total_grid = yoloWidth * yoloHeight; 190 | int bnIdx = idx / total_grid; 191 | idx = idx - total_grid * bnIdx; 192 | int info_len_i = 5 + classes; 193 | const float* curInput = input + bnIdx * (info_len_i * total_grid * CHECK_COUNT); 194 | 195 | for (int k = 0; k < CHECK_COUNT; ++k) { 196 | float box_prob = Logist(curInput[idx + k * info_len_i * total_grid + 4 * total_grid]); 197 | if (box_prob < IGNORE_THRESH) continue; 198 | int class_id = 0; 199 | float max_cls_prob = 0.0; 200 | for (int i = 5; i < info_len_i; ++i) { 201 | float p = Logist(curInput[idx + k * info_len_i * total_grid + i * total_grid]); 202 | if (p > max_cls_prob) { 203 | max_cls_prob = p; 204 | class_id = i - 5; 205 | } 206 | } 207 | float *res_count = output + bnIdx * outputElem; 208 | int count = (int)atomicAdd(res_count, 1); 209 | if (count >= maxoutobject) return; 210 | char *data = (char*)res_count + sizeof(float) + count * sizeof(Detection); 211 | Detection *det = (Detection*)(data); 212 | 213 | int row = idx / yoloWidth; 214 | int col = idx % yoloWidth; 215 | 216 | //Location 217 | // pytorch: 218 | // y = x[i].sigmoid() 219 | // y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i] # xy 220 | // y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh 221 | // X: (sigmoid(tx) + cx)/FeaturemapW * netwidth 222 | det->bbox[0] = (col - 0.5f + 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 0 * total_grid])) * netwidth / yoloWidth; 223 | det->bbox[1] = (row - 0.5f + 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 1 * total_grid])) * netheight / yoloHeight; 224 | 225 | // W: (Pw * e^tw) / FeaturemapW * netwidth 226 | // v5: https://github.com/ultralytics/yolov5/issues/471 227 | det->bbox[2] = 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 2 * total_grid]); 228 | det->bbox[2] = det->bbox[2] * det->bbox[2] * anchors[2 * k]; 229 | det->bbox[3] = 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 3 * total_grid]); 230 | det->bbox[3] = det->bbox[3] * det->bbox[3] * anchors[2 * k + 1]; 231 | det->conf = box_prob * max_cls_prob; 232 | det->class_id = class_id; 233 | } 234 | } 235 | 236 | void YoloLayerPlugin::forwardGpu(const float* const* inputs, float *output, cudaStream_t stream, int batchSize) 237 | { 238 | int outputElem = 1 + mMaxOutObject * sizeof(Detection) / sizeof(float); 239 | for (int idx = 0; idx < batchSize; ++idx) { 240 | CUDA_CHECK(cudaMemset(output + idx * outputElem, 0, sizeof(float))); 241 | } 242 | int numElem = 0; 243 | for (unsigned int i = 0; i < mYoloKernel.size(); ++i) { 244 | const auto& yolo = mYoloKernel[i]; 245 | numElem = yolo.width * yolo.height * batchSize; 246 | if (numElem < mThreadCount) mThreadCount = numElem; 247 | 248 | //printf("Net: %d %d \n", mYoloV5NetWidth, mYoloV5NetHeight); 249 | CalDetection << < (numElem + mThreadCount - 1) / mThreadCount, mThreadCount, 0, stream >> > 250 | (inputs[i], output, numElem, mYoloV5NetWidth, mYoloV5NetHeight, mMaxOutObject, yolo.width, yolo.height, (float*)mAnchor[i], mClassCount, outputElem); 251 | } 252 | } 253 | 254 | 255 | int YoloLayerPlugin::enqueue(int batchSize, const void* const* inputs, void* TRT_CONST_ENQUEUE* outputs, void* workspace, cudaStream_t stream) TRT_NOEXCEPT 256 | { 257 | forwardGpu((const float* const*)inputs, (float*)outputs[0], stream, batchSize); 258 | return 0; 259 | } 260 | 261 | PluginFieldCollection YoloPluginCreator::mFC{}; 262 | std::vector YoloPluginCreator::mPluginAttributes; 263 | 264 | YoloPluginCreator::YoloPluginCreator() 265 | { 266 | mPluginAttributes.clear(); 267 | 268 | mFC.nbFields = mPluginAttributes.size(); 269 | mFC.fields = mPluginAttributes.data(); 270 | } 271 | 272 | const char* YoloPluginCreator::getPluginName() const TRT_NOEXCEPT 273 | { 274 | return "YoloLayer_TRT"; 275 | } 276 | 277 | const char* YoloPluginCreator::getPluginVersion() const TRT_NOEXCEPT 278 | { 279 | return "1"; 280 | } 281 | 282 | const PluginFieldCollection* YoloPluginCreator::getFieldNames() TRT_NOEXCEPT 283 | { 284 | return &mFC; 285 | } 286 | 287 | IPluginV2IOExt* YoloPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) TRT_NOEXCEPT 288 | { 289 | assert(fc->nbFields == 2); 290 | assert(strcmp(fc->fields[0].name, "netinfo") == 0); 291 | assert(strcmp(fc->fields[1].name, "kernels") == 0); 292 | int *p_netinfo = (int*)(fc->fields[0].data); 293 | int class_count = p_netinfo[0]; 294 | int input_w = p_netinfo[1]; 295 | int input_h = p_netinfo[2]; 296 | int max_output_object_count = p_netinfo[3]; 297 | std::vector kernels(fc->fields[1].length); 298 | memcpy(&kernels[0], fc->fields[1].data, kernels.size() * sizeof(Yolo::YoloKernel)); 299 | YoloLayerPlugin* obj = new YoloLayerPlugin(class_count, input_w, input_h, max_output_object_count, kernels); 300 | obj->setPluginNamespace(mNamespace.c_str()); 301 | return obj; 302 | } 303 | 304 | IPluginV2IOExt* YoloPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength) TRT_NOEXCEPT 305 | { 306 | // This object will be deleted when the network is destroyed, which will 307 | // call YoloLayerPlugin::destroy() 308 | YoloLayerPlugin* obj = new YoloLayerPlugin(serialData, serialLength); 309 | obj->setPluginNamespace(mNamespace.c_str()); 310 | return obj; 311 | } 312 | } 313 | 314 | -------------------------------------------------------------------------------- /yolo/include/yololayer.h: -------------------------------------------------------------------------------- 1 | #ifndef _YOLO_LAYER_H 2 | #define _YOLO_LAYER_H 3 | 4 | #include 5 | #include 6 | #include 7 | #include "macros.h" 8 | 9 | namespace Yolo 10 | { 11 | static constexpr int CHECK_COUNT = 3; 12 | static constexpr float IGNORE_THRESH = 0.1f; 13 | struct YoloKernel 14 | { 15 | int width; 16 | int height; 17 | float anchors[CHECK_COUNT * 2]; 18 | }; 19 | static constexpr int MAX_OUTPUT_BBOX_COUNT = 1000; 20 | static constexpr int CLASS_NUM = 2; 21 | static constexpr int INPUT_H = 640; // yolov5's input height and width must be divisible by 32. 22 | static constexpr int INPUT_W = 640; 23 | 24 | static constexpr int LOCATIONS = 4; 25 | struct alignas(float) Detection { 26 | //center_x center_y w h 27 | float bbox[LOCATIONS]; 28 | float conf; // bbox_conf * cls_conf 29 | float class_id; 30 | }; 31 | } 32 | 33 | namespace nvinfer1 34 | { 35 | class API YoloLayerPlugin : public IPluginV2IOExt 36 | { 37 | public: 38 | YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut, const std::vector& vYoloKernel); 39 | YoloLayerPlugin(const void* data, size_t length); 40 | ~YoloLayerPlugin(); 41 | 42 | int getNbOutputs() const TRT_NOEXCEPT override 43 | { 44 | return 1; 45 | } 46 | 47 | Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) TRT_NOEXCEPT override; 48 | 49 | int initialize() TRT_NOEXCEPT override; 50 | 51 | virtual void terminate() TRT_NOEXCEPT override {}; 52 | 53 | virtual size_t getWorkspaceSize(int maxBatchSize) const TRT_NOEXCEPT override { return 0; } 54 | 55 | virtual int enqueue(int batchSize, const void* const* inputs, void*TRT_CONST_ENQUEUE* outputs, void* workspace, cudaStream_t stream) TRT_NOEXCEPT override; 56 | 57 | virtual size_t getSerializationSize() const TRT_NOEXCEPT override; 58 | 59 | virtual void serialize(void* buffer) const TRT_NOEXCEPT override; 60 | 61 | bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const TRT_NOEXCEPT override { 62 | return inOut[pos].format == TensorFormat::kLINEAR && inOut[pos].type == DataType::kFLOAT; 63 | } 64 | 65 | const char* getPluginType() const TRT_NOEXCEPT override; 66 | 67 | const char* getPluginVersion() const TRT_NOEXCEPT override; 68 | 69 | void destroy() TRT_NOEXCEPT override; 70 | 71 | IPluginV2IOExt* clone() const TRT_NOEXCEPT override; 72 | 73 | void setPluginNamespace(const char* pluginNamespace) TRT_NOEXCEPT override; 74 | 75 | const char* getPluginNamespace() const TRT_NOEXCEPT override; 76 | 77 | DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const TRT_NOEXCEPT override; 78 | 79 | bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const TRT_NOEXCEPT override; 80 | 81 | bool canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT override; 82 | 83 | void attachToContext( 84 | cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) TRT_NOEXCEPT override; 85 | 86 | void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) TRT_NOEXCEPT override; 87 | 88 | void detachFromContext() TRT_NOEXCEPT override; 89 | 90 | private: 91 | void forwardGpu(const float* const* inputs, float *output, cudaStream_t stream, int batchSize = 1); 92 | int mThreadCount = 256; 93 | const char* mPluginNamespace; 94 | int mKernelCount; 95 | int mClassCount; 96 | int mYoloV5NetWidth; 97 | int mYoloV5NetHeight; 98 | int mMaxOutObject; 99 | std::vector mYoloKernel; 100 | void** mAnchor; 101 | }; 102 | 103 | class API YoloPluginCreator : public IPluginCreator 104 | { 105 | public: 106 | YoloPluginCreator(); 107 | 108 | ~YoloPluginCreator() override = default; 109 | 110 | const char* getPluginName() const TRT_NOEXCEPT override; 111 | 112 | const char* getPluginVersion() const TRT_NOEXCEPT override; 113 | 114 | const PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override; 115 | 116 | IPluginV2IOExt* createPlugin(const char* name, const PluginFieldCollection* fc) TRT_NOEXCEPT override; 117 | 118 | IPluginV2IOExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) TRT_NOEXCEPT override; 119 | 120 | void setPluginNamespace(const char* libNamespace) TRT_NOEXCEPT override 121 | { 122 | mNamespace = libNamespace; 123 | } 124 | 125 | const char* getPluginNamespace() const TRT_NOEXCEPT override 126 | { 127 | return mNamespace.c_str(); 128 | } 129 | 130 | private: 131 | std::string mNamespace; 132 | static PluginFieldCollection mFC; 133 | static std::vector mPluginAttributes; 134 | }; 135 | REGISTER_TENSORRT_PLUGIN(YoloPluginCreator); 136 | }; 137 | 138 | #endif // _YOLO_LAYER_H 139 | -------------------------------------------------------------------------------- /yolo/include/yolov5_lib.h: -------------------------------------------------------------------------------- 1 | 2 | #pragma once 3 | #include "deepsort.h" 4 | #ifdef __cplusplus 5 | extern "C" 6 | { 7 | #endif 8 | 9 | void * yolov5_trt_create(const char * engine_name); 10 | 11 | int yolov5_trt_detect(void *h, cv::Mat &img, float threshold,std::vector& det); 12 | 13 | void yolov5_trt_destroy(void *h); 14 | 15 | #ifdef __cplusplus 16 | } 17 | #endif 18 | 19 | 20 | -------------------------------------------------------------------------------- /yolo/src/calibrator.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include "calibrator.h" 6 | #include "cuda_utils.h" 7 | #include "utils.h" 8 | 9 | Int8EntropyCalibrator2::Int8EntropyCalibrator2(int batchsize, int input_w, int input_h, const char* img_dir, const char* calib_table_name, const char* input_blob_name, bool read_cache) 10 | : batchsize_(batchsize) 11 | , input_w_(input_w) 12 | , input_h_(input_h) 13 | , img_idx_(0) 14 | , img_dir_(img_dir) 15 | , calib_table_name_(calib_table_name) 16 | , input_blob_name_(input_blob_name) 17 | , read_cache_(read_cache) 18 | { 19 | input_count_ = 3 * input_w * input_h * batchsize; 20 | CUDA_CHECK(cudaMalloc(&device_input_, input_count_ * sizeof(float))); 21 | read_files_in_dir(img_dir, img_files_); 22 | } 23 | 24 | Int8EntropyCalibrator2::~Int8EntropyCalibrator2() 25 | { 26 | CUDA_CHECK(cudaFree(device_input_)); 27 | } 28 | 29 | int Int8EntropyCalibrator2::getBatchSize() const TRT_NOEXCEPT 30 | { 31 | return batchsize_; 32 | } 33 | 34 | bool Int8EntropyCalibrator2::getBatch(void* bindings[], const char* names[], int nbBindings) TRT_NOEXCEPT 35 | { 36 | if (img_idx_ + batchsize_ > (int)img_files_.size()) { 37 | return false; 38 | } 39 | 40 | std::vector input_imgs_; 41 | for (int i = img_idx_; i < img_idx_ + batchsize_; i++) { 42 | std::cout << img_files_[i] << " " << i << std::endl; 43 | cv::Mat temp = cv::imread(img_dir_ + img_files_[i]); 44 | if (temp.empty()){ 45 | std::cerr << "Fatal error: image cannot open!" << std::endl; 46 | return false; 47 | } 48 | cv::Mat pr_img = preprocess_img(temp, input_w_, input_h_); 49 | input_imgs_.push_back(pr_img); 50 | } 51 | img_idx_ += batchsize_; 52 | cv::Mat blob = cv::dnn::blobFromImages(input_imgs_, 1.0 / 255.0, cv::Size(input_w_, input_h_), cv::Scalar(0, 0, 0), true, false); 53 | 54 | CUDA_CHECK(cudaMemcpy(device_input_, blob.ptr(0), input_count_ * sizeof(float), cudaMemcpyHostToDevice)); 55 | assert(!strcmp(names[0], input_blob_name_)); 56 | bindings[0] = device_input_; 57 | return true; 58 | } 59 | 60 | const void* Int8EntropyCalibrator2::readCalibrationCache(size_t& length) TRT_NOEXCEPT 61 | { 62 | std::cout << "reading calib cache: " << calib_table_name_ << std::endl; 63 | calib_cache_.clear(); 64 | std::ifstream input(calib_table_name_, std::ios::binary); 65 | input >> std::noskipws; 66 | if (read_cache_ && input.good()) 67 | { 68 | std::copy(std::istream_iterator(input), std::istream_iterator(), std::back_inserter(calib_cache_)); 69 | } 70 | length = calib_cache_.size(); 71 | return length ? calib_cache_.data() : nullptr; 72 | } 73 | 74 | void Int8EntropyCalibrator2::writeCalibrationCache(const void* cache, size_t length) TRT_NOEXCEPT 75 | { 76 | std::cout << "writing calib cache: " << calib_table_name_ << " size: " << length << std::endl; 77 | std::ofstream output(calib_table_name_, std::ios::binary); 78 | output.write(reinterpret_cast(cache), length); 79 | } 80 | 81 | -------------------------------------------------------------------------------- /yolo/src/yolov5_lib.cpp: -------------------------------------------------------------------------------- 1 | //yolov5_lib.cpp 2 | 3 | #include 4 | #include 5 | #include "cuda_runtime_api.h" 6 | #include "logging.h" 7 | #include "common.hpp" 8 | #include "yolov5_lib.h" 9 | #include "cuda_utils.h" 10 | #include "utils.h" 11 | #include "datatype.h" 12 | #define USE_FP16 // comment out this if want to use FP32 13 | #define DEVICE 0 // GPU id 14 | #define NMS_THRESH 0.4 15 | #define CONF_THRESH 0.5 16 | #define BATCH_SIZE 1 17 | 18 | // stuff we know about the network and the input/output blobs 19 | static const int INPUT_H = Yolo::INPUT_H; 20 | static const int INPUT_W = Yolo::INPUT_W; 21 | static const int CLASS_NUM = Yolo::CLASS_NUM; 22 | static const int OUTPUT_SIZE = Yolo::MAX_OUTPUT_BBOX_COUNT * sizeof(Yolo::Detection) / sizeof(float) + 1; // we assume the yololayer outputs no more than MAX_OUTPUT_BBOX_COUNT boxes that conf >= 0.1 23 | const char* INPUT_BLOB_NAME = "data"; 24 | const char* OUTPUT_BLOB_NAME = "prob"; 25 | static Logger gLogger; 26 | 27 | 28 | static void doInference(IExecutionContext& context, cudaStream_t& stream, void **buffers, float* input, float* output, int batchSize) { 29 | // DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host 30 | CUDA_CHECK(cudaMemcpyAsync(buffers[0], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream)); 31 | //cudaMemcpyAsync(buffers[0], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream); 32 | context.enqueue(batchSize, buffers, stream, nullptr); 33 | CUDA_CHECK(cudaMemcpyAsync(output, buffers[1], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream)); 34 | //cudaMemcpyAsync(output, buffers[1], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream); 35 | cudaStreamSynchronize(stream); 36 | } 37 | 38 | 39 | typedef struct 40 | { 41 | 42 | float *data; 43 | float *prob; 44 | IRuntime *runtime; 45 | ICudaEngine *engine; 46 | IExecutionContext *exe_context; 47 | void* buffers[2]; 48 | cudaStream_t cuda_stream; 49 | int inputIndex; 50 | int outputIndex; 51 | char result_json_str[16384]; 52 | 53 | }Yolov5TRTContext; 54 | 55 | typedef struct{ 56 | int class_id; 57 | int x1; 58 | int y1; 59 | int x2; 60 | int y2; 61 | float conf; 62 | 63 | }DeepsortContext; 64 | 65 | void * yolov5_trt_create(const char * engine_name) 66 | { 67 | size_t size = 0; 68 | char *trtModelStream = NULL; 69 | Yolov5TRTContext * trt_ctx = NULL; 70 | 71 | trt_ctx = new Yolov5TRTContext(); 72 | 73 | std::ifstream file(engine_name, std::ios::binary); 74 | printf("yolov5_trt_create ... \n"); 75 | if (file.good()) { 76 | file.seekg(0, file.end); 77 | size = file.tellg(); 78 | file.seekg(0, file.beg); 79 | trtModelStream = new char[size]; 80 | assert(trtModelStream); 81 | file.read(trtModelStream, size); 82 | file.close(); 83 | }else 84 | return NULL; 85 | 86 | trt_ctx->data = new float[BATCH_SIZE * 3 * INPUT_H * INPUT_W]; 87 | trt_ctx->prob = new float[BATCH_SIZE * OUTPUT_SIZE]; 88 | trt_ctx->runtime = createInferRuntime(gLogger); 89 | assert(trt_ctx->runtime != nullptr); 90 | 91 | printf("yolov5_trt_create cuda engine... \n"); 92 | trt_ctx->engine = trt_ctx->runtime->deserializeCudaEngine(trtModelStream, size); 93 | assert(trt_ctx->engine != nullptr); 94 | trt_ctx->exe_context = trt_ctx->engine->createExecutionContext(); 95 | 96 | 97 | delete[] trtModelStream; 98 | assert(trt_ctx->engine->getNbBindings() == 2); 99 | 100 | // In order to bind the buffers, we need to know the names of the input and output tensors. 101 | // Note that indices are guaranteed to be less than IEngine::getNbBindings() 102 | trt_ctx->inputIndex = trt_ctx->engine->getBindingIndex(INPUT_BLOB_NAME); 103 | trt_ctx->outputIndex = trt_ctx->engine->getBindingIndex(OUTPUT_BLOB_NAME); 104 | 105 | assert(trt_ctx->inputIndex == 0); 106 | assert(trt_ctx->outputIndex == 1); 107 | // Create GPU buffers on device 108 | 109 | printf("yolov5_trt_create buffer ... \n"); 110 | CUDA_CHECK(cudaMalloc(&trt_ctx->buffers[trt_ctx->inputIndex], BATCH_SIZE * 3 * INPUT_H * INPUT_W * sizeof(float))); 111 | //cudaMalloc(&trt_ctx->buffers[trt_ctx->inputIndex], BATCH_SIZE * 3 * INPUT_H * INPUT_W * sizeof(float)); 112 | CUDA_CHECK(cudaMalloc(&trt_ctx->buffers[trt_ctx->outputIndex], BATCH_SIZE * OUTPUT_SIZE * sizeof(float))); 113 | //cudaMalloc(&trt_ctx->buffers[trt_ctx->outputIndex], BATCH_SIZE * OUTPUT_SIZE * sizeof(float)); 114 | // Create stream 115 | 116 | printf("yolov5_trt_create stream ... \n"); 117 | CUDA_CHECK(cudaStreamCreate(&trt_ctx->cuda_stream)); 118 | //cudaStreamCreate(&trt_ctx->cuda_stream); 119 | printf("yolov5_trt_create done ... \n"); 120 | return (void *)trt_ctx; 121 | 122 | 123 | } 124 | 125 | 126 | int yolov5_trt_detect(void *h, cv::Mat &img, float threshold,std::vector& det) 127 | { 128 | Yolov5TRTContext *trt_ctx; 129 | int i; 130 | int delay_preprocess; 131 | int delay_infer; 132 | 133 | trt_ctx = (Yolov5TRTContext *)h; 134 | 135 | 136 | trt_ctx->result_json_str[0] = 0; 137 | // whether det is empty , if not, empty det 138 | if (!det.empty()) det.clear(); 139 | if (img.empty()) return 0; 140 | 141 | auto start0 = std::chrono::system_clock::now(); 142 | 143 | //printf("yolov5_trt_detect start preprocess img \n"); 144 | cv::Mat pr_img = preprocess_img(img, INPUT_W, INPUT_H); 145 | //std::cout<<"after preprocess_img pr_img size:"<data[i] = (float)uc_pixel[2] / 255.0; 156 | trt_ctx->data[i + INPUT_H * INPUT_W] = (float)uc_pixel[1] / 255.0; 157 | trt_ctx->data[i + 2 * INPUT_H * INPUT_W] = (float)uc_pixel[0] / 255.0; 158 | uc_pixel += 3; 159 | ++i; 160 | } 161 | } 162 | auto end0 = std::chrono::system_clock::now(); 163 | 164 | delay_preprocess = std::chrono::duration_cast(end0 - start0).count(); 165 | 166 | // Run inference 167 | //printf("yolov5_trt_detect start do inference\n"); 168 | auto start = std::chrono::system_clock::now(); 169 | doInference(*trt_ctx->exe_context, trt_ctx->cuda_stream, trt_ctx->buffers, trt_ctx->data, trt_ctx->prob, BATCH_SIZE); 170 | 171 | auto end = std::chrono::system_clock::now(); 172 | delay_infer = std::chrono::duration_cast(end - start).count(); 173 | 174 | std::cout <<"delay_proress:" << delay_preprocess << "ms, " << "delay_infer:" << delay_infer << "ms" << std::endl; 175 | 176 | //printf("yolov5_trt_detect start do process infer result \n"); 177 | 178 | int fcount = 1; 179 | int str_len; 180 | std::vector> batch_res(1); 181 | auto& res = batch_res[0]; 182 | nms(res, &trt_ctx->prob[0], threshold, NMS_THRESH); 183 | 184 | 185 | i = 0; 186 | for(i = 0 ; i < res.size(); i++){ 187 | int x1, y1, x2, y2; 188 | int class_id; 189 | float conf; 190 | cv::Rect r = get_rect(img, res[i].bbox); 191 | 192 | DetectBox dd(r.x,r.y,r.x + r.width,r.y + r.height,(float)res[i].conf,(int)res[i].class_id); 193 | det.push_back(dd); 194 | } 195 | return 1; 196 | } 197 | 198 | 199 | void yolov5_trt_destroy(void *h) 200 | { 201 | Yolov5TRTContext *trt_ctx; 202 | 203 | trt_ctx = (Yolov5TRTContext *)h; 204 | 205 | // Release stream and buffers 206 | cudaStreamDestroy(trt_ctx->cuda_stream); 207 | CUDA_CHECK(cudaFree(trt_ctx->buffers[trt_ctx->inputIndex])); 208 | //cudaFree(trt_ctx->buffers[trt_ctx->inputIndex]); 209 | CUDA_CHECK(cudaFree(trt_ctx->buffers[trt_ctx->outputIndex])); 210 | //cudaFree(trt_ctx->buffers[trt_ctx->outputIndex]) 211 | // Destroy the engine 212 | trt_ctx->exe_context->destroy(); 213 | trt_ctx->engine->destroy(); 214 | trt_ctx->runtime->destroy(); 215 | 216 | delete trt_ctx->data; 217 | delete trt_ctx->prob; 218 | 219 | delete trt_ctx; 220 | 221 | } 222 | 223 | --------------------------------------------------------------------------------