├── LICENSE ├── README.md ├── data ├── bus.jpg └── zidane.jpg ├── docker ├── Dockerfile └── readme.md ├── docs ├── API-Build.md ├── Jetson.md ├── Normal.md ├── Pose.md ├── Segment.md └── star.md ├── requirements.txt ├── scripts ├── .pre-commit-config.yaml ├── build.py ├── config.py ├── export-det.py ├── export-pose.py ├── export-seg.py ├── gen_pkl.py ├── infer-det-without-torch.py ├── infer-det.py ├── infer-pose-without-torch.py ├── infer-pose.py ├── infer-seg-without-torch.py ├── infer-seg.py ├── models │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-38.pyc │ │ ├── common.cpython-38.pyc │ │ └── engine.cpython-38.pyc │ ├── api.py │ ├── common.py │ ├── cudart_api.py │ ├── engine.py │ ├── pycuda_api.py │ ├── torch_utils.py │ └── utils.py └── trt-profile.py ├── src ├── deepstream │ ├── CMakeLists.txt │ ├── README.md │ ├── config_yoloV8.txt │ ├── custom_bbox_parser │ │ └── nvdsparsebbox_yoloV8.cpp │ ├── deepstream_app_config.txt │ └── labels.txt ├── detect │ ├── end2end │ │ ├── CMakeLists.txt │ │ ├── include │ │ │ ├── common.hpp │ │ │ └── yolov8.hpp │ │ └── main.cpp │ └── normal │ │ ├── CMakeLists.txt │ │ ├── include │ │ ├── common.hpp │ │ └── yolov8.hpp │ │ └── main.cpp ├── jetson │ ├── detect │ │ ├── CMakeLists.txt │ │ ├── include │ │ │ ├── common.hpp │ │ │ └── yolov8.hpp │ │ └── main.cpp │ ├── pose │ │ ├── CMakeLists.txt │ │ ├── include │ │ │ ├── common.hpp │ │ │ └── yolov8-pose.hpp │ │ └── main.cpp │ └── segment │ │ ├── CMakeLists.txt │ │ ├── include │ │ ├── common.hpp │ │ └── yolov8-seg.hpp │ │ └── main.cpp ├── pose │ └── normal │ │ ├── CMakeLists.txt │ │ ├── include │ │ ├── common.hpp │ │ └── yolov8-pose.hpp │ │ └── main.cpp └── segment │ ├── normal │ ├── CMakeLists.txt │ ├── include │ │ ├── common.hpp │ │ └── yolov8-seg.hpp │ └── main.cpp │ └── simple │ ├── CMakeLists.txt │ ├── include │ ├── common.hpp │ └── yolov8-seg.hpp │ └── main.cpp └── weights ├── .readme.md ├── yolov8n-pose.onnx ├── yolov8n-seg.onnx └── yolov8n.onnx /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 lin 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | 2 | # `YOLOv8-ROS-TensorRT-CPP` 3 | 4 | ros1-noetic版本请移步[noetic-devel分支](https://github.com/linClubs/YOLOv8-ROS-TensorRT/tree/noetic-devel) 5 | 6 | ros2-galactic版本请移步[galactic-devel分支](https://github.com/linClubs/YOLOv8-ROS-TensorRT/tree/galactic-devel) 7 | 8 | 9 | # 1. 环境配置 10 | 11 | ## 1.1 显卡驱动 `cuda` `cudnn`安装 12 | 13 | [安装参考链接](https://blog.csdn.net/h904798869/article/details/131719404) 14 | 15 | ## 1.2 基础环境安装 16 | ~~~python 17 | # 1 创建python虚拟环境 18 | conda create -n yolov8 python=3.8 19 | # 2 激活虚拟环境 20 | 21 | # 3 安装torch-1.10.0 torchvision==0.11 22 | pip install torch==1.10.0+cu113 torchvision==0.11.0+cu113 torchaudio==0.10.0 -f https://download.pytorch.org/whl/torch_stable.html 23 | 24 | # 4 其他依赖 25 | pip install numpy==1.23.5 onnx==1.14.1 onnxsim opencv-python==4.5.4.58 ultralytics onnxruntime==1.16.0 26 | # mpmath, flatbuffers, sympy, humanfriendly, coloredlogs 27 | 28 | # 5 tensorrt安装参考下面 29 | ~~~ 30 | 31 | + `TensorRT`安装 32 | 33 | [安装参考安装TensorRT部分](https://blog.csdn.net/h904798869/article/details/131719404) 34 | 35 | [下载地址](https://link.csdn.net/?target=https%3A%2F%2Fdeveloper.nvidia.com%2Fnvidia-tensorrt-download) 需要注册账号登录才能下载 36 | 37 | ~~~python 38 | # 1 解压TensorRT得到TensorRT-8.5.3.1目录 39 | tar -xf TensorRT-8.5.3.1.Linux.x86_64-gnu.cuda-11.8.cudnn8.6.tar.gz 40 | 41 | # 2 安装python版tensorrt 42 | # 2.1 进入TensorRT-8.5.3.1/python目录 43 | cd TensorRT-8.5.3.1/python 44 | # 2.2 安装python3.8支持的tensorrt, 选择cp38即可 45 | pip install tensorrt-8.5.3.1-cp38-none-linux_x86_64.whl 46 | 47 | # 3 c++ 版TensorRT直接使用这个文件, 48 | # 修改CMakeLists.txt的TensorRT_ROOT路径, 改成自己对应TensorRT所在的路径即可使用 49 | set(TensorRT_ROOT /path/to/TensorRT-8.5.3.1) 50 | ~~~ 51 | 52 | # 1.3 报错汇总 53 | 54 | 1. 运行`python`脚本报错找不到`TensorRT`相关的库文件 55 | 56 | + 修改方法2： 57 | 将缺的`TensorRT`相关的库文件放进`/usr/lib` 58 | 59 | ~~~python 60 | # 1 错误1 61 | ImportError: libnvinfer.so.8: cannot open shared object file: No such file or directory 62 | # 修改将libnvinfer.so.8放入/usr/lib 63 | sudo cp TensorRT-8.5.3.1/targets/x86_64-linux-gnu/lib/libnvinfer.so.8 /usr/lib 64 | 65 | # 2 错误2 66 | mportError: libnvonnxparser.so.8: cannot open shared object file: No such file or directory 67 | # 修改将libnvonnxparser.so.8放入/usr/lib 68 | sudo cp TensorRT-8.5.3.1/targets/x86_64-linux-gnu/lib/libnvonnxparser.so.8 /usr/lib 69 | 70 | # 错误3 71 | libnvparsers.so.8：cannot open shared object file: No such file or directory 72 | # 修改将libnvparsers.so.8放入/usr/lib 73 | sudo cp TensorRT-8.5.3.1/targets/x86_64-linux-gnu/lib/libnvparsers.so.8 /usr/lib 74 | 75 | # 错误4 76 | libcudnn.so.8 cannot open shared object file 77 | # 把/usr/local/cuda/lib64/libcudnn.so.8文件放入/usr/lib 78 | 79 | # 错误5 c++推理代码运行报错 80 | fatal error: NvInferPlugin.h: No such file or directory 81 | # 修改CMakeLists.txt的TensorRT_ROOT路径 82 | set(TensorRT_ROOT /root/share/TensorRT-8.5.3.1) 83 | set(TensorRT_INCLUDE_DIRS ${TensorRT_ROOT}/include) 84 | set(TensorRT_LIBRARIES ${TensorRT_ROOT}/lib) 85 | 86 | # 报错6 87 | TypeError: pybind11::init(): factory function returned nullptr 88 | ~~~ 89 | 90 | 91 | # 2 pt2onnx 92 | 93 | `pt`文件转`onnx`文件 94 | 95 | ## 2.1 下载`pt`文件 96 | 97 | 下载`pt`文件详情见[官方地址readme](https://github.com/ultralytics/ultralytics/blob/main/README.md) 98 | 99 | 可以直接点击下面链接下载`yolov8n`相关模型： 100 | 101 | [yolov8n-det](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n.pt) 102 | 103 | [yolov8n-seg](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n-seg.pt) 104 | 105 | [yolov8n-pose](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n-pose.pt) 106 | 107 | 108 | ## 1.2 `det and seg` 模型转换 109 | ~~~python 110 | # 1 进入scripts目录 111 | cd scripts 112 | 113 | # 2 det 114 | python export-det.py --weights ../weights/yolov8n.pt --iou-thres 0.65 --conf-thres 0.2 --topk 100 --opset 11 --sim --input-shape 1 3 640 640 --device cuda:0 115 | 116 | # 3 seg 117 | python export-seg.py --weights ../weights/yolov8n-seg.pt --opset 11 --sim --input-shape 1 3 640 640 --device cuda:0 118 | 119 | # --topk最大检测框数 120 | # -input-shape输入尺寸 121 | ~~~ 122 | 123 | ## 1.3 `pose`模型转换 124 | 125 | 1. `scripts`目录下新建一个转换脚本`export-pose.py`文件，填入下面内容 126 | 127 | ~~~python 128 | from ultralytics import YOLO 129 | # Load a model 130 | model = YOLO("../weights/yolov8n-pose.pt") # load a pretrained model (recommended for training) 131 | success = model.export(format="onnx", opset=11, simplify=True) # export the model to onnx format 132 | assert success 133 | ~~~ 134 | 135 | 2. 运行`export-pose.py`生成`onnx`文件 136 | 137 | ~~~python 138 | python export-pose.py 139 | ~~~ 140 | 141 | 生成的`onnx`与对应的`pt`在同级目录 142 | 143 | 144 | # 2 `onnx2engine` 145 | 146 | + `onnx`文件转`engine`格式文件 147 | 148 | ~~~python 149 | # 1 进入scripts目录 150 | cd scripts 151 | 152 | # 2 det 153 | python build.py --weights ../weights/yolov8n.onnx --iou-thres 0.65 --conf-thres 0.25 --topk 100 --fp16 --device cuda:0 154 | 155 | # 3 seg, 增加--seg参数 156 | python build.py --weights ../weights/yolov8n-seg.onnx --iou-thres 0.65 --conf-thres 0.25 --topk 100 --fp16 --device cuda:0 --seg 157 | 158 | # 4 pose 使用tensorRT 159 | cd tensorRT/bin 160 | ./trtexec --onnx=yolov8n-pose.onnx --saveEngine=yolov8n-pose.engine 161 | ~~~ 162 | 163 | + 等待时间比较久，只要电脑没开，都正常 164 | 165 | # 3 `c++`模型推理 166 | 167 | 1. 配置`TensorRT`环境变量 168 | 169 | + c++编译前,需要修改`CMakeLists.txt`中`TensorRT_ROOT`路径, 改成自己对应TensorRT所在的路径即可使用 170 | 171 | ~~~c 172 | set(TensorRT_ROOT /path/to/TensorRT-8.5.3.1) 173 | ~~~ 174 | 175 | 2. 编译源码 176 | 177 | + `src/detect/end2end`检测代码为例 178 | 179 | ~~~python 180 | # 1 进入目录 181 | cd src/detect/end2end 182 | 183 | # 2 cmake编译 184 | mkdir build && cd build && cmake .. && make -j 185 | 186 | # 3 pose代码就进入src/pose/normal进行cmake编译 187 | ~~~ 188 | 189 | 3. `inference`推理 190 | 191 | + 注意权重路径和待检测图像路径给正确即可 192 | 193 | ~~~python 194 | # det-python版本 需要进入scripts路径运行 195 | python3 infer-det.py --engine yolov8s.engine --imgs data --show --out-dir outputs 196 | --device cuda:0 197 | 198 | # 1 检测 det-c++版本 199 | # 1.1 infer image 200 | ./yolov8 yolov8n.engine data/bus.jpg 201 | # 1.2 infer images 202 | ./yolov8 yolov8n.engine data 203 | # 1.3 infer video 204 | ./yolov8 yolov8n.engine data/test.mp4 # the video path 205 | 206 | # 2 分割 seg 207 | ./yolov8-seg weights/yolov8n-seg.engine data/zidane.jpg 208 | 209 | # 3 姿态预测 pose 210 | ./yolov8-pose weights/yolov8n-pose.engine data/bus.jpg 211 | ~~~ 212 | 213 | --- 214 | -------------------------------------------------------------------------------- /data/bus.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linClubs/YOLOv8-ROS-TensorRT/0b23e9753b2cf0a2f26eda2335d867f1f2f69dda/data/bus.jpg -------------------------------------------------------------------------------- /data/zidane.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linClubs/YOLOv8-ROS-TensorRT/0b23e9753b2cf0a2f26eda2335d867f1f2f69dda/data/zidane.jpg -------------------------------------------------------------------------------- /docker/Dockerfile: -------------------------------------------------------------------------------- 1 | 2 | FROM nvidia/cuda:11.3.1-devel-ubuntu20.04 3 | 4 | # 1 setup environment 5 | ENV LANG C.UTF-8 6 | ENV LC_ALL C.UTF-8 7 | ENV DEBIAN_FRONTEND noninteractive 8 | 9 | # 2 apt换源 10 | RUN echo "deb http://mirrors.ustc.edu.cn/ubuntu/ focal main restricted universe multiverse" > /etc/apt/sources.list \ 11 | && echo "deb http://mirrors.ustc.edu.cn/ubuntu/ focal-updates main restricted universe multiverse" >> /etc/apt/sources.list \ 12 | && echo "deb http://mirrors.ustc.edu.cn/ubuntu/ focal-backports main restricted universe multiverse" >> /etc/apt/sources.list \ 13 | && echo "deb http://mirrors.ustc.edu.cn/ubuntu/ focal-security main restricted universe multiverse" >> /etc/apt/sources.list 14 | 15 | # 3 更新源和安装wget和git 16 | RUN apt update && apt install wget cmake -y 17 | 18 | # ImportError: libGL.so.1:找不到 19 | RUN apt-get install ffmpeg libsm6 libxext6 -y 20 | 21 | # 4 安装conda并初始化 22 | RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh \ 23 | && chmod +x Miniconda3-latest-Linux-x86_64.sh \ 24 | && ./Miniconda3-latest-Linux-x86_64.sh -b \ 25 | && rm -rf Miniconda3-latest-Linux-x86_64.sh \ 26 | && ~/miniconda3/bin/conda init 27 | 28 | # 5 conda和pip换国内源 29 | RUN ~/miniconda3/bin/conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/ \ 30 | && ~/miniconda3/bin/conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge \ 31 | && ~/miniconda3/bin/conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/msys2/ \ 32 | && ~/miniconda3/bin/conda config --set show_channel_urls yes \ 33 | && mkdir ~/.pip && touch ~/.pip/pip.conf \ 34 | && echo "[global]" > ~/.pip/pip.conf \ 35 | && echo "index-url = https://pypi.tuna.tsinghua.edu.cn/simple/" >> ~/.pip/pip.conf 36 | 37 | # 6 其他依赖 38 | RUN apt install git vim sudo kmod 39 | RUN apt install libopencv-dev 40 | 41 | WORKDIR /root -------------------------------------------------------------------------------- /docker/readme.md: -------------------------------------------------------------------------------- 1 | ~~~python 2 | # 构建镜像 3 | docker build -f Dockerfile -t yolo:v1 . 4 | 5 | # 新建容器 6 | sudo docker run -it --gpus all -v ~/share:/root/share --name test yolo:v1 /bin/bash 7 | ~~~ -------------------------------------------------------------------------------- /docs/API-Build.md: -------------------------------------------------------------------------------- 1 | # Build TensorRT Engine By TensorRT Python API 2 | 3 | When you want to build engine by API. You should generate the pickle weights parameters first. 4 | 5 | ``` shell 6 | python3 gen_pkl.py -w yolov8s.pt -o yolov8s.pkl 7 | ``` 8 | 9 | You will get a `yolov8s.pkl` which contain the operators' parameters. 10 | 11 | And you can rebuild `yolov8s` model in TensorRT api. 12 | 13 | ``` 14 | python3 build.py \ 15 | --weights yolov8s.pkl \ 16 | --iou-thres 0.65 \ 17 | --conf-thres 0.25 \ 18 | --topk 100 \ 19 | --fp16 \ 20 | --input-shape 1 3 640 640 \ 21 | --device cuda:0 22 | ``` 23 | 24 | ***Notice !!!*** 25 | 26 | Now we only support static input shape model build by TensorRT api. 27 | You'd best give the legal `input-shape`. 28 | 29 | ***Notice !!!*** 30 | 31 | Now we don't support YOLOv8-seg model building by API. It will be supported later. 32 | -------------------------------------------------------------------------------- /docs/Jetson.md: -------------------------------------------------------------------------------- 1 | # YOLOv8 on Jetson 2 | 3 | Only test on `Jetson-NX 4GB` 4 | 5 | ENVS: 6 | 7 | - Jetpack 4.6.3 8 | - CUDA-10.2 9 | - CUDNN-8.2.1 10 | - TensorRT-8.2.1 11 | - DeepStream-6.0.1 12 | - OpenCV-4.1.1 13 | - CMake-3.10.2 14 | 15 | If you have other environment-related issues, please discuss in issue. 16 | 17 | ## End2End Detection 18 | 19 | ### 1. Export Detection End2End ONNX 20 | 21 | `yolov8s.pt` is your trained pytorch model, or the official pre-trained model. 22 | 23 | Do not use any model other than pytorch model. 24 | Do not use [`build.py`](../build.py) to export engine if you don't know how to install pytorch and other environments on 25 | jetson. 26 | 27 | ***!!! Please use the PC to execute the following script !!!*** 28 | 29 | ```shell 30 | # Export yolov8s.pt to yolov8s.onnx 31 | python3 export-det.py --weights yolov8s.pt --sim 32 | ``` 33 | 34 | ***!!! Please use the Jetson to execute the following script !!!*** 35 | 36 | ```shell 37 | # Using trtexec tools for export engine 38 | /usr/src/tensorrt/bin/trtexec \ 39 | --onnx=yolov8s.onnx \ 40 | --saveEngine=yolov8s.engine 41 | ``` 42 | 43 | After executing the above command, you will get an engine named `yolov8s.engine` . 44 | 45 | ### 2. Inference with c++ 46 | 47 | It is highly recommended to use C++ inference on Jetson. 48 | Here is a demo: [`csrc/jetson/detect`](../csrc/jetson/detect) . 49 | 50 | #### Build: 51 | 52 | Please modify `CLASS_NAMES` and `COLORS` in [`main.cpp`](../csrc/jetson/detect/main.cpp) for yourself. 53 | 54 | And build: 55 | 56 | ``` shell 57 | export root=${PWD} 58 | cd src/jetson/detect 59 | mkdir build 60 | cmake .. 61 | make 62 | mv yolov8 ${root} 63 | cd ${root} 64 | ``` 65 | 66 | Usage: 67 | 68 | ``` shell 69 | # infer image 70 | ./yolov8 yolov8s.engine data/bus.jpg 71 | # infer images 72 | ./yolov8 yolov8s.engine data 73 | # infer video 74 | ./yolov8 yolov8s.engine data/test.mp4 # the video path 75 | ``` 76 | 77 | ## Speedup Segmention 78 | 79 | ### 1. Export Segmention Speedup ONNX 80 | 81 | `yolov8s-seg.pt` is your trained pytorch model, or the official pre-trained model. 82 | 83 | Do not use any model other than pytorch model. 84 | Do not use [`build.py`](../build.py) to export engine if you don't know how to install pytorch and other environments on 85 | jetson. 86 | 87 | ***!!! Please use the PC to execute the following script !!!*** 88 | 89 | ```shell 90 | # Export yolov8s-seg.pt to yolov8s-seg.onnx 91 | python3 export-seg.py --weights yolov8s-seg.pt --sim 92 | ``` 93 | 94 | ***!!! Please use the Jetson to execute the following script !!!*** 95 | 96 | ```shell 97 | # Using trtexec tools for export engine 98 | /usr/src/tensorrt/bin/trtexec \ 99 | --onnx=yolov8s-seg.onnx \ 100 | --saveEngine=yolov8s-seg.engine 101 | ``` 102 | 103 | After executing the above command, you will get an engine named `yolov8s-seg.engine` . 104 | 105 | ### 2. Inference with c++ 106 | 107 | It is highly recommended to use C++ inference on Jetson. 108 | Here is a demo: [`csrc/jetson/segment`](../csrc/jetson/segment) . 109 | 110 | #### Build: 111 | 112 | Please modify `CLASS_NAMES` and `COLORS` and postprocess parameters in [`main.cpp`](../csrc/jetson/segment/main.cpp) for 113 | yourself. 114 | 115 | ```c++ 116 | int topk = 100; 117 | int seg_h = 160; // yolov8 model proto height 118 | int seg_w = 160; // yolov8 model proto width 119 | int seg_channels = 32; // yolov8 model proto channels 120 | float score_thres = 0.25f; 121 | float iou_thres = 0.65f; 122 | ``` 123 | 124 | And build: 125 | 126 | ``` shell 127 | export root=${PWD} 128 | cd src/jetson/segment 129 | mkdir build 130 | cmake .. 131 | make 132 | mv yolov8-seg ${root} 133 | cd ${root} 134 | ``` 135 | 136 | Usage: 137 | 138 | ``` shell 139 | # infer image 140 | ./yolov8-seg yolov8s-seg.engine data/bus.jpg 141 | # infer images 142 | ./yolov8-seg yolov8s-seg.engine data 143 | # infer video 144 | ./yolov8-seg yolov8s-seg.engine data/test.mp4 # the video path 145 | ``` 146 | 147 | ## Normal Posture 148 | 149 | ### 1. Export Posture Normal ONNX 150 | 151 | `yolov8s-pose.pt` is your trained pytorch model, or the official pre-trained model. 152 | 153 | Do not use any model other than pytorch model. 154 | Do not use [`build.py`](../build.py) to export engine if you don't know how to install pytorch and other environments on 155 | jetson. 156 | 157 | ***!!! Please use the PC to execute the following script !!!*** 158 | 159 | ```shell 160 | # Export yolov8s-pose.pt to yolov8s-pose.onnx 161 | yolo export model=yolov8s-pose.pt format=onnx simplify=True 162 | ``` 163 | 164 | ***!!! Please use the Jetson to execute the following script !!!*** 165 | 166 | ```shell 167 | # Using trtexec tools for export engine 168 | /usr/src/tensorrt/bin/trtexec \ 169 | --onnx=yolov8s-pose.onnx \ 170 | --saveEngine=yolov8s-pose.engine 171 | ``` 172 | 173 | After executing the above command, you will get an engine named `yolov8s-pose.engine` . 174 | 175 | ### 2. Inference with c++ 176 | 177 | It is highly recommended to use C++ inference on Jetson. 178 | Here is a demo: [`csrc/jetson/pose`](../csrc/jetson/pose) . 179 | 180 | #### Build: 181 | 182 | Please modify `KPS_COLORS` and `SKELETON` and `LIMB_COLORS` and postprocess parameters 183 | in [`main.cpp`](../csrc/jetson/pose/main.cpp) for yourself. 184 | 185 | ```c++ 186 | int topk = 100; 187 | float score_thres = 0.25f; 188 | float iou_thres = 0.65f; 189 | ``` 190 | 191 | And build: 192 | 193 | ``` shell 194 | export root=${PWD} 195 | cd src/jetson/pose 196 | mkdir build 197 | cmake .. 198 | make 199 | mv yolov8-pose ${root} 200 | cd ${root} 201 | ``` 202 | 203 | Usage: 204 | 205 | ``` shell 206 | # infer image 207 | ./yolov8-pose yolov8s-pose.engine data/bus.jpg 208 | # infer images 209 | ./yolov8-pose yolov8s-pose.engine data 210 | # infer video 211 | ./yolov8-pose yolov8s-pose.engine data/test.mp4 # the video path 212 | ``` 213 | -------------------------------------------------------------------------------- /docs/Normal.md: -------------------------------------------------------------------------------- 1 | # Normal Usage of [`ultralytics`](https://github.com/ultralytics/ultralytics) 2 | 3 | ## Export TensorRT Engine 4 | 5 | ### 1. ONNX -> TensorRT 6 | 7 | You can export your onnx model by `ultralytics` API. 8 | 9 | ``` shell 10 | yolo export model=yolov8s.pt format=onnx opset=11 simplify=True 11 | ``` 12 | 13 | or run this python script: 14 | 15 | ```python 16 | from ultralytics import YOLO 17 | 18 | # Load a model 19 | model = YOLO("yolov8s.pt") # load a pretrained model (recommended for training) 20 | success = model.export(format="onnx", opset=11, simplify=True) # export the model to onnx format 21 | assert success 22 | ``` 23 | 24 | Then build engine by Trtexec Tools. 25 | 26 | You can export TensorRT engine by [`trtexec`](https://github.com/NVIDIA/TensorRT/tree/main/samples/trtexec) tools. 27 | 28 | Usage: 29 | 30 | ``` shell 31 | /usr/src/tensorrt/bin/trtexec \ 32 | --onnx=yolov8s.onnx \ 33 | --saveEngine=yolov8s.engine \ 34 | --fp16 35 | ``` 36 | 37 | ### 2. Direct to TensorRT (NOT RECOMMAND!!) 38 | 39 | Usage: 40 | 41 | ```shell 42 | yolo export model=yolov8s.pt format=engine device=0 43 | ``` 44 | 45 | or run python script: 46 | 47 | ```python 48 | 49 | from ultralytics import YOLO 50 | 51 | # Load a model 52 | model = YOLO("yolov8s.pt") # load a pretrained model (recommended for training) 53 | success = model.export(format="engine", device=0) # export the model to engine format 54 | assert success 55 | ``` 56 | 57 | After executing the above script, you will get an engine named `yolov8s.engine` . 58 | 59 | ## Inference with c++ 60 | 61 | You can infer with c++ in [`csrc/detect/normal`](../csrc/detect/normal) . 62 | 63 | ### Build: 64 | 65 | Please set you own librarys in [`CMakeLists.txt`](../csrc/detect/normal/CMakeLists.txt) and modify `CLASS_NAMES` 66 | and `COLORS` in [`main.cpp`](../csrc/detect/normal/main.cpp). 67 | 68 | Besides, you can modify the postprocess parameters such as `num_labels` and `score_thres` and `iou_thres` and `topk` 69 | in [`main.cpp`](../csrc/detect/normal/main.cpp). 70 | 71 | ```c++ 72 | int num_labels = 80; 73 | int topk = 100; 74 | float score_thres = 0.25f; 75 | float iou_thres = 0.65f; 76 | ``` 77 | 78 | And build: 79 | 80 | ``` shell 81 | export root=${PWD} 82 | cd src/detect/normal 83 | mkdir build 84 | cmake .. 85 | make 86 | mv yolov8 ${root} 87 | cd ${root} 88 | ``` 89 | 90 | Usage: 91 | 92 | ``` shell 93 | # infer image 94 | ./yolov8 yolov8s.engine data/bus.jpg 95 | # infer images 96 | ./yolov8 yolov8s.engine data 97 | # infer video 98 | ./yolov8 yolov8s.engine data/test.mp4 # the video path 99 | ``` 100 | -------------------------------------------------------------------------------- /docs/Pose.md: -------------------------------------------------------------------------------- 1 | # YOLOv8-pose Model with TensorRT 2 | 3 | The yolov8-pose model conversion route is : 4 | YOLOv8 PyTorch model -> ONNX -> TensorRT Engine 5 | 6 | ***Notice !!!*** We don't support TensorRT API building !!! 7 | 8 | # Export Orin ONNX model by ultralytics 9 | 10 | You can leave this repo and use the original `ultralytics` repo for onnx export. 11 | 12 | ### 1. ONNX -> TensorRT 13 | 14 | You can export your onnx model by `ultralytics` API. 15 | 16 | ``` shell 17 | yolo export model=yolov8s-pose.pt format=onnx opset=11 simplify=True 18 | ``` 19 | 20 | or run this python script: 21 | 22 | ```python 23 | from ultralytics import YOLO 24 | 25 | # Load a model 26 | model = YOLO("yolov8s-pose.pt") # load a pretrained model (recommended for training) 27 | success = model.export(format="onnx", opset=11, simplify=True) # export the model to onnx format 28 | assert success 29 | ``` 30 | 31 | Then build engine by Trtexec Tools. 32 | 33 | You can export TensorRT engine by [`trtexec`](https://github.com/NVIDIA/TensorRT/tree/main/samples/trtexec) tools. 34 | 35 | Usage: 36 | 37 | ``` shell 38 | /usr/src/tensorrt/bin/trtexec \ 39 | --onnx=yolov8s-pose.onnx \ 40 | --saveEngine=yolov8s-pose.engine \ 41 | --fp16 42 | ``` 43 | 44 | ### 2. Direct to TensorRT (NOT RECOMMAND!!) 45 | 46 | Usage: 47 | 48 | ```shell 49 | yolo export model=yolov8s-pose.pt format=engine device=0 50 | ``` 51 | 52 | or run python script: 53 | 54 | ```python 55 | from ultralytics import YOLO 56 | 57 | # Load a model 58 | model = YOLO("yolov8s-pose.pt") # load a pretrained model (recommended for training) 59 | success = model.export(format="engine", device=0) # export the model to engine format 60 | assert success 61 | ``` 62 | 63 | After executing the above script, you will get an engine named `yolov8s-pose.engine` . 64 | 65 | # Inference 66 | 67 | ## Infer with python script 68 | 69 | You can infer images with the engine by [`infer-pose.py`](../infer-pose.py) . 70 | 71 | Usage: 72 | 73 | ``` shell 74 | python3 infer-pose.py \ 75 | --engine yolov8s-pose.engine \ 76 | --imgs data \ 77 | --show \ 78 | --out-dir outputs \ 79 | --device cuda:0 80 | ``` 81 | 82 | #### Description of all arguments 83 | 84 | - `--engine` : The Engine you export. 85 | - `--imgs` : The images path you want to detect. 86 | - `--show` : Whether to show detection results. 87 | - `--out-dir` : Where to save detection results images. It will not work when use `--show` flag. 88 | - `--device` : The CUDA deivce you use. 89 | 90 | ## Inference with c++ 91 | 92 | You can infer with c++ in [`csrc/pose/normal`](../csrc/pose/normal) . 93 | 94 | ### Build: 95 | 96 | Please set you own librarys in [`CMakeLists.txt`](../csrc/pose/normal/CMakeLists.txt) and modify `KPS_COLORS` 97 | and `SKELETON` and `LIMB_COLORS` in [`main.cpp`](../csrc/pose/normal/main.cpp). 98 | 99 | Besides, you can modify the postprocess parameters such as `score_thres` and `iou_thres` and `topk` 100 | in [`main.cpp`](../csrc/pose/normal/main.cpp). 101 | 102 | ```c++ 103 | int topk = 100; 104 | float score_thres = 0.25f; 105 | float iou_thres = 0.65f; 106 | ``` 107 | 108 | And build: 109 | 110 | ``` shell 111 | export root=${PWD} 112 | cd src/pose/normal 113 | mkdir build 114 | cmake .. 115 | make 116 | mv yolov8-pose ${root} 117 | cd ${root} 118 | ``` 119 | 120 | Usage: 121 | 122 | ``` shell 123 | # infer image 124 | ./yolov8-pose yolov8s-pose.engine data/bus.jpg 125 | # infer images 126 | ./yolov8-pose yolov8s-pose.engine data 127 | # infer video 128 | ./yolov8-pose yolov8s-pose.engine data/test.mp4 # the video path 129 | ``` 130 | -------------------------------------------------------------------------------- /docs/star.md: -------------------------------------------------------------------------------- 1 | ## Star History 2 | 3 | [![YOLOv8-TensorRT](https://api.star-history.com/svg?repos=triple-Mu/YOLOv8-TensorRT&type=Date)](https://star-history.com/#triple-Mu/YOLOv8-TensorRT&Date) 4 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | numpy<=1.23.5 2 | onnx==1.14.1 3 | onnxsim==1.16.0 4 | opencv-python 5 | # torch 6 | # torchvision 7 | ultralytics 8 | # tensorrt 9 | # cuda-python 10 | # pycuda 11 | onnxruntime==1.16.0 12 | -------------------------------------------------------------------------------- /scripts/.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | repos: 2 | - repo: https://github.com/PyCQA/flake8 3 | rev: 5.0.4 4 | hooks: 5 | - id: flake8 6 | - repo: https://github.com/PyCQA/isort 7 | rev: 5.11.5 8 | hooks: 9 | - id: isort 10 | - repo: https://github.com/pre-commit/mirrors-yapf 11 | rev: v0.32.0 12 | hooks: 13 | - id: yapf 14 | - repo: https://github.com/pre-commit/pre-commit-hooks 15 | rev: v4.3.0 16 | hooks: 17 | - id: trailing-whitespace 18 | - id: end-of-file-fixer 19 | - id: requirements-txt-fixer 20 | - id: double-quote-string-fixer 21 | - id: check-merge-conflict 22 | - id: fix-encoding-pragma 23 | args: ["--remove"] 24 | - id: mixed-line-ending 25 | args: ["--fix=lf"] 26 | -------------------------------------------------------------------------------- /scripts/build.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | from models import EngineBuilder 4 | 5 | 6 | def parse_args(): 7 | parser = argparse.ArgumentParser() 8 | parser.add_argument('--weights', 9 | type=str, 10 | required=True, 11 | help='Weights file') 12 | parser.add_argument('--iou-thres', 13 | type=float, 14 | default=0.65, 15 | help='IOU threshoud for NMS plugin') 16 | parser.add_argument('--conf-thres', 17 | type=float, 18 | default=0.25, 19 | help='CONF threshoud for NMS plugin') 20 | parser.add_argument('--topk', 21 | type=int, 22 | default=100, 23 | help='Max number of detection bboxes') 24 | parser.add_argument('--input-shape', 25 | nargs='+', 26 | type=int, 27 | default=[1, 3, 640, 640], 28 | help='Model input shape only for api builder') 29 | parser.add_argument('--fp16', 30 | action='store_true', 31 | help='Build model with fp16 mode') 32 | parser.add_argument('--device', 33 | type=str, 34 | default='cuda:0', 35 | help='TensorRT builder device') 36 | 37 | parser.add_argument('--seg', 38 | action='store_true', 39 | help='Build seg model by onnx') 40 | args = parser.parse_args() 41 | assert len(args.input_shape) == 4 42 | return args 43 | 44 | 45 | def main(args): 46 | builder = EngineBuilder(args.weights, args.device) 47 | builder.seg = args.seg 48 | builder.build(fp16=args.fp16, 49 | input_shape=args.input_shape, 50 | iou_thres=args.iou_thres, 51 | conf_thres=args.conf_thres, 52 | topk=args.topk) 53 | 54 | 55 | if __name__ == '__main__': 56 | args = parse_args() 57 | main(args) 58 | -------------------------------------------------------------------------------- /scripts/config.py: -------------------------------------------------------------------------------- 1 | import random 2 | 3 | import numpy as np 4 | 5 | random.seed(0) 6 | 7 | # detection model classes 8 | CLASSES = ('person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 9 | 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 10 | 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 11 | 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 12 | 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 13 | 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 14 | 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 15 | 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 16 | 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 17 | 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 18 | 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 19 | 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 20 | 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 21 | 'scissors', 'teddy bear', 'hair drier', 'toothbrush') 22 | 23 | # colors for per classes 24 | COLORS = { 25 | cls: [random.randint(0, 255) for _ in range(3)] 26 | for i, cls in enumerate(CLASSES) 27 | } 28 | 29 | # colors for segment masks 30 | MASK_COLORS = np.array([(255, 56, 56), (255, 157, 151), (255, 112, 31), 31 | (255, 178, 29), (207, 210, 49), (72, 249, 10), 32 | (146, 204, 23), (61, 219, 134), (26, 147, 52), 33 | (0, 212, 187), (44, 153, 168), (0, 194, 255), 34 | (52, 69, 147), (100, 115, 255), (0, 24, 236), 35 | (132, 56, 255), (82, 0, 133), (203, 56, 255), 36 | (255, 149, 200), (255, 55, 199)], 37 | dtype=np.float32) / 255. 38 | 39 | KPS_COLORS = [[0, 255, 0], [0, 255, 0], [0, 255, 0], [0, 255, 0], [0, 255, 0], 40 | [255, 128, 0], [255, 128, 0], [255, 128, 0], [255, 128, 0], 41 | [255, 128, 0], [255, 128, 0], [51, 153, 255], [51, 153, 255], 42 | [51, 153, 255], [51, 153, 255], [51, 153, 255], [51, 153, 255]] 43 | 44 | SKELETON = [[16, 14], [14, 12], [17, 15], [15, 13], [12, 13], [6, 12], [7, 13], 45 | [6, 7], [6, 8], [7, 9], [8, 10], [9, 11], [2, 3], [1, 2], [1, 3], 46 | [2, 4], [3, 5], [4, 6], [5, 7]] 47 | 48 | LIMB_COLORS = [[51, 153, 255], [51, 153, 255], [51, 153, 255], [51, 153, 255], 49 | [255, 51, 255], [255, 51, 255], [255, 51, 255], [255, 128, 0], 50 | [255, 128, 0], [255, 128, 0], [255, 128, 0], [255, 128, 0], 51 | [0, 255, 0], [0, 255, 0], [0, 255, 0], [0, 255, 0], [0, 255, 0], 52 | [0, 255, 0], [0, 255, 0]] 53 | 54 | # alpha for segment masks 55 | ALPHA = 0.5 56 | -------------------------------------------------------------------------------- /scripts/export-det.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from io import BytesIO 3 | 4 | import onnx 5 | import torch 6 | from ultralytics import YOLO 7 | 8 | from models.common import PostDetect, optim 9 | 10 | try: 11 | import onnxsim 12 | except ImportError: 13 | onnxsim = None 14 | 15 | 16 | def parse_args(): 17 | parser = argparse.ArgumentParser() 18 | parser.add_argument('-w', 19 | '--weights', 20 | type=str, 21 | required=True, 22 | help='PyTorch yolov8 weights') 23 | parser.add_argument('--iou-thres', 24 | type=float, 25 | default=0.65, 26 | help='IOU threshoud for NMS plugin') 27 | parser.add_argument('--conf-thres', 28 | type=float, 29 | default=0.25, 30 | help='CONF threshoud for NMS plugin') 31 | parser.add_argument('--topk', 32 | type=int, 33 | default=100, 34 | help='Max number of detection bboxes') 35 | parser.add_argument('--opset', 36 | type=int, 37 | default=11, 38 | help='ONNX opset version') 39 | parser.add_argument('--sim', 40 | action='store_true', 41 | help='simplify onnx model') 42 | parser.add_argument('--input-shape', 43 | nargs='+', 44 | type=int, 45 | default=[1, 3, 640, 640], 46 | help='Model input shape only for api builder') 47 | parser.add_argument('--device', 48 | type=str, 49 | default='cpu', 50 | help='Export ONNX device') 51 | args = parser.parse_args() 52 | assert len(args.input_shape) == 4 53 | PostDetect.conf_thres = args.conf_thres 54 | PostDetect.iou_thres = args.iou_thres 55 | PostDetect.topk = args.topk 56 | return args 57 | 58 | 59 | def main(args): 60 | b = args.input_shape[0] 61 | YOLOv8 = YOLO(args.weights) 62 | model = YOLOv8.model.fuse().eval() 63 | for m in model.modules(): 64 | optim(m) 65 | m.to(args.device) 66 | model.to(args.device) 67 | fake_input = torch.randn(args.input_shape).to(args.device) 68 | for _ in range(2): 69 | model(fake_input) 70 | save_path = args.weights.replace('.pt', '.onnx') 71 | with BytesIO() as f: 72 | torch.onnx.export( 73 | model, 74 | fake_input, 75 | f, 76 | opset_version=args.opset, 77 | input_names=['images'], 78 | output_names=['num_dets', 'bboxes', 'scores', 'labels']) 79 | f.seek(0) 80 | onnx_model = onnx.load(f) 81 | onnx.checker.check_model(onnx_model) 82 | shapes = [b, 1, b, args.topk, 4, b, args.topk, b, args.topk] 83 | for i in onnx_model.graph.output: 84 | for j in i.type.tensor_type.shape.dim: 85 | j.dim_param = str(shapes.pop(0)) 86 | if args.sim: 87 | try: 88 | onnx_model, check = onnxsim.simplify(onnx_model) 89 | assert check, 'assert check failed' 90 | except Exception as e: 91 | print(f'Simplifier failure: {e}') 92 | onnx.save(onnx_model, save_path) 93 | print(f'ONNX export success, saved as {save_path}') 94 | 95 | 96 | if __name__ == '__main__': 97 | main(parse_args()) 98 | -------------------------------------------------------------------------------- /scripts/export-pose.py: -------------------------------------------------------------------------------- 1 | from ultralytics import YOLO 2 | 3 | # Load a model 4 | model = YOLO("../weights/yolov8n-pose.pt") # load a pretrained model (recommended for training) 5 | success = model.export(format="onnx", opset=11, simplify=True) # export the model to onnx format 6 | assert success -------------------------------------------------------------------------------- /scripts/export-seg.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from io import BytesIO 3 | 4 | import onnx 5 | import torch 6 | from ultralytics import YOLO 7 | 8 | from models.common import optim 9 | 10 | try: 11 | import onnxsim 12 | except ImportError: 13 | onnxsim = None 14 | 15 | 16 | def parse_args(): 17 | parser = argparse.ArgumentParser() 18 | parser.add_argument('-w', 19 | '--weights', 20 | type=str, 21 | required=True, 22 | help='PyTorch yolov8 weights') 23 | parser.add_argument('--opset', 24 | type=int, 25 | default=11, 26 | help='ONNX opset version') 27 | parser.add_argument('--sim', 28 | action='store_true', 29 | help='simplify onnx model') 30 | parser.add_argument('--input-shape', 31 | nargs='+', 32 | type=int, 33 | default=[1, 3, 640, 640], 34 | help='Model input shape only for api builder') 35 | parser.add_argument('--device', 36 | type=str, 37 | default='cpu', 38 | help='Export ONNX device') 39 | 40 | args = parser.parse_args() 41 | assert len(args.input_shape) == 4 42 | return args 43 | 44 | 45 | def main(args): 46 | YOLOv8 = YOLO(args.weights) 47 | model = YOLOv8.model.fuse().eval() 48 | for m in model.modules(): 49 | optim(m) 50 | m.to(args.device) 51 | model.to(args.device) 52 | fake_input = torch.randn(args.input_shape).to(args.device) 53 | for _ in range(2): 54 | model(fake_input) 55 | save_path = args.weights.replace('.pt', '.onnx') 56 | with BytesIO() as f: 57 | torch.onnx.export(model, 58 | fake_input, 59 | f, 60 | opset_version=args.opset, 61 | input_names=['images'], 62 | output_names=['outputs', 'proto']) 63 | f.seek(0) 64 | onnx_model = onnx.load(f) 65 | onnx.checker.check_model(onnx_model) 66 | if args.sim: 67 | try: 68 | onnx_model, check = onnxsim.simplify(onnx_model) 69 | assert check, 'assert check failed' 70 | except Exception as e: 71 | print(f'Simplifier failure: {e}') 72 | onnx.save(onnx_model, save_path) 73 | print(f'ONNX export success, saved as {save_path}') 74 | 75 | 76 | if __name__ == '__main__': 77 | main(parse_args()) 78 | -------------------------------------------------------------------------------- /scripts/gen_pkl.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import pickle 3 | from collections import OrderedDict 4 | 5 | from ultralytics import YOLO 6 | 7 | 8 | def parse_args(): 9 | parser = argparse.ArgumentParser() 10 | parser.add_argument('-w', 11 | '--weights', 12 | type=str, 13 | required=True, 14 | help='YOLOv8 pytorch weights') 15 | parser.add_argument('-o', 16 | '--output', 17 | type=str, 18 | required=True, 19 | help='Output file') 20 | args = parser.parse_args() 21 | return args 22 | 23 | 24 | args = parse_args() 25 | 26 | model = YOLO(args.weights) 27 | model.model.fuse() 28 | YOLOv8 = model.model.model 29 | 30 | strides = YOLOv8[-1].stride.detach().cpu().numpy() 31 | reg_max = YOLOv8[-1].dfl.conv.weight.shape[1] 32 | 33 | state_dict = OrderedDict(GD=model.model.yaml['depth_multiple'], 34 | GW=model.model.yaml['width_multiple'], 35 | strides=strides, 36 | reg_max=reg_max) 37 | 38 | for name, value in YOLOv8.state_dict().items(): 39 | value = value.detach().cpu().numpy() 40 | i = int(name.split('.')[0]) 41 | layer = YOLOv8[i] 42 | module_name = layer.type.split('.')[-1] 43 | stem = module_name + '.' + name 44 | state_dict[stem] = value 45 | 46 | with open(args.output, 'wb') as f: 47 | pickle.dump(state_dict, f) 48 | -------------------------------------------------------------------------------- /scripts/infer-det-without-torch.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from pathlib import Path 3 | 4 | import cv2 5 | import numpy as np 6 | 7 | from config import CLASSES, COLORS 8 | from models.utils import blob, det_postprocess, letterbox, path_to_list 9 | 10 | 11 | def main(args: argparse.Namespace) -> None: 12 | if args.method == 'cudart': 13 | from models.cudart_api import TRTEngine 14 | elif args.method == 'pycuda': 15 | from models.pycuda_api import TRTEngine 16 | else: 17 | raise NotImplementedError 18 | 19 | Engine = TRTEngine(args.engine) 20 | H, W = Engine.inp_info[0].shape[-2:] 21 | 22 | images = path_to_list(args.imgs) 23 | save_path = Path(args.out_dir) 24 | 25 | if not args.show and not save_path.exists(): 26 | save_path.mkdir(parents=True, exist_ok=True) 27 | 28 | for image in images: 29 | save_image = save_path / image.name 30 | bgr = cv2.imread(str(image)) 31 | draw = bgr.copy() 32 | bgr, ratio, dwdh = letterbox(bgr, (W, H)) 33 | rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) 34 | tensor = blob(rgb, return_seg=False) 35 | dwdh = np.array(dwdh * 2, dtype=np.float32) 36 | tensor = np.ascontiguousarray(tensor) 37 | # inference 38 | data = Engine(tensor) 39 | 40 | bboxes, scores, labels = det_postprocess(data) 41 | if bboxes.size == 0: 42 | # if no bounding box 43 | print(f'{image}: no object!') 44 | continue 45 | bboxes -= dwdh 46 | bboxes /= ratio 47 | 48 | for (bbox, score, label) in zip(bboxes, scores, labels): 49 | bbox = bbox.round().astype(np.int32).tolist() 50 | cls_id = int(label) 51 | cls = CLASSES[cls_id] 52 | color = COLORS[cls] 53 | cv2.rectangle(draw, bbox[:2], bbox[2:], color, 2) 54 | cv2.putText(draw, 55 | f'{cls}:{score:.3f}', (bbox[0], bbox[1] - 2), 56 | cv2.FONT_HERSHEY_SIMPLEX, 57 | 0.75, [225, 255, 255], 58 | thickness=2) 59 | if args.show: 60 | cv2.imshow('result', draw) 61 | cv2.waitKey(0) 62 | else: 63 | cv2.imwrite(str(save_image), draw) 64 | 65 | 66 | def parse_args(): 67 | parser = argparse.ArgumentParser() 68 | parser.add_argument('--engine', type=str, help='Engine file') 69 | parser.add_argument('--imgs', type=str, help='Images file') 70 | parser.add_argument('--show', 71 | action='store_true', 72 | help='Show the detection results') 73 | parser.add_argument('--out-dir', 74 | type=str, 75 | default='./output', 76 | help='Path to output file') 77 | parser.add_argument('--method', 78 | type=str, 79 | default='cudart', 80 | help='CUDART pipeline') 81 | args = parser.parse_args() 82 | return args 83 | 84 | 85 | if __name__ == '__main__': 86 | args = parse_args() 87 | main(args) 88 | -------------------------------------------------------------------------------- /scripts/infer-det.py: -------------------------------------------------------------------------------- 1 | from models import TRTModule # isort:skip 2 | import argparse 3 | from pathlib import Path 4 | 5 | import cv2 6 | import torch 7 | 8 | from config import CLASSES, COLORS 9 | from models.torch_utils import det_postprocess 10 | from models.utils import blob, letterbox, path_to_list 11 | 12 | 13 | def main(args: argparse.Namespace) -> None: 14 | device = torch.device(args.device) 15 | Engine = TRTModule(args.engine, device) 16 | H, W = Engine.inp_info[0].shape[-2:] 17 | 18 | # set desired output names order 19 | Engine.set_desired(['num_dets', 'bboxes', 'scores', 'labels']) 20 | 21 | images = path_to_list(args.imgs) 22 | save_path = Path(args.out_dir) 23 | 24 | if not args.show and not save_path.exists(): 25 | save_path.mkdir(parents=True, exist_ok=True) 26 | 27 | for image in images: 28 | save_image = save_path / image.name 29 | bgr = cv2.imread(str(image)) 30 | draw = bgr.copy() 31 | bgr, ratio, dwdh = letterbox(bgr, (W, H)) 32 | rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) 33 | tensor = blob(rgb, return_seg=False) 34 | dwdh = torch.asarray(dwdh * 2, dtype=torch.float32, device=device) 35 | tensor = torch.asarray(tensor, device=device) 36 | # inference 37 | data = Engine(tensor) 38 | 39 | bboxes, scores, labels = det_postprocess(data) 40 | if bboxes.numel() == 0: 41 | # if no bounding box 42 | print(f'{image}: no object!') 43 | continue 44 | bboxes -= dwdh 45 | bboxes /= ratio 46 | 47 | for (bbox, score, label) in zip(bboxes, scores, labels): 48 | bbox = bbox.round().int().tolist() 49 | cls_id = int(label) 50 | cls = CLASSES[cls_id] 51 | color = COLORS[cls] 52 | cv2.rectangle(draw, bbox[:2], bbox[2:], color, 2) 53 | cv2.putText(draw, 54 | f'{cls}:{score:.3f}', (bbox[0], bbox[1] - 2), 55 | cv2.FONT_HERSHEY_SIMPLEX, 56 | 0.75, [225, 255, 255], 57 | thickness=2) 58 | if args.show: 59 | cv2.imshow('result', draw) 60 | cv2.waitKey(0) 61 | else: 62 | cv2.imwrite(str(save_image), draw) 63 | 64 | 65 | def parse_args() -> argparse.Namespace: 66 | parser = argparse.ArgumentParser() 67 | parser.add_argument('--engine', type=str, help='Engine file') 68 | parser.add_argument('--imgs', type=str, help='Images file') 69 | parser.add_argument('--show', 70 | action='store_true', 71 | help='Show the detection results') 72 | parser.add_argument('--out-dir', 73 | type=str, 74 | default='./output', 75 | help='Path to output file') 76 | parser.add_argument('--device', 77 | type=str, 78 | default='cuda:0', 79 | help='TensorRT infer device') 80 | args = parser.parse_args() 81 | return args 82 | 83 | 84 | if __name__ == '__main__': 85 | args = parse_args() 86 | main(args) 87 | -------------------------------------------------------------------------------- /scripts/infer-pose-without-torch.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from pathlib import Path 3 | 4 | import cv2 5 | import numpy as np 6 | 7 | from config import COLORS, KPS_COLORS, LIMB_COLORS, SKELETON 8 | from models.utils import blob, letterbox, path_to_list, pose_postprocess 9 | 10 | 11 | def main(args: argparse.Namespace) -> None: 12 | if args.method == 'cudart': 13 | from models.cudart_api import TRTEngine 14 | elif args.method == 'pycuda': 15 | from models.pycuda_api import TRTEngine 16 | else: 17 | raise NotImplementedError 18 | 19 | Engine = TRTEngine(args.engine) 20 | H, W = Engine.inp_info[0].shape[-2:] 21 | 22 | images = path_to_list(args.imgs) 23 | save_path = Path(args.out_dir) 24 | 25 | if not args.show and not save_path.exists(): 26 | save_path.mkdir(parents=True, exist_ok=True) 27 | 28 | for image in images: 29 | save_image = save_path / image.name 30 | bgr = cv2.imread(str(image)) 31 | draw = bgr.copy() 32 | bgr, ratio, dwdh = letterbox(bgr, (W, H)) 33 | dw, dh = int(dwdh[0]), int(dwdh[1]) 34 | rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) 35 | tensor = blob(rgb, return_seg=False) 36 | dwdh = np.array(dwdh * 2, dtype=np.float32) 37 | tensor = np.ascontiguousarray(tensor) 38 | # inference 39 | data = Engine(tensor) 40 | 41 | bboxes, scores, kpts = pose_postprocess(data, args.conf_thres, 42 | args.iou_thres) 43 | if bboxes.size == 0: 44 | # if no bounding box 45 | print(f'{image}: no object!') 46 | continue 47 | bboxes -= dwdh 48 | bboxes /= ratio 49 | 50 | for (bbox, score, kpt) in zip(bboxes, scores, kpts): 51 | bbox = bbox.round().astype(np.int32).tolist() 52 | color = COLORS['person'] 53 | cv2.rectangle(draw, bbox[:2], bbox[2:], color, 2) 54 | cv2.putText(draw, 55 | f'person:{score:.3f}', (bbox[0], bbox[1] - 2), 56 | cv2.FONT_HERSHEY_SIMPLEX, 57 | 0.75, [225, 255, 255], 58 | thickness=2) 59 | for i in range(19): 60 | if i < 17: 61 | px, py, ps = kpt[i] 62 | if ps > 0.5: 63 | kcolor = KPS_COLORS[i] 64 | px = round(float(px - dw) / ratio) 65 | py = round(float(py - dh) / ratio) 66 | cv2.circle(draw, (px, py), 5, kcolor, -1) 67 | xi, yi = SKELETON[i] 68 | pos1_s = kpt[xi - 1][2] 69 | pos2_s = kpt[yi - 1][2] 70 | if pos1_s > 0.5 and pos2_s > 0.5: 71 | limb_color = LIMB_COLORS[i] 72 | pos1_x = round(float(kpt[xi - 1][0] - dw) / ratio) 73 | pos1_y = round(float(kpt[xi - 1][1] - dh) / ratio) 74 | 75 | pos2_x = round(float(kpt[yi - 1][0] - dw) / ratio) 76 | pos2_y = round(float(kpt[yi - 1][1] - dh) / ratio) 77 | 78 | cv2.line(draw, (pos1_x, pos1_y), (pos2_x, pos2_y), 79 | limb_color, 2) 80 | if args.show: 81 | cv2.imshow('result', draw) 82 | cv2.waitKey(0) 83 | else: 84 | cv2.imwrite(str(save_image), draw) 85 | 86 | 87 | def parse_args(): 88 | parser = argparse.ArgumentParser() 89 | parser.add_argument('--engine', type=str, help='Engine file') 90 | parser.add_argument('--imgs', type=str, help='Images file') 91 | parser.add_argument('--show', 92 | action='store_true', 93 | help='Show the detection results') 94 | parser.add_argument('--out-dir', 95 | type=str, 96 | default='./output', 97 | help='Path to output file') 98 | parser.add_argument('--conf-thres', 99 | type=float, 100 | default=0.25, 101 | help='Confidence threshold') 102 | parser.add_argument('--iou-thres', 103 | type=float, 104 | default=0.65, 105 | help='Confidence threshold') 106 | parser.add_argument('--method', 107 | type=str, 108 | default='cudart', 109 | help='CUDART pipeline') 110 | args = parser.parse_args() 111 | return args 112 | 113 | 114 | if __name__ == '__main__': 115 | args = parse_args() 116 | main(args) 117 | -------------------------------------------------------------------------------- /scripts/infer-pose.py: -------------------------------------------------------------------------------- 1 | from models import TRTModule # isort:skip 2 | import argparse 3 | from pathlib import Path 4 | 5 | import cv2 6 | import torch 7 | 8 | from config import COLORS, KPS_COLORS, LIMB_COLORS, SKELETON 9 | from models.torch_utils import pose_postprocess 10 | from models.utils import blob, letterbox, path_to_list 11 | 12 | 13 | def main(args: argparse.Namespace) -> None: 14 | device = torch.device(args.device) 15 | Engine = TRTModule(args.engine, device) 16 | H, W = Engine.inp_info[0].shape[-2:] 17 | 18 | images = path_to_list(args.imgs) 19 | save_path = Path(args.out_dir) 20 | 21 | if not args.show and not save_path.exists(): 22 | save_path.mkdir(parents=True, exist_ok=True) 23 | 24 | for image in images: 25 | save_image = save_path / image.name 26 | bgr = cv2.imread(str(image)) 27 | draw = bgr.copy() 28 | bgr, ratio, dwdh = letterbox(bgr, (W, H)) 29 | dw, dh = int(dwdh[0]), int(dwdh[1]) 30 | rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) 31 | tensor = blob(rgb, return_seg=False) 32 | dwdh = torch.asarray(dwdh * 2, dtype=torch.float32, device=device) 33 | tensor = torch.asarray(tensor, device=device) 34 | # inference 35 | data = Engine(tensor) 36 | 37 | bboxes, scores, kpts = pose_postprocess(data, args.conf_thres, 38 | args.iou_thres) 39 | if bboxes.numel() == 0: 40 | # if no bounding box 41 | print(f'{image}: no object!') 42 | continue 43 | bboxes -= dwdh 44 | bboxes /= ratio 45 | 46 | for (bbox, score, kpt) in zip(bboxes, scores, kpts): 47 | bbox = bbox.round().int().tolist() 48 | color = COLORS['person'] 49 | cv2.rectangle(draw, bbox[:2], bbox[2:], color, 2) 50 | cv2.putText(draw, 51 | f'person:{score:.3f}', (bbox[0], bbox[1] - 2), 52 | cv2.FONT_HERSHEY_SIMPLEX, 53 | 0.75, [225, 255, 255], 54 | thickness=2) 55 | for i in range(19): 56 | if i < 17: 57 | px, py, ps = kpt[i] 58 | if ps > 0.5: 59 | kcolor = KPS_COLORS[i] 60 | px = round(float(px - dw) / ratio) 61 | py = round(float(py - dh) / ratio) 62 | cv2.circle(draw, (px, py), 5, kcolor, -1) 63 | xi, yi = SKELETON[i] 64 | pos1_s = kpt[xi - 1][2] 65 | pos2_s = kpt[yi - 1][2] 66 | if pos1_s > 0.5 and pos2_s > 0.5: 67 | limb_color = LIMB_COLORS[i] 68 | pos1_x = round(float(kpt[xi - 1][0] - dw) / ratio) 69 | pos1_y = round(float(kpt[xi - 1][1] - dh) / ratio) 70 | 71 | pos2_x = round(float(kpt[yi - 1][0] - dw) / ratio) 72 | pos2_y = round(float(kpt[yi - 1][1] - dh) / ratio) 73 | 74 | cv2.line(draw, (pos1_x, pos1_y), (pos2_x, pos2_y), 75 | limb_color, 2) 76 | if args.show: 77 | cv2.imshow('result', draw) 78 | cv2.waitKey(0) 79 | else: 80 | cv2.imwrite(str(save_image), draw) 81 | 82 | 83 | def parse_args() -> argparse.Namespace: 84 | parser = argparse.ArgumentParser() 85 | parser.add_argument('--engine', type=str, help='Engine file') 86 | parser.add_argument('--imgs', type=str, help='Images file') 87 | parser.add_argument('--show', 88 | action='store_true', 89 | help='Show the detection results') 90 | parser.add_argument('--out-dir', 91 | type=str, 92 | default='./output', 93 | help='Path to output file') 94 | parser.add_argument('--conf-thres', 95 | type=float, 96 | default=0.25, 97 | help='Confidence threshold') 98 | parser.add_argument('--iou-thres', 99 | type=float, 100 | default=0.65, 101 | help='Confidence threshold') 102 | parser.add_argument('--device', 103 | type=str, 104 | default='cuda:0', 105 | help='TensorRT infer device') 106 | args = parser.parse_args() 107 | return args 108 | 109 | 110 | if __name__ == '__main__': 111 | args = parse_args() 112 | main(args) 113 | -------------------------------------------------------------------------------- /scripts/infer-seg-without-torch.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from pathlib import Path 3 | 4 | import cv2 5 | import numpy as np 6 | 7 | from config import ALPHA, CLASSES, COLORS, MASK_COLORS 8 | from models.utils import blob, letterbox, path_to_list, seg_postprocess 9 | 10 | 11 | def main(args: argparse.Namespace) -> None: 12 | if args.method == 'cudart': 13 | from models.cudart_api import TRTEngine 14 | elif args.method == 'pycuda': 15 | from models.pycuda_api import TRTEngine 16 | else: 17 | raise NotImplementedError 18 | 19 | Engine = TRTEngine(args.engine) 20 | H, W = Engine.inp_info[0].shape[-2:] 21 | 22 | images = path_to_list(args.imgs) 23 | save_path = Path(args.out_dir) 24 | 25 | if not args.show and not save_path.exists(): 26 | save_path.mkdir(parents=True, exist_ok=True) 27 | 28 | for image in images: 29 | save_image = save_path / image.name 30 | bgr = cv2.imread(str(image)) 31 | draw = bgr.copy() 32 | bgr, ratio, dwdh = letterbox(bgr, (W, H)) 33 | dw, dh = int(dwdh[0]), int(dwdh[1]) 34 | rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) 35 | tensor, seg_img = blob(rgb, return_seg=True) 36 | dwdh = np.array(dwdh * 2, dtype=np.float32) 37 | tensor = np.ascontiguousarray(tensor) 38 | # inference 39 | data = Engine(tensor) 40 | 41 | seg_img = seg_img[dh:H - dh, dw:W - dw, [2, 1, 0]] 42 | bboxes, scores, labels, masks = seg_postprocess( 43 | data, bgr.shape[:2], args.conf_thres, args.iou_thres) 44 | if bboxes.size == 0: 45 | # if no bounding box 46 | print(f'{image}: no object!') 47 | continue 48 | masks = masks[:, dh:H - dh, dw:W - dw, :] 49 | mask_colors = MASK_COLORS[labels % len(MASK_COLORS)] 50 | mask_colors = mask_colors.reshape(-1, 1, 1, 3) * ALPHA 51 | mask_colors = masks @ mask_colors 52 | inv_alph_masks = (1 - masks * 0.5).cumprod(0) 53 | mcs = (mask_colors * inv_alph_masks).sum(0) * 2 54 | seg_img = (seg_img * inv_alph_masks[-1] + mcs) * 255 55 | draw = cv2.resize(seg_img.astype(np.uint8), draw.shape[:2][::-1]) 56 | 57 | bboxes -= dwdh 58 | bboxes /= ratio 59 | 60 | for (bbox, score, label) in zip(bboxes, scores, labels): 61 | bbox = bbox.round().astype(np.int32).tolist() 62 | cls_id = int(label) 63 | cls = CLASSES[cls_id] 64 | color = COLORS[cls] 65 | cv2.rectangle(draw, bbox[:2], bbox[2:], color, 2) 66 | cv2.putText(draw, 67 | f'{cls}:{score:.3f}', (bbox[0], bbox[1] - 2), 68 | cv2.FONT_HERSHEY_SIMPLEX, 69 | 0.75, [225, 255, 255], 70 | thickness=2) 71 | if args.show: 72 | cv2.imshow('result', draw) 73 | cv2.waitKey(0) 74 | else: 75 | cv2.imwrite(str(save_image), draw) 76 | 77 | 78 | def parse_args(): 79 | parser = argparse.ArgumentParser() 80 | parser.add_argument('--engine', type=str, help='Engine file') 81 | parser.add_argument('--imgs', type=str, help='Images file') 82 | parser.add_argument('--show', 83 | action='store_true', 84 | help='Show the detection results') 85 | parser.add_argument('--out-dir', 86 | type=str, 87 | default='./output', 88 | help='Path to output file') 89 | parser.add_argument('--conf-thres', 90 | type=float, 91 | default=0.25, 92 | help='Confidence threshold') 93 | parser.add_argument('--iou-thres', 94 | type=float, 95 | default=0.65, 96 | help='Confidence threshold') 97 | parser.add_argument('--method', 98 | type=str, 99 | default='cudart', 100 | help='CUDART pipeline') 101 | args = parser.parse_args() 102 | return args 103 | 104 | 105 | if __name__ == '__main__': 106 | args = parse_args() 107 | main(args) 108 | -------------------------------------------------------------------------------- /scripts/infer-seg.py: -------------------------------------------------------------------------------- 1 | from models import TRTModule # isort:skip 2 | import argparse 3 | from pathlib import Path 4 | 5 | import cv2 6 | import numpy as np 7 | import torch 8 | 9 | from config import ALPHA, CLASSES, COLORS, MASK_COLORS 10 | from models.torch_utils import seg_postprocess 11 | from models.utils import blob, letterbox, path_to_list 12 | 13 | 14 | def main(args: argparse.Namespace) -> None: 15 | device = torch.device(args.device) 16 | Engine = TRTModule(args.engine, device) 17 | H, W = Engine.inp_info[0].shape[-2:] 18 | 19 | # set desired output names order 20 | Engine.set_desired(['outputs', 'proto']) 21 | 22 | images = path_to_list(args.imgs) 23 | save_path = Path(args.out_dir) 24 | 25 | if not args.show and not save_path.exists(): 26 | save_path.mkdir(parents=True, exist_ok=True) 27 | 28 | for image in images: 29 | save_image = save_path / image.name 30 | bgr = cv2.imread(str(image)) 31 | draw = bgr.copy() 32 | bgr, ratio, dwdh = letterbox(bgr, (W, H)) 33 | dw, dh = int(dwdh[0]), int(dwdh[1]) 34 | rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) 35 | tensor, seg_img = blob(rgb, return_seg=True) 36 | dwdh = torch.asarray(dwdh * 2, dtype=torch.float32, device=device) 37 | tensor = torch.asarray(tensor, device=device) 38 | # inference 39 | data = Engine(tensor) 40 | 41 | seg_img = torch.asarray(seg_img[dh:H - dh, dw:W - dw, [2, 1, 0]], 42 | device=device) 43 | bboxes, scores, labels, masks = seg_postprocess( 44 | data, bgr.shape[:2], args.conf_thres, args.iou_thres) 45 | if bboxes.numel() == 0: 46 | # if no bounding box 47 | print(f'{image}: no object!') 48 | continue 49 | masks = masks[:, dh:H - dh, dw:W - dw, :] 50 | indices = (labels % len(MASK_COLORS)).long() 51 | mask_colors = torch.asarray(MASK_COLORS, device=device)[indices] 52 | mask_colors = mask_colors.view(-1, 1, 1, 3) * ALPHA 53 | mask_colors = masks @ mask_colors 54 | inv_alph_masks = (1 - masks * 0.5).cumprod(0) 55 | mcs = (mask_colors * inv_alph_masks).sum(0) * 2 56 | seg_img = (seg_img * inv_alph_masks[-1] + mcs) * 255 57 | draw = cv2.resize(seg_img.cpu().numpy().astype(np.uint8), 58 | draw.shape[:2][::-1]) 59 | 60 | bboxes -= dwdh 61 | bboxes /= ratio 62 | 63 | for (bbox, score, label) in zip(bboxes, scores, labels): 64 | bbox = bbox.round().int().tolist() 65 | cls_id = int(label) 66 | cls = CLASSES[cls_id] 67 | color = COLORS[cls] 68 | cv2.rectangle(draw, bbox[:2], bbox[2:], color, 2) 69 | cv2.putText(draw, 70 | f'{cls}:{score:.3f}', (bbox[0], bbox[1] - 2), 71 | cv2.FONT_HERSHEY_SIMPLEX, 72 | 0.75, [225, 255, 255], 73 | thickness=2) 74 | if args.show: 75 | cv2.imshow('result', draw) 76 | cv2.waitKey(0) 77 | else: 78 | cv2.imwrite(str(save_image), draw) 79 | 80 | 81 | def parse_args() -> argparse.Namespace: 82 | parser = argparse.ArgumentParser() 83 | parser.add_argument('--engine', type=str, help='Engine file') 84 | parser.add_argument('--imgs', type=str, help='Images file') 85 | parser.add_argument('--show', 86 | action='store_true', 87 | help='Show the detection results') 88 | parser.add_argument('--out-dir', 89 | type=str, 90 | default='./output', 91 | help='Path to output file') 92 | parser.add_argument('--conf-thres', 93 | type=float, 94 | default=0.25, 95 | help='Confidence threshold') 96 | parser.add_argument('--iou-thres', 97 | type=float, 98 | default=0.65, 99 | help='Confidence threshold') 100 | parser.add_argument('--device', 101 | type=str, 102 | default='cuda:0', 103 | help='TensorRT infer device') 104 | args = parser.parse_args() 105 | return args 106 | 107 | 108 | if __name__ == '__main__': 109 | args = parse_args() 110 | main(args) 111 | -------------------------------------------------------------------------------- /scripts/models/__init__.py: -------------------------------------------------------------------------------- 1 | from .engine import EngineBuilder, TRTModule, TRTProfilerV0, TRTProfilerV1 # isort:skip # noqa: E501 2 | import warnings 3 | 4 | import torch 5 | 6 | warnings.filterwarnings(action='ignore', category=torch.jit.TracerWarning) 7 | warnings.filterwarnings(action='ignore', category=torch.jit.ScriptWarning) 8 | warnings.filterwarnings(action='ignore', category=UserWarning) 9 | warnings.filterwarnings(action='ignore', category=FutureWarning) 10 | warnings.filterwarnings(action='ignore', category=DeprecationWarning) 11 | __all__ = ['EngineBuilder', 'TRTModule', 'TRTProfilerV0', 'TRTProfilerV1'] 12 | -------------------------------------------------------------------------------- /scripts/models/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linClubs/YOLOv8-ROS-TensorRT/0b23e9753b2cf0a2f26eda2335d867f1f2f69dda/scripts/models/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /scripts/models/__pycache__/common.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linClubs/YOLOv8-ROS-TensorRT/0b23e9753b2cf0a2f26eda2335d867f1f2f69dda/scripts/models/__pycache__/common.cpython-38.pyc -------------------------------------------------------------------------------- /scripts/models/__pycache__/engine.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linClubs/YOLOv8-ROS-TensorRT/0b23e9753b2cf0a2f26eda2335d867f1f2f69dda/scripts/models/__pycache__/engine.cpython-38.pyc -------------------------------------------------------------------------------- /scripts/models/common.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple 2 | 3 | import torch 4 | import torch.nn as nn 5 | from torch import Graph, Tensor, Value 6 | 7 | 8 | def make_anchors(feats: Tensor, 9 | strides: Tensor, 10 | grid_cell_offset: float = 0.5) -> Tuple[Tensor, Tensor]: 11 | anchor_points, stride_tensor = [], [] 12 | assert feats is not None 13 | dtype, device = feats[0].dtype, feats[0].device 14 | for i, stride in enumerate(strides): 15 | _, _, h, w = feats[i].shape 16 | sx = torch.arange(end=w, device=device, 17 | dtype=dtype) + grid_cell_offset # shift x 18 | sy = torch.arange(end=h, device=device, 19 | dtype=dtype) + grid_cell_offset # shift y 20 | sy, sx = torch.meshgrid(sy, sx) 21 | anchor_points.append(torch.stack((sx, sy), -1).view(-1, 2)) 22 | stride_tensor.append( 23 | torch.full((h * w, 1), stride, dtype=dtype, device=device)) 24 | return torch.cat(anchor_points), torch.cat(stride_tensor) 25 | 26 | 27 | class TRT_NMS(torch.autograd.Function): 28 | 29 | @staticmethod 30 | def forward( 31 | ctx: Graph, 32 | boxes: Tensor, 33 | scores: Tensor, 34 | iou_threshold: float = 0.65, 35 | score_threshold: float = 0.25, 36 | max_output_boxes: int = 100, 37 | background_class: int = -1, 38 | box_coding: int = 0, 39 | plugin_version: str = '1', 40 | score_activation: int = 0 41 | ) -> Tuple[Tensor, Tensor, Tensor, Tensor]: 42 | batch_size, num_boxes, num_classes = scores.shape 43 | num_dets = torch.randint(0, 44 | max_output_boxes, (batch_size, 1), 45 | dtype=torch.int32) 46 | boxes = torch.randn(batch_size, max_output_boxes, 4) 47 | scores = torch.randn(batch_size, max_output_boxes) 48 | labels = torch.randint(0, 49 | num_classes, (batch_size, max_output_boxes), 50 | dtype=torch.int32) 51 | 52 | return num_dets, boxes, scores, labels 53 | 54 | @staticmethod 55 | def symbolic( 56 | g, 57 | boxes: Value, 58 | scores: Value, 59 | iou_threshold: float = 0.45, 60 | score_threshold: float = 0.25, 61 | max_output_boxes: int = 100, 62 | background_class: int = -1, 63 | box_coding: int = 0, 64 | score_activation: int = 0, 65 | plugin_version: str = '1') -> Tuple[Value, Value, Value, Value]: 66 | out = g.op('TRT::EfficientNMS_TRT', 67 | boxes, 68 | scores, 69 | iou_threshold_f=iou_threshold, 70 | score_threshold_f=score_threshold, 71 | max_output_boxes_i=max_output_boxes, 72 | background_class_i=background_class, 73 | box_coding_i=box_coding, 74 | plugin_version_s=plugin_version, 75 | score_activation_i=score_activation, 76 | outputs=4) 77 | nums_dets, boxes, scores, classes = out 78 | return nums_dets, boxes, scores, classes 79 | 80 | 81 | class C2f(nn.Module): 82 | 83 | def __init__(self, *args, **kwargs): 84 | super().__init__() 85 | 86 | def forward(self, x): 87 | x = self.cv1(x) 88 | x = [x, x[:, self.c:, ...]] 89 | x.extend(m(x[-1]) for m in self.m) 90 | x.pop(1) 91 | return self.cv2(torch.cat(x, 1)) 92 | 93 | 94 | class PostDetect(nn.Module): 95 | export = True 96 | shape = None 97 | dynamic = False 98 | iou_thres = 0.65 99 | conf_thres = 0.25 100 | topk = 100 101 | 102 | def __init__(self, *args, **kwargs): 103 | super().__init__() 104 | 105 | def forward(self, x): 106 | shape = x[0].shape 107 | b, res, b_reg_num = shape[0], [], self.reg_max * 4 108 | for i in range(self.nl): 109 | res.append(torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)) 110 | if self.dynamic or self.shape != shape: 111 | self.anchors, self.strides = (x.transpose( 112 | 0, 1) for x in make_anchors(x, self.stride, 0.5)) 113 | self.shape = shape 114 | x = [i.view(b, self.no, -1) for i in res] 115 | y = torch.cat(x, 2) 116 | boxes, scores = y[:, :b_reg_num, ...], y[:, b_reg_num:, ...].sigmoid() 117 | boxes = boxes.view(b, 4, self.reg_max, -1).permute(0, 1, 3, 2) 118 | boxes = boxes.softmax(-1) @ torch.arange(self.reg_max).to(boxes) 119 | boxes0, boxes1 = -boxes[:, :2, ...], boxes[:, 2:, ...] 120 | boxes = self.anchors.repeat(b, 2, 1) + torch.cat([boxes0, boxes1], 1) 121 | boxes = boxes * self.strides 122 | 123 | return TRT_NMS.apply(boxes.transpose(1, 2), scores.transpose(1, 2), 124 | self.iou_thres, self.conf_thres, self.topk) 125 | 126 | 127 | class PostSeg(nn.Module): 128 | export = True 129 | shape = None 130 | dynamic = False 131 | 132 | def __init__(self, *args, **kwargs): 133 | super().__init__() 134 | 135 | def forward(self, x): 136 | p = self.proto(x[0]) # mask protos 137 | bs = p.shape[0] # batch size 138 | mc = torch.cat( 139 | [self.cv4[i](x[i]).view(bs, self.nm, -1) for i in range(self.nl)], 140 | 2) # mask coefficients 141 | boxes, scores, labels = self.forward_det(x) 142 | out = torch.cat([boxes, scores, labels.float(), mc.transpose(1, 2)], 2) 143 | return out, p.flatten(2) 144 | 145 | def forward_det(self, x): 146 | shape = x[0].shape 147 | b, res, b_reg_num = shape[0], [], self.reg_max * 4 148 | for i in range(self.nl): 149 | res.append(torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)) 150 | if self.dynamic or self.shape != shape: 151 | self.anchors, self.strides = \ 152 | (x.transpose(0, 1) for x in make_anchors(x, self.stride, 0.5)) 153 | self.shape = shape 154 | x = [i.view(b, self.no, -1) for i in res] 155 | y = torch.cat(x, 2) 156 | boxes, scores = y[:, :b_reg_num, ...], y[:, b_reg_num:, ...].sigmoid() 157 | boxes = boxes.view(b, 4, self.reg_max, -1).permute(0, 1, 3, 2) 158 | boxes = boxes.softmax(-1) @ torch.arange(self.reg_max).to(boxes) 159 | boxes0, boxes1 = -boxes[:, :2, ...], boxes[:, 2:, ...] 160 | boxes = self.anchors.repeat(b, 2, 1) + torch.cat([boxes0, boxes1], 1) 161 | boxes = boxes * self.strides 162 | scores, labels = scores.transpose(1, 2).max(dim=-1, keepdim=True) 163 | return boxes.transpose(1, 2), scores, labels 164 | 165 | 166 | def optim(module: nn.Module): 167 | s = str(type(module))[6:-2].split('.')[-1] 168 | if s == 'Detect': 169 | setattr(module, '__class__', PostDetect) 170 | elif s == 'Segment': 171 | setattr(module, '__class__', PostSeg) 172 | elif s == 'C2f': 173 | setattr(module, '__class__', C2f) 174 | -------------------------------------------------------------------------------- /scripts/models/cudart_api.py: -------------------------------------------------------------------------------- 1 | import os 2 | import warnings 3 | from dataclasses import dataclass 4 | from pathlib import Path 5 | from typing import List, Optional, Tuple, Union 6 | 7 | import numpy as np 8 | import tensorrt as trt 9 | from cuda import cudart 10 | from numpy import ndarray 11 | 12 | os.environ['CUDA_MODULE_LOADING'] = 'LAZY' 13 | warnings.filterwarnings(action='ignore', category=DeprecationWarning) 14 | 15 | 16 | @dataclass 17 | class Tensor: 18 | name: str 19 | dtype: np.dtype 20 | shape: Tuple 21 | cpu: ndarray 22 | gpu: int 23 | 24 | 25 | class TRTEngine: 26 | 27 | def __init__(self, weight: Union[str, Path]) -> None: 28 | self.weight = Path(weight) if isinstance(weight, str) else weight 29 | status, self.stream = cudart.cudaStreamCreate() 30 | assert status.value == 0 31 | self.__init_engine() 32 | self.__init_bindings() 33 | self.__warm_up() 34 | 35 | def __init_engine(self) -> None: 36 | logger = trt.Logger(trt.Logger.WARNING) 37 | trt.init_libnvinfer_plugins(logger, namespace='') 38 | with trt.Runtime(logger) as runtime: 39 | model = runtime.deserialize_cuda_engine(self.weight.read_bytes()) 40 | 41 | context = model.create_execution_context() 42 | 43 | names = [model.get_binding_name(i) for i in range(model.num_bindings)] 44 | self.num_bindings = model.num_bindings 45 | self.bindings: List[int] = [0] * self.num_bindings 46 | num_inputs, num_outputs = 0, 0 47 | 48 | for i in range(model.num_bindings): 49 | if model.binding_is_input(i): 50 | num_inputs += 1 51 | else: 52 | num_outputs += 1 53 | 54 | self.num_inputs = num_inputs 55 | self.num_outputs = num_outputs 56 | self.model = model 57 | self.context = context 58 | self.input_names = names[:num_inputs] 59 | self.output_names = names[num_inputs:] 60 | 61 | def __init_bindings(self) -> None: 62 | dynamic = False 63 | inp_info = [] 64 | out_info = [] 65 | out_ptrs = [] 66 | for i, name in enumerate(self.input_names): 67 | assert self.model.get_binding_name(i) == name 68 | dtype = trt.nptype(self.model.get_binding_dtype(i)) 69 | shape = tuple(self.model.get_binding_shape(i)) 70 | if -1 in shape: 71 | dynamic |= True 72 | if not dynamic: 73 | cpu = np.empty(shape, dtype) 74 | status, gpu = cudart.cudaMallocAsync(cpu.nbytes, self.stream) 75 | assert status.value == 0 76 | cudart.cudaMemcpyAsync( 77 | gpu, cpu.ctypes.data, cpu.nbytes, 78 | cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, self.stream) 79 | else: 80 | cpu, gpu = np.empty(0), 0 81 | inp_info.append(Tensor(name, dtype, shape, cpu, gpu)) 82 | for i, name in enumerate(self.output_names): 83 | i += self.num_inputs 84 | assert self.model.get_binding_name(i) == name 85 | dtype = trt.nptype(self.model.get_binding_dtype(i)) 86 | shape = tuple(self.model.get_binding_shape(i)) 87 | if not dynamic: 88 | cpu = np.empty(shape, dtype=dtype) 89 | status, gpu = cudart.cudaMallocAsync(cpu.nbytes, self.stream) 90 | assert status.value == 0 91 | cudart.cudaMemcpyAsync( 92 | gpu, cpu.ctypes.data, cpu.nbytes, 93 | cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, self.stream) 94 | out_ptrs.append(gpu) 95 | else: 96 | cpu, gpu = np.empty(0), 0 97 | out_info.append(Tensor(name, dtype, shape, cpu, gpu)) 98 | 99 | self.is_dynamic = dynamic 100 | self.inp_info = inp_info 101 | self.out_info = out_info 102 | self.out_ptrs = out_ptrs 103 | 104 | def __warm_up(self) -> None: 105 | if self.is_dynamic: 106 | print('You engine has dynamic axes, please warm up by yourself !') 107 | return 108 | for _ in range(10): 109 | inputs = [] 110 | for i in self.inp_info: 111 | inputs.append(i.cpu) 112 | self.__call__(inputs) 113 | 114 | def set_profiler(self, profiler: Optional[trt.IProfiler]) -> None: 115 | self.context.profiler = profiler \ 116 | if profiler is not None else trt.Profiler() 117 | 118 | def __call__(self, *inputs) -> Union[Tuple, ndarray]: 119 | 120 | assert len(inputs) == self.num_inputs 121 | contiguous_inputs: List[ndarray] = [ 122 | np.ascontiguousarray(i) for i in inputs 123 | ] 124 | 125 | for i in range(self.num_inputs): 126 | 127 | if self.is_dynamic: 128 | self.context.set_binding_shape( 129 | i, tuple(contiguous_inputs[i].shape)) 130 | status, self.inp_info[i].gpu = cudart.cudaMallocAsync( 131 | contiguous_inputs[i].nbytes, self.stream) 132 | assert status.value == 0 133 | cudart.cudaMemcpyAsync( 134 | self.inp_info[i].gpu, contiguous_inputs[i].ctypes.data, 135 | contiguous_inputs[i].nbytes, 136 | cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, self.stream) 137 | self.bindings[i] = self.inp_info[i].gpu 138 | 139 | output_gpu_ptrs: List[int] = [] 140 | outputs: List[ndarray] = [] 141 | 142 | for i in range(self.num_outputs): 143 | j = i + self.num_inputs 144 | if self.is_dynamic: 145 | shape = tuple(self.context.get_binding_shape(j)) 146 | dtype = self.out_info[i].dtype 147 | cpu = np.empty(shape, dtype=dtype) 148 | status, gpu = cudart.cudaMallocAsync(cpu.nbytes, self.stream) 149 | assert status.value == 0 150 | cudart.cudaMemcpyAsync( 151 | gpu, cpu.ctypes.data, cpu.nbytes, 152 | cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, self.stream) 153 | else: 154 | cpu = self.out_info[i].cpu 155 | gpu = self.out_info[i].gpu 156 | outputs.append(cpu) 157 | output_gpu_ptrs.append(gpu) 158 | self.bindings[j] = gpu 159 | 160 | self.context.execute_async_v2(self.bindings, self.stream) 161 | cudart.cudaStreamSynchronize(self.stream) 162 | 163 | for i, o in enumerate(output_gpu_ptrs): 164 | cudart.cudaMemcpyAsync( 165 | outputs[i].ctypes.data, o, outputs[i].nbytes, 166 | cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, self.stream) 167 | 168 | return tuple(outputs) if len(outputs) > 1 else outputs[0] 169 | -------------------------------------------------------------------------------- /scripts/models/pycuda_api.py: -------------------------------------------------------------------------------- 1 | import os 2 | import warnings 3 | from dataclasses import dataclass 4 | from pathlib import Path 5 | from typing import List, Optional, Tuple, Union 6 | 7 | import numpy as np 8 | import pycuda.autoinit # noqa F401 9 | import pycuda.driver as cuda 10 | import tensorrt as trt 11 | from numpy import ndarray 12 | 13 | os.environ['CUDA_MODULE_LOADING'] = 'LAZY' 14 | warnings.filterwarnings(action='ignore', category=DeprecationWarning) 15 | 16 | 17 | @dataclass 18 | class Tensor: 19 | name: str 20 | dtype: np.dtype 21 | shape: Tuple 22 | cpu: ndarray 23 | gpu: int 24 | 25 | 26 | class TRTEngine: 27 | 28 | def __init__(self, weight: Union[str, Path]) -> None: 29 | self.weight = Path(weight) if isinstance(weight, str) else weight 30 | self.stream = cuda.Stream(0) 31 | self.__init_engine() 32 | self.__init_bindings() 33 | self.__warm_up() 34 | 35 | def __init_engine(self) -> None: 36 | logger = trt.Logger(trt.Logger.WARNING) 37 | trt.init_libnvinfer_plugins(logger, namespace='') 38 | with trt.Runtime(logger) as runtime: 39 | model = runtime.deserialize_cuda_engine(self.weight.read_bytes()) 40 | 41 | context = model.create_execution_context() 42 | 43 | names = [model.get_binding_name(i) for i in range(model.num_bindings)] 44 | self.num_bindings = model.num_bindings 45 | self.bindings: List[int] = [0] * self.num_bindings 46 | num_inputs, num_outputs = 0, 0 47 | 48 | for i in range(model.num_bindings): 49 | if model.binding_is_input(i): 50 | num_inputs += 1 51 | else: 52 | num_outputs += 1 53 | 54 | self.num_inputs = num_inputs 55 | self.num_outputs = num_outputs 56 | self.model = model 57 | self.context = context 58 | self.input_names = names[:num_inputs] 59 | self.output_names = names[num_inputs:] 60 | 61 | def __init_bindings(self) -> None: 62 | dynamic = False 63 | inp_info = [] 64 | out_info = [] 65 | out_ptrs = [] 66 | for i, name in enumerate(self.input_names): 67 | assert self.model.get_binding_name(i) == name 68 | dtype = trt.nptype(self.model.get_binding_dtype(i)) 69 | shape = tuple(self.model.get_binding_shape(i)) 70 | if -1 in shape: 71 | dynamic |= True 72 | if not dynamic: 73 | cpu = np.empty(shape, dtype) 74 | gpu = cuda.mem_alloc(cpu.nbytes) 75 | cuda.memcpy_htod_async(gpu, cpu, self.stream) 76 | else: 77 | cpu, gpu = np.empty(0), 0 78 | inp_info.append(Tensor(name, dtype, shape, cpu, gpu)) 79 | for i, name in enumerate(self.output_names): 80 | i += self.num_inputs 81 | assert self.model.get_binding_name(i) == name 82 | dtype = trt.nptype(self.model.get_binding_dtype(i)) 83 | shape = tuple(self.model.get_binding_shape(i)) 84 | if not dynamic: 85 | cpu = np.empty(shape, dtype=dtype) 86 | gpu = cuda.mem_alloc(cpu.nbytes) 87 | cuda.memcpy_htod_async(gpu, cpu, self.stream) 88 | out_ptrs.append(gpu) 89 | else: 90 | cpu, gpu = np.empty(0), 0 91 | out_info.append(Tensor(name, dtype, shape, cpu, gpu)) 92 | 93 | self.is_dynamic = dynamic 94 | self.inp_info = inp_info 95 | self.out_info = out_info 96 | self.out_ptrs = out_ptrs 97 | 98 | def __warm_up(self) -> None: 99 | if self.is_dynamic: 100 | print('You engine has dynamic axes, please warm up by yourself !') 101 | return 102 | for _ in range(10): 103 | inputs = [] 104 | for i in self.inp_info: 105 | inputs.append(i.cpu) 106 | self.__call__(inputs) 107 | 108 | def set_profiler(self, profiler: Optional[trt.IProfiler]) -> None: 109 | self.context.profiler = profiler \ 110 | if profiler is not None else trt.Profiler() 111 | 112 | def __call__(self, *inputs) -> Union[Tuple, ndarray]: 113 | 114 | assert len(inputs) == self.num_inputs 115 | contiguous_inputs: List[ndarray] = [ 116 | np.ascontiguousarray(i) for i in inputs 117 | ] 118 | 119 | for i in range(self.num_inputs): 120 | 121 | if self.is_dynamic: 122 | self.context.set_binding_shape( 123 | i, tuple(contiguous_inputs[i].shape)) 124 | self.inp_info[i].gpu = cuda.mem_alloc( 125 | contiguous_inputs[i].nbytes) 126 | 127 | cuda.memcpy_htod_async(self.inp_info[i].gpu, contiguous_inputs[i], 128 | self.stream) 129 | self.bindings[i] = int(self.inp_info[i].gpu) 130 | 131 | output_gpu_ptrs: List[int] = [] 132 | outputs: List[ndarray] = [] 133 | 134 | for i in range(self.num_outputs): 135 | j = i + self.num_inputs 136 | if self.is_dynamic: 137 | shape = tuple(self.context.get_binding_shape(j)) 138 | dtype = self.out_info[i].dtype 139 | cpu = np.empty(shape, dtype=dtype) 140 | gpu = cuda.mem_alloc(cpu.nbytes) 141 | cuda.memcpy_htod_async(gpu, cpu, self.stream) 142 | else: 143 | cpu = self.out_info[i].cpu 144 | gpu = self.out_info[i].gpu 145 | outputs.append(cpu) 146 | output_gpu_ptrs.append(gpu) 147 | self.bindings[j] = int(gpu) 148 | 149 | self.context.execute_async_v2(self.bindings, self.stream.handle) 150 | self.stream.synchronize() 151 | 152 | for i, o in enumerate(output_gpu_ptrs): 153 | cuda.memcpy_dtoh_async(outputs[i], o, self.stream) 154 | 155 | return tuple(outputs) if len(outputs) > 1 else outputs[0] 156 | -------------------------------------------------------------------------------- /scripts/models/torch_utils.py: -------------------------------------------------------------------------------- 1 | from typing import List, Tuple, Union 2 | 3 | import torch 4 | import torch.nn.functional as F 5 | from torch import Tensor 6 | from torchvision.ops import batched_nms, nms 7 | 8 | 9 | def seg_postprocess( 10 | data: Tuple[Tensor], 11 | shape: Union[Tuple, List], 12 | conf_thres: float = 0.25, 13 | iou_thres: float = 0.65) \ 14 | -> Tuple[Tensor, Tensor, Tensor, Tensor]: 15 | assert len(data) == 2 16 | h, w = shape[0] // 4, shape[1] // 4 # 4x downsampling 17 | outputs, proto = data[0][0], data[1][0] 18 | bboxes, scores, labels, maskconf = outputs.split([4, 1, 1, 32], 1) 19 | scores, labels = scores.squeeze(), labels.squeeze() 20 | idx = scores > conf_thres 21 | if not idx.any(): # no bounding boxes or seg were created 22 | return bboxes.new_zeros((0, 4)), scores.new_zeros( 23 | (0, )), labels.new_zeros((0, )), bboxes.new_zeros((0, 0, 0, 0)) 24 | bboxes, scores, labels, maskconf = \ 25 | bboxes[idx], scores[idx], labels[idx], maskconf[idx] 26 | idx = batched_nms(bboxes, scores, labels, iou_thres) 27 | bboxes, scores, labels, maskconf = \ 28 | bboxes[idx], scores[idx], labels[idx].int(), maskconf[idx] 29 | masks = (maskconf @ proto).sigmoid().view(-1, h, w) 30 | masks = crop_mask(masks, bboxes / 4.) 31 | masks = F.interpolate(masks[None], 32 | shape, 33 | mode='bilinear', 34 | align_corners=False)[0] 35 | masks = masks.gt_(0.5)[..., None] 36 | return bboxes, scores, labels, masks 37 | 38 | 39 | def pose_postprocess( 40 | data: Union[Tuple, Tensor], 41 | conf_thres: float = 0.25, 42 | iou_thres: float = 0.65) \ 43 | -> Tuple[Tensor, Tensor, Tensor]: 44 | if isinstance(data, tuple): 45 | assert len(data) == 1 46 | data = data[0] 47 | outputs = torch.transpose(data[0], 0, 1).contiguous() 48 | bboxes, scores, kpts = outputs.split([4, 1, 51], 1) 49 | scores, kpts = scores.squeeze(), kpts.squeeze() 50 | idx = scores > conf_thres 51 | if not idx.any(): # no bounding boxes or seg were created 52 | return bboxes.new_zeros((0, 4)), scores.new_zeros( 53 | (0, )), bboxes.new_zeros((0, 0, 0)) 54 | bboxes, scores, kpts = bboxes[idx], scores[idx], kpts[idx] 55 | xycenter, wh = bboxes.chunk(2, -1) 56 | bboxes = torch.cat([xycenter - 0.5 * wh, xycenter + 0.5 * wh], -1) 57 | idx = nms(bboxes, scores, iou_thres) 58 | bboxes, scores, kpts = bboxes[idx], scores[idx], kpts[idx] 59 | return bboxes, scores, kpts.reshape(idx.shape[0], -1, 3) 60 | 61 | 62 | def det_postprocess(data: Tuple[Tensor, Tensor, Tensor, Tensor]): 63 | assert len(data) == 4 64 | iou_thres: float = 0.65 65 | num_dets, bboxes, scores, labels = data[0][0], data[1][0], data[2][ 66 | 0], data[3][0] 67 | nums = num_dets.item() 68 | if nums == 0: 69 | return bboxes.new_zeros((0, 4)), scores.new_zeros( 70 | (0, )), labels.new_zeros((0, )) 71 | # check score negative 72 | scores[scores < 0] = 1 + scores[scores < 0] 73 | # add nms 74 | idx = nms(bboxes, scores, iou_thres) 75 | bboxes, scores, labels = bboxes[idx], scores[idx], labels[idx] 76 | bboxes = bboxes[:nums] 77 | scores = scores[:nums] 78 | labels = labels[:nums] 79 | 80 | return bboxes, scores, labels 81 | 82 | 83 | def crop_mask(masks: Tensor, bboxes: Tensor) -> Tensor: 84 | n, h, w = masks.shape 85 | x1, y1, x2, y2 = torch.chunk(bboxes[:, :, None], 4, 1) # x1 shape(1,1,n) 86 | r = torch.arange(w, device=masks.device, 87 | dtype=x1.dtype)[None, None, :] # rows shape(1,w,1) 88 | c = torch.arange(h, device=masks.device, 89 | dtype=x1.dtype)[None, :, None] # cols shape(h,1,1) 90 | 91 | return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2)) 92 | -------------------------------------------------------------------------------- /scripts/models/utils.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | from typing import List, Tuple, Union 3 | 4 | import cv2 5 | import numpy as np 6 | from numpy import ndarray 7 | from torchvision.ops import nms 8 | 9 | # image suffixs 10 | SUFFIXS = ('.bmp', '.dng', '.jpeg', '.jpg', '.mpo', '.png', '.tif', '.tiff', 11 | '.webp', '.pfm') 12 | 13 | 14 | def letterbox(im: ndarray, 15 | new_shape: Union[Tuple, List] = (640, 640), 16 | color: Union[Tuple, List] = (114, 114, 114)) \ 17 | -> Tuple[ndarray, float, Tuple[float, float]]: 18 | # Resize and pad image while meeting stride-multiple constraints 19 | shape = im.shape[:2] # current shape [height, width] 20 | if isinstance(new_shape, int): 21 | new_shape = (new_shape, new_shape) 22 | # new_shape: [width, height] 23 | 24 | # Scale ratio (new / old) 25 | r = min(new_shape[0] / shape[1], new_shape[1] / shape[0]) 26 | # Compute padding [width, height] 27 | new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) 28 | dw, dh = new_shape[0] - new_unpad[0], new_shape[1] - new_unpad[ 29 | 1] # wh padding 30 | 31 | dw /= 2 # divide padding into 2 sides 32 | dh /= 2 33 | 34 | if shape[::-1] != new_unpad: # resize 35 | im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR) 36 | top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) 37 | left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) 38 | im = cv2.copyMakeBorder(im, 39 | top, 40 | bottom, 41 | left, 42 | right, 43 | cv2.BORDER_CONSTANT, 44 | value=color) # add border 45 | return im, r, (dw, dh) 46 | 47 | 48 | def blob(im: ndarray, return_seg: bool = False) -> Union[ndarray, Tuple]: 49 | seg = None 50 | if return_seg: 51 | seg = im.astype(np.float32) / 255 52 | im = im.transpose([2, 0, 1]) 53 | im = im[np.newaxis, ...] 54 | im = np.ascontiguousarray(im).astype(np.float32) / 255 55 | if return_seg: 56 | return im, seg 57 | else: 58 | return im 59 | 60 | 61 | def sigmoid(x): 62 | return 1. / (1. + np.exp(-x)) 63 | 64 | 65 | def path_to_list(images_path: Union[str, Path]) -> List: 66 | if isinstance(images_path, str): 67 | images_path = Path(images_path) 68 | assert images_path.exists() 69 | if images_path.is_dir(): 70 | images = [ 71 | i.absolute() for i in images_path.iterdir() if i.suffix in SUFFIXS 72 | ] 73 | else: 74 | assert images_path.suffix in SUFFIXS 75 | images = [images_path.absolute()] 76 | return images 77 | 78 | 79 | def crop_mask(masks: ndarray, bboxes: ndarray) -> ndarray: 80 | n, h, w = masks.shape 81 | x1, y1, x2, y2 = np.split(bboxes[:, :, None], [1, 2, 3], 82 | 1) # x1 shape(1,1,n) 83 | r = np.arange(w, dtype=x1.dtype)[None, None, :] # rows shape(1,w,1) 84 | c = np.arange(h, dtype=x1.dtype)[None, :, None] # cols shape(h,1,1) 85 | 86 | return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2)) 87 | 88 | 89 | def det_postprocess(data: Tuple[ndarray, ndarray, ndarray, ndarray]): 90 | assert len(data) == 4 91 | iou_thres: float = 0.65 92 | num_dets, bboxes, scores, labels = (i[0] for i in data) 93 | nums = num_dets.item() 94 | if nums == 0: 95 | return np.empty((0, 4), dtype=np.float32), np.empty( 96 | (0, ), dtype=np.float32), np.empty((0, ), dtype=np.int32) 97 | # check score negative 98 | scores[scores < 0] = 1 + scores[scores < 0] 99 | # add nms 100 | idx = nms(bboxes, scores, iou_thres) 101 | bboxes, scores, labels = bboxes[idx], scores[idx], labels[idx] 102 | 103 | bboxes = bboxes[:nums] 104 | scores = scores[:nums] 105 | labels = labels[:nums] 106 | return bboxes, scores, labels 107 | 108 | 109 | def seg_postprocess( 110 | data: Tuple[ndarray], 111 | shape: Union[Tuple, List], 112 | conf_thres: float = 0.25, 113 | iou_thres: float = 0.65) \ 114 | -> Tuple[ndarray, ndarray, ndarray, ndarray]: 115 | assert len(data) == 2 116 | h, w = shape[0] // 4, shape[1] // 4 # 4x downsampling 117 | outputs, proto = (i[0] for i in data) 118 | bboxes, scores, labels, maskconf = np.split(outputs, [4, 5, 6], 1) 119 | scores, labels = scores.squeeze(), labels.squeeze() 120 | idx = scores > conf_thres 121 | if not idx.any(): # no bounding boxes or seg were created 122 | return np.empty((0, 4), dtype=np.float32), \ 123 | np.empty((0,), dtype=np.float32), \ 124 | np.empty((0,), dtype=np.int32), \ 125 | np.empty((0, 0, 0, 0), dtype=np.int32) 126 | 127 | bboxes, scores, labels, maskconf = \ 128 | bboxes[idx], scores[idx], labels[idx], maskconf[idx] 129 | cvbboxes = np.concatenate([bboxes[:, :2], bboxes[:, 2:] - bboxes[:, :2]], 130 | 1) 131 | labels = labels.astype(np.int32) 132 | v0, v1 = map(int, (cv2.__version__).split('.')[:2]) 133 | assert v0 == 4, 'OpenCV version is wrong' 134 | if v1 > 6: 135 | idx = cv2.dnn.NMSBoxesBatched(cvbboxes, scores, labels, conf_thres, 136 | iou_thres) 137 | else: 138 | idx = cv2.dnn.NMSBoxes(cvbboxes, scores, conf_thres, iou_thres) 139 | bboxes, scores, labels, maskconf = \ 140 | bboxes[idx], scores[idx], labels[idx], maskconf[idx] 141 | masks = sigmoid(maskconf @ proto).reshape(-1, h, w) 142 | masks = crop_mask(masks, bboxes / 4.) 143 | masks = masks.transpose([1, 2, 0]) 144 | masks = cv2.resize(masks, (shape[1], shape[0]), 145 | interpolation=cv2.INTER_LINEAR) 146 | masks = masks.transpose(2, 0, 1) 147 | masks = np.ascontiguousarray((masks > 0.5)[..., None], dtype=np.float32) 148 | return bboxes, scores, labels, masks 149 | 150 | 151 | def pose_postprocess( 152 | data: Union[Tuple, ndarray], 153 | conf_thres: float = 0.25, 154 | iou_thres: float = 0.65) \ 155 | -> Tuple[ndarray, ndarray, ndarray]: 156 | if isinstance(data, tuple): 157 | assert len(data) == 1 158 | data = data[0] 159 | outputs = np.transpose(data[0], (1, 0)) 160 | bboxes, scores, kpts = np.split(outputs, [4, 5], 1) 161 | scores, kpts = scores.squeeze(), kpts.squeeze() 162 | idx = scores > conf_thres 163 | if not idx.any(): # no bounding boxes or seg were created 164 | return np.empty((0, 4), dtype=np.float32), np.empty( 165 | (0, ), dtype=np.float32), np.empty((0, 0, 0), dtype=np.float32) 166 | bboxes, scores, kpts = bboxes[idx], scores[idx], kpts[idx] 167 | xycenter, wh = np.split(bboxes, [ 168 | 2, 169 | ], -1) 170 | cvbboxes = np.concatenate([xycenter - 0.5 * wh, wh], -1) 171 | idx = cv2.dnn.NMSBoxes(cvbboxes, scores, conf_thres, iou_thres) 172 | cvbboxes, scores, kpts = cvbboxes[idx], scores[idx], kpts[idx] 173 | cvbboxes[:, 2:] += cvbboxes[:, :2] 174 | return cvbboxes, scores, kpts.reshape(idx.shape[0], -1, 3) 175 | -------------------------------------------------------------------------------- /scripts/trt-profile.py: -------------------------------------------------------------------------------- 1 | from models import TRTModule, TRTProfilerV0 # isort:skip 2 | import argparse 3 | 4 | import torch 5 | 6 | 7 | def profile(args): 8 | device = torch.device(args.device) 9 | Engine = TRTModule(args.engine, device) 10 | profiler = TRTProfilerV0() 11 | Engine.set_profiler(profiler) 12 | random_input = torch.randn(Engine.inp_info[0].shape, device=device) 13 | _ = Engine(random_input) 14 | 15 | 16 | def parse_args(): 17 | parser = argparse.ArgumentParser() 18 | parser.add_argument('--engine', type=str, help='Engine file') 19 | 20 | parser.add_argument('--device', 21 | type=str, 22 | default='cuda:0', 23 | help='TensorRT infer device') 24 | args = parser.parse_args() 25 | return args 26 | 27 | 28 | if __name__ == '__main__': 29 | args = parse_args() 30 | profile(args) 31 | -------------------------------------------------------------------------------- /src/deepstream/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.1) 2 | 3 | set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86 89 90) 4 | set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc) 5 | 6 | project(nvdsinfer_custom_bbox_yoloV8 LANGUAGES CXX CUDA) 7 | 8 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -O3 -Wall -Werror -shared -fPIC") 9 | set(CMAKE_CXX_STANDARD 14) 10 | set(CMAKE_BUILD_TYPE Release) 11 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 12 | 13 | # CUDA 14 | find_package(CUDA REQUIRED) 15 | message(STATUS "CUDA Libs: \n${CUDA_LIBRARIES}\n") 16 | get_filename_component(CUDA_LIB_DIR ${CUDA_LIBRARIES} DIRECTORY) 17 | message(STATUS "CUDA Headers: \n${CUDA_INCLUDE_DIRS}\n") 18 | 19 | 20 | # TensorRT 21 | set(TensorRT_INCLUDE_DIRS /usr/include/x86_64-linux-gnu) 22 | set(TensorRT_LIBRARIES /usr/lib/x86_64-linux-gnu) 23 | 24 | message(STATUS "TensorRT Libs: \n${TensorRT_LIBRARIES}\n") 25 | message(STATUS "TensorRT Headers: \n${TensorRT_INCLUDE_DIRS}\n") 26 | 27 | 28 | set(DEEPSTREAM /opt/nvidia/deepstream/deepstream) 29 | set(DS_LIBRARIES ${DEEPSTREAM}/lib) 30 | set(DS_INCLUDE_DIRS ${DEEPSTREAM}/sources/includes) 31 | message(STATUS "DS Libs: \n${DS_LIBRARIES}\n") 32 | message(STATUS "DS Headers: \n${DS_INCLUDE_DIRS}\n") 33 | 34 | 35 | 36 | list(APPEND INCLUDE_DIRS 37 | ${CUDA_INCLUDE_DIRS} 38 | ${TensorRT_INCLUDE_DIRS} 39 | ${DS_INCLUDE_DIRS} 40 | ) 41 | 42 | list(APPEND ALL_LIBS 43 | ${CUDA_LIBRARIES} 44 | ${CUDA_LIB_DIR} 45 | ${TensorRT_LIBRARIES} 46 | ${DS_LIBRARIES} 47 | ) 48 | 49 | include_directories(${INCLUDE_DIRS}) 50 | 51 | add_library( 52 | ${PROJECT_NAME} 53 | SHARED 54 | custom_bbox_parser/nvdsparsebbox_yoloV8.cpp 55 | ) 56 | 57 | target_link_libraries(${PROJECT_NAME} PRIVATE nvinfer nvinfer_plugin) 58 | -------------------------------------------------------------------------------- /src/deepstream/README.md: -------------------------------------------------------------------------------- 1 | # Start DeepStream Using the engine build from [`YOLOv8-TensorRT`](https://github.com/triple-Mu/YOLOv8-TensorRT) 2 | 3 | ## 1. Build you own TensorRT engine from `trtexec` or [`build.py`](https://github.com/triple-Mu/YOLOv8-TensorRT/blob/main/build.py) 4 | 5 | For example, if you have built an engine named `yolov8s.engine`. 6 | 7 | ## 2. Compile deepstream plugin 8 | 9 | First, modify the [`CMakeLists.txt`](https://github.com/triple-Mu/YOLOv8-TensorRT/blob/main/csrc/deepstream/CMakeLists.txt) 10 | 11 | ```cmake 12 | # Set your own TensorRT path 13 | set(TensorRT_INCLUDE_DIRS /usr/include/x86_64-linux-gnu) 14 | set(TensorRT_LIBRARIES /usr/lib/x86_64-linux-gnu) 15 | # Set your own DeepStream path 16 | set(DEEPSTREAM /opt/nvidia/deepstream/deepstream) 17 | ``` 18 | 19 | Second, build deepstream plugin 20 | 21 | ```shell 22 | mkdir build 23 | cd build 24 | cmake .. 25 | make 26 | ``` 27 | You will get a lib `libnvdsinfer_custom_bbox_yoloV8.so` in `build`. 28 | 29 | 30 | ## 3. Modify the deepstream config 31 | 32 | The net config is [`config_yoloV8.txt`](config_yoloV8.txt). Please modify by your own model. 33 | 34 | ```text 35 | net-scale-factor=0.0039215697906911373 # the normalize param == 1/255 36 | model-engine-file=./yolov8s.engine # the engine path you build 37 | labelfile-path=./labels.txt # the class name path 38 | num-detected-classes=80 # the number of classes 39 | output-blob-names=num_dets;bboxes;scores;labels # the model output names 40 | custom-lib-path=./build/libnvdsinfer_custom_bbox_yoloV8.so # the deepstream plugin you build 41 | ``` 42 | 43 | The deepstream config is [`deepstream_app_config.txt`](deepstream_app_config.txt). 44 | 45 | ```text 46 | **** 47 | [source0] 48 | enable=1 49 | #Type - 1=CameraV4L2 2=URI 3=MultiURI 50 | type=3 51 | uri=file://./sample_1080p_h264.mp4 # the video path or stream you want to detect 52 | **** 53 | **** 54 | config-file=config_yoloV8.txt # the net config path 55 | ``` 56 | 57 | You can get more information from [`deepstream offical`](https://developer.nvidia.com/deepstream-sdk). 58 | 59 | ## 4. Runing detector ! 60 | 61 | ```shell 62 | deepstream-app -c deepstream_app_config.txt 63 | ``` 64 | -------------------------------------------------------------------------------- /src/deepstream/config_yoloV8.txt: -------------------------------------------------------------------------------- 1 | ################################################################################ 2 | # Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. 3 | # 4 | # Permission is hereby granted, free of charge, to any person obtaining a 5 | # copy of this software and associated documentation files (the "Software"), 6 | # to deal in the Software without restriction, including without limitation 7 | # the rights to use, copy, modify, merge, publish, distribute, sublicense, 8 | # and/or sell copies of the Software, and to permit persons to whom the 9 | # Software is furnished to do so, subject to the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be included in 12 | # all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 17 | # THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 20 | # DEALINGS IN THE SOFTWARE. 21 | ################################################################################ 22 | 23 | # Following properties are mandatory when engine files are not specified: 24 | # int8-calib-file(Only in INT8), model-file-format 25 | # Caffemodel mandatory properties: model-file, proto-file, output-blob-names 26 | # UFF: uff-file, input-dims, uff-input-blob-name, output-blob-names 27 | # ONNX: onnx-file 28 | # 29 | # Mandatory properties for detectors: 30 | # num-detected-classes 31 | # 32 | # Optional properties for detectors: 33 | # enable-dbscan(Default=false), interval(Primary mode only, Default=0) 34 | # custom-lib-path 35 | # parse-bbox-func-name 36 | # 37 | # Mandatory properties for classifiers: 38 | # classifier-threshold, is-classifier 39 | # 40 | # Optional properties for classifiers: 41 | # classifier-async-mode(Secondary mode only, Default=false) 42 | # 43 | # Optional properties in secondary mode: 44 | # operate-on-gie-id(Default=0), operate-on-class-ids(Defaults to all classes), 45 | # input-object-min-width, input-object-min-height, input-object-max-width, 46 | # input-object-max-height 47 | # 48 | # Following properties are always recommended: 49 | # batch-size(Default=1) 50 | # 51 | # Other optional properties: 52 | # net-scale-factor(Default=1), network-mode(Default=0 i.e FP32), 53 | # model-color-format(Default=0 i.e. RGB) model-engine-file, labelfile-path, 54 | # mean-file, gie-unique-id(Default=0), offsets, gie-mode (Default=1 i.e. primary), 55 | # custom-lib-path, network-mode(Default=0 i.e FP32) 56 | # 57 | # The values in the config file are overridden by values set through GObject 58 | # properties. 59 | 60 | [property] 61 | net-scale-factor=0.0039215697906911373 62 | gpu-id=0 63 | #0=RGB, 1=BGR 64 | model-color-format=0 65 | model-engine-file=./yolov8s.engine 66 | labelfile-path=./labels.txt 67 | ## 0=FP32, 1=INT8, 2=FP16 mode 68 | network-mode=2 69 | num-detected-classes=80 70 | gie-unique-id=1 71 | is-classifier=0 72 | maintain-aspect-ratio=1 73 | output-blob-names=num_dets;bboxes;scores;labels 74 | parse-bbox-func-name=NvDsInferParseCustomYoloV8 75 | custom-lib-path=./build/libnvdsinfer_custom_bbox_yoloV8.so 76 | -------------------------------------------------------------------------------- /src/deepstream/custom_bbox_parser/nvdsparsebbox_yoloV8.cpp: -------------------------------------------------------------------------------- 1 | #include "nvdsinfer_custom_impl.h" 2 | #include 3 | #include 4 | 5 | /** 6 | * Function expected by DeepStream for decoding the TinyYOLOv2 output. 7 | * 8 | * C-linkage [extern "C"] was written to prevent name-mangling. This function must return true after 9 | * adding all bounding boxes to the objectList vector. 10 | * 11 | * @param [outputLayersInfo] std::vector of NvDsInferLayerInfo objects with information about the output layer. 12 | * @param [networkInfo] NvDsInferNetworkInfo object with information about the TinyYOLOv2 network. 13 | * @param [detectionParams] NvDsInferParseDetectionParams with information about some config params. 14 | * @param [objectList] std::vector of NvDsInferParseObjectInfo objects to which bounding box information must 15 | * be stored. 16 | * 17 | * @return true 18 | */ 19 | 20 | // This is just the function prototype. The definition is written at the end of the file. 21 | extern "C" bool NvDsInferParseCustomYoloV8(std::vector const& outputLayersInfo, 22 | NvDsInferNetworkInfo const& networkInfo, 23 | NvDsInferParseDetectionParams const& detectionParams, 24 | std::vector& objectList); 25 | 26 | static __inline__ float bbox_clip(const float& val, const float& minVal = 0.f, const float& maxVal = 1280.f) 27 | { 28 | assert(minVal <= maxVal); 29 | return std::max(std::min(val, (maxVal - 1)), minVal); 30 | } 31 | 32 | static std::vector decodeYoloV8Tensor(const int* num_dets, 33 | const float* bboxes, 34 | const float* scores, 35 | const int* labels, 36 | const unsigned int& img_w, 37 | const unsigned int& img_h) 38 | { 39 | std::vector bboxInfo; 40 | size_t nums = num_dets[0]; 41 | for (size_t i = 0; i < nums; i++) { 42 | float x0 = (bboxes[i * 4]); 43 | float y0 = (bboxes[i * 4 + 1]); 44 | float x1 = (bboxes[i * 4 + 2]); 45 | float y1 = (bboxes[i * 4 + 3]); 46 | x0 = bbox_clip(x0, 0.f, img_w); 47 | y0 = bbox_clip(y0, 0.f, img_h); 48 | x1 = bbox_clip(x1, 0.f, img_w); 49 | y1 = bbox_clip(y1, 0.f, img_h); 50 | NvDsInferParseObjectInfo obj; 51 | obj.left = x0; 52 | obj.top = y0; 53 | obj.width = x1 - x0; 54 | obj.height = y1 - y0; 55 | obj.detectionConfidence = scores[i]; 56 | obj.classId = labels[i]; 57 | bboxInfo.push_back(obj); 58 | } 59 | 60 | return bboxInfo; 61 | } 62 | 63 | /* C-linkage to prevent name-mangling */ 64 | extern "C" bool NvDsInferParseCustomYoloV8(std::vector const& outputLayersInfo, 65 | NvDsInferNetworkInfo const& networkInfo, 66 | NvDsInferParseDetectionParams const& detectionParams, 67 | std::vector& objectList) 68 | { 69 | 70 | // Some assertions and error checking. 71 | 72 | if (outputLayersInfo.empty() || outputLayersInfo.size() != 4) { 73 | std::cerr << "Could not find output layer in bbox parsing" << std::endl; 74 | return false; 75 | } 76 | 77 | // Obtaining the output layer. 78 | const NvDsInferLayerInfo& num_dets = outputLayersInfo[0]; 79 | const NvDsInferLayerInfo& bboxes = outputLayersInfo[1]; 80 | const NvDsInferLayerInfo& scores = outputLayersInfo[2]; 81 | const NvDsInferLayerInfo& labels = outputLayersInfo[3]; 82 | 83 | // num_dets(int) bboxes(float) scores(float) labels(int) 84 | assert(num_dets.dims.numDims == 2); 85 | assert(bboxes.dims.numDims == 3); 86 | assert(scores.dims.numDims == 2); 87 | assert(labels.dims.numDims == 2); 88 | 89 | // Decoding the output tensor of YOLOv8 to the NvDsInferParseObjectInfo format. 90 | std::vector objects = decodeYoloV8Tensor((const int*)(num_dets.buffer), 91 | (const float*)(bboxes.buffer), 92 | (const float*)(scores.buffer), 93 | (const int*)(labels.buffer), 94 | networkInfo.width, 95 | networkInfo.height); 96 | 97 | objectList.clear(); 98 | objectList = objects; 99 | return true; 100 | } 101 | 102 | /* Check that the custom function has been defined correctly */ 103 | CHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseCustomYoloV8); 104 | -------------------------------------------------------------------------------- /src/deepstream/deepstream_app_config.txt: -------------------------------------------------------------------------------- 1 | [application] 2 | enable-perf-measurement=1 3 | perf-measurement-interval-sec=5 4 | #gie-kitti-output-dir=streamscl 5 | 6 | [tiled-display] 7 | enable=1 8 | rows=1 9 | columns=1 10 | # show 11 | width=1920 12 | height=1080 13 | gpu-id=0 14 | #(0): nvbuf-mem-default - Default memory allocated, specific to particular platform 15 | #(1): nvbuf-mem-cuda-pinned - Allocate Pinned/Host cuda memory, applicable for Tesla 16 | #(2): nvbuf-mem-cuda-device - Allocate Device cuda memory, applicable for Tesla 17 | #(3): nvbuf-mem-cuda-unified - Allocate Unified cuda memory, applicable for Tesla 18 | #(4): nvbuf-mem-surface-array - Allocate Surface Array memory, applicable for Jetson 19 | nvbuf-memory-type=0 20 | 21 | [source0] 22 | enable=1 23 | #Type - 1=CameraV4L2 2=URI 3=MultiURI 24 | type=3 25 | uri=file://./sample_1080p_h264.mp4 26 | num-sources=1 27 | drop-frame-interval=1 28 | gpu-id=0 29 | # (0): memtype_device - Memory type Device 30 | # (1): memtype_pinned - Memory type Host Pinned 31 | # (2): memtype_unified - Memory type Unified 32 | cudadec-memtype=0 33 | 34 | [sink0] 35 | enable=1 36 | #Type - 1=FakeSink 2=EglSink 3=File 37 | type=2 38 | sync=0 39 | source-id=0 40 | gpu-id=0 41 | nvbuf-memory-type=0 42 | 43 | #[sink1] 44 | #enable=1 45 | #Type - 1=FakeSink 2=EglSink 3=File 4=RTSPStreaming 46 | #type=4 47 | ##1=h264 2=h265 48 | #codec=1 49 | #sync=0 50 | #bitrate=4000000 51 | ## set below properties in case of RTSPStreaming 52 | #rtsp-port=8554 53 | #udp-port=5400 54 | 55 | #[sink2] 56 | #enable=1 57 | ##Type - 1=FakeSink 2=EglSink 3=File 4=RTSPStreaming 58 | #type=3 59 | ##1=h264 2=h265 60 | #codec=1 61 | #sync=0 62 | #bitrate=4000000 63 | #container=1 64 | #output-file=./output.mp4 65 | 66 | [osd] 67 | enable=1 68 | gpu-id=0 69 | border-width=1 70 | text-size=15 71 | text-color=1;1;1;1; 72 | text-bg-color=0.3;0.3;0.3;1 73 | font=Serif 74 | show-clock=0 75 | clock-x-offset=800 76 | clock-y-offset=820 77 | clock-text-size=12 78 | clock-color=1;0;0;1 79 | nvbuf-memory-type=0 80 | 81 | [streammux] 82 | gpu-id=0 83 | ##Boolean property to inform muxer that sources are live 84 | live-source=0 85 | batch-size=1 86 | ##time out in usec, to wait after the first buffer is available 87 | ##to push the batch even if the complete batch is not formed 88 | batched-push-timeout=40000 89 | ## Set muxer output width and height 90 | width=1920 91 | height=1080 92 | ##Enable to maintain aspect ratio wrt source, and allow black borders, works 93 | ##along with width, height properties 94 | enable-padding=0 95 | nvbuf-memory-type=0 96 | 97 | # config-file property is mandatory for any gie section. 98 | # Other properties are optional and if set will override the properties set in 99 | # the infer config file. 100 | [primary-gie] 101 | enable=1 102 | gpu-id=0 103 | model-engine-file=yolov8s.engine 104 | labelfile-path=./labels.txt 105 | batch-size=1 106 | #Required by the app for OSD, not a plugin property 107 | bbox-border-color0=1;0;0;1 108 | bbox-border-color1=0;1;1;1 109 | bbox-border-color2=0;0;1;1 110 | bbox-border-color3=0;1;0;1 111 | gie-unique-id=1 112 | nvbuf-memory-type=0 113 | config-file=config_yoloV8.txt 114 | 115 | [tests] 116 | file-loop=0 117 | -------------------------------------------------------------------------------- /src/deepstream/labels.txt: -------------------------------------------------------------------------------- 1 | person 2 | bicycle 3 | car 4 | motorbike 5 | aeroplane 6 | bus 7 | train 8 | truck 9 | boat 10 | traffic light 11 | fire hydrant 12 | stop sign 13 | parking meter 14 | bench 15 | bird 16 | cat 17 | dog 18 | horse 19 | sheep 20 | cow 21 | elephant 22 | bear 23 | zebra 24 | giraffe 25 | backpack 26 | umbrella 27 | handbag 28 | tie 29 | suitcase 30 | frisbee 31 | skis 32 | snowboard 33 | sports ball 34 | kite 35 | baseball bat 36 | baseball glove 37 | skateboard 38 | surfboard 39 | tennis racket 40 | bottle 41 | wine glass 42 | cup 43 | fork 44 | knife 45 | spoon 46 | bowl 47 | banana 48 | apple 49 | sandwich 50 | orange 51 | broccoli 52 | carrot 53 | hot dog 54 | pizza 55 | donut 56 | cake 57 | chair 58 | sofa 59 | pottedplant 60 | bed 61 | diningtable 62 | toilet 63 | tvmonitor 64 | laptop 65 | mouse 66 | remote 67 | keyboard 68 | cell phone 69 | microwave 70 | oven 71 | toaster 72 | sink 73 | refrigerator 74 | book 75 | clock 76 | vase 77 | scissors 78 | teddy bear 79 | hair drier 80 | toothbrush 81 | -------------------------------------------------------------------------------- /src/detect/end2end/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.1) 2 | 3 | set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86 89 90) 4 | set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc) 5 | 6 | project(yolov8 LANGUAGES CXX CUDA) 7 | 8 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -O3") 9 | set(CMAKE_CXX_STANDARD 14) 10 | set(CMAKE_BUILD_TYPE Release) 11 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 12 | 13 | # CUDA 14 | find_package(CUDA REQUIRED) 15 | message(STATUS "CUDA Libs: \n${CUDA_LIBRARIES}\n") 16 | get_filename_component(CUDA_LIB_DIR ${CUDA_LIBRARIES} DIRECTORY) 17 | message(STATUS "CUDA Headers: \n${CUDA_INCLUDE_DIRS}\n") 18 | 19 | # OpenCV 20 | find_package(OpenCV REQUIRED) 21 | message(STATUS "OpenCV Libs: \n${OpenCV_LIBS}\n") 22 | message(STATUS "OpenCV Libraries: \n${OpenCV_LIBRARIES}\n") 23 | message(STATUS "OpenCV Headers: \n${OpenCV_INCLUDE_DIRS}\n") 24 | 25 | # TensorRT # TensorRT_ROOT的路径设置成自己的 26 | # set(TensorRT_ROOT /root/share/TensorRT-8.5.3.1) 27 | set(TensorRT_ROOT /home/lin/software/TensorRT-8.5.3.1) 28 | set(TensorRT_INCLUDE_DIRS ${TensorRT_ROOT}/include) 29 | set(TensorRT_LIBRARIES ${TensorRT_ROOT}/lib) 30 | 31 | message(STATUS "TensorRT Libs: \n${TensorRT_LIBRARIES}\n") 32 | message(STATUS "TensorRT Headers: \n${TensorRT_INCLUDE_DIRS}\n") 33 | 34 | list(APPEND INCLUDE_DIRS 35 | ${CUDA_INCLUDE_DIRS} 36 | ${OpenCV_INCLUDE_DIRS} 37 | ${TensorRT_INCLUDE_DIRS} 38 | include 39 | ) 40 | 41 | list(APPEND ALL_LIBS 42 | ${CUDA_LIBRARIES} 43 | ${CUDA_LIB_DIR} 44 | ${OpenCV_LIBRARIES} 45 | ${TensorRT_LIBRARIES} 46 | ) 47 | 48 | include_directories(${INCLUDE_DIRS}) 49 | 50 | add_executable(${PROJECT_NAME} 51 | main.cpp 52 | include/yolov8.hpp 53 | include/common.hpp 54 | ) 55 | 56 | target_link_directories(${PROJECT_NAME} PUBLIC ${ALL_LIBS}) 57 | target_link_libraries(${PROJECT_NAME} PRIVATE nvinfer nvinfer_plugin cudart ${OpenCV_LIBS}) 58 | -------------------------------------------------------------------------------- /src/detect/end2end/include/common.hpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 1/24/23. 3 | // 4 | 5 | #ifndef DETECT_END2END_COMMON_HPP 6 | #define DETECT_END2END_COMMON_HPP 7 | #include "NvInfer.h" 8 | #include "opencv2/opencv.hpp" 9 | #include 10 | #include 11 | 12 | #define CHECK(call) \ 13 | do { \ 14 | const cudaError_t error_code = call; \ 15 | if (error_code != cudaSuccess) { \ 16 | printf("CUDA Error:\n"); \ 17 | printf(" File: %s\n", __FILE__); \ 18 | printf(" Line: %d\n", __LINE__); \ 19 | printf(" Error code: %d\n", error_code); \ 20 | printf(" Error text: %s\n", cudaGetErrorString(error_code)); \ 21 | exit(1); \ 22 | } \ 23 | } while (0) 24 | 25 | class Logger: public nvinfer1::ILogger { 26 | public: 27 | nvinfer1::ILogger::Severity reportableSeverity; 28 | 29 | explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO): 30 | reportableSeverity(severity) 31 | { 32 | } 33 | 34 | void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override 35 | { 36 | if (severity > reportableSeverity) { 37 | return; 38 | } 39 | switch (severity) { 40 | case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: 41 | std::cerr << "INTERNAL_ERROR: "; 42 | break; 43 | case nvinfer1::ILogger::Severity::kERROR: 44 | std::cerr << "ERROR: "; 45 | break; 46 | case nvinfer1::ILogger::Severity::kWARNING: 47 | std::cerr << "WARNING: "; 48 | break; 49 | case nvinfer1::ILogger::Severity::kINFO: 50 | std::cerr << "INFO: "; 51 | break; 52 | default: 53 | std::cerr << "VERBOSE: "; 54 | break; 55 | } 56 | std::cerr << msg << std::endl; 57 | } 58 | }; 59 | 60 | inline int get_size_by_dims(const nvinfer1::Dims& dims) 61 | { 62 | int size = 1; 63 | for (int i = 0; i < dims.nbDims; i++) { 64 | size *= dims.d[i]; 65 | } 66 | return size; 67 | } 68 | 69 | inline int type_to_size(const nvinfer1::DataType& dataType) 70 | { 71 | switch (dataType) { 72 | case nvinfer1::DataType::kFLOAT: 73 | return 4; 74 | case nvinfer1::DataType::kHALF: 75 | return 2; 76 | case nvinfer1::DataType::kINT32: 77 | return 4; 78 | case nvinfer1::DataType::kINT8: 79 | return 1; 80 | case nvinfer1::DataType::kBOOL: 81 | return 1; 82 | default: 83 | return 4; 84 | } 85 | } 86 | 87 | inline static float clamp(float val, float min, float max) 88 | { 89 | return val > min ? (val < max ? val : max) : min; 90 | } 91 | 92 | inline bool IsPathExist(const std::string& path) 93 | { 94 | if (access(path.c_str(), 0) == F_OK) { 95 | return true; 96 | } 97 | return false; 98 | } 99 | 100 | inline bool IsFile(const std::string& path) 101 | { 102 | if (!IsPathExist(path)) { 103 | printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str()); 104 | return false; 105 | } 106 | struct stat buffer; 107 | return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode)); 108 | } 109 | 110 | inline bool IsFolder(const std::string& path) 111 | { 112 | if (!IsPathExist(path)) { 113 | return false; 114 | } 115 | struct stat buffer; 116 | return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode)); 117 | } 118 | 119 | namespace det { 120 | struct Binding { 121 | size_t size = 1; 122 | size_t dsize = 1; 123 | nvinfer1::Dims dims; 124 | std::string name; 125 | }; 126 | 127 | struct Object { 128 | cv::Rect_ rect; 129 | int label = 0; 130 | float prob = 0.0; 131 | }; 132 | 133 | struct PreParam { 134 | float ratio = 1.0f; 135 | float dw = 0.0f; 136 | float dh = 0.0f; 137 | float height = 0; 138 | float width = 0; 139 | }; 140 | } // namespace det 141 | #endif // DETECT_END2END_COMMON_HPP 142 | -------------------------------------------------------------------------------- /src/detect/end2end/main.cpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 1/20/23. 3 | // 4 | #include "chrono" 5 | #include "opencv2/opencv.hpp" 6 | #include "yolov8.hpp" 7 | 8 | const std::vector CLASS_NAMES = { 9 | "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", 10 | "truck", "boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", 11 | "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", 12 | "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", 13 | "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat", 14 | "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup", 15 | "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", 16 | "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", 17 | "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", 18 | "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven", 19 | "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", 20 | "teddy bear", "hair drier", "toothbrush"}; 21 | 22 | const std::vector> COLORS = { 23 | {0, 114, 189}, {217, 83, 25}, {237, 177, 32}, {126, 47, 142}, {119, 172, 48}, {77, 190, 238}, 24 | {162, 20, 47}, {76, 76, 76}, {153, 153, 153}, {255, 0, 0}, {255, 128, 0}, {191, 191, 0}, 25 | {0, 255, 0}, {0, 0, 255}, {170, 0, 255}, {85, 85, 0}, {85, 170, 0}, {85, 255, 0}, 26 | {170, 85, 0}, {170, 170, 0}, {170, 255, 0}, {255, 85, 0}, {255, 170, 0}, {255, 255, 0}, 27 | {0, 85, 128}, {0, 170, 128}, {0, 255, 128}, {85, 0, 128}, {85, 85, 128}, {85, 170, 128}, 28 | {85, 255, 128}, {170, 0, 128}, {170, 85, 128}, {170, 170, 128}, {170, 255, 128}, {255, 0, 128}, 29 | {255, 85, 128}, {255, 170, 128}, {255, 255, 128}, {0, 85, 255}, {0, 170, 255}, {0, 255, 255}, 30 | {85, 0, 255}, {85, 85, 255}, {85, 170, 255}, {85, 255, 255}, {170, 0, 255}, {170, 85, 255}, 31 | {170, 170, 255}, {170, 255, 255}, {255, 0, 255}, {255, 85, 255}, {255, 170, 255}, {85, 0, 0}, 32 | {128, 0, 0}, {170, 0, 0}, {212, 0, 0}, {255, 0, 0}, {0, 43, 0}, {0, 85, 0}, 33 | {0, 128, 0}, {0, 170, 0}, {0, 212, 0}, {0, 255, 0}, {0, 0, 43}, {0, 0, 85}, 34 | {0, 0, 128}, {0, 0, 170}, {0, 0, 212}, {0, 0, 255}, {0, 0, 0}, {36, 36, 36}, 35 | {73, 73, 73}, {109, 109, 109}, {146, 146, 146}, {182, 182, 182}, {219, 219, 219}, {0, 114, 189}, 36 | {80, 183, 189}, {128, 128, 0}}; 37 | 38 | int main(int argc, char** argv) 39 | { 40 | // cuda:0 41 | cudaSetDevice(0); 42 | 43 | const std::string engine_file_path{argv[1]}; 44 | const std::string path{argv[2]}; 45 | 46 | std::vector imagePathList; 47 | bool isVideo{false}; 48 | 49 | assert(argc == 3); 50 | 51 | auto yolov8 = new YOLOv8(engine_file_path); 52 | yolov8->make_pipe(true); 53 | 54 | if (IsFile(path)) { 55 | std::string suffix = path.substr(path.find_last_of('.') + 1); 56 | if (suffix == "jpg" || suffix == "jpeg" || suffix == "png") { 57 | imagePathList.push_back(path); 58 | } 59 | else if (suffix == "mp4" || suffix == "avi" || suffix == "m4v" || suffix == "mpeg" || suffix == "mov" 60 | || suffix == "mkv") { 61 | isVideo = true; 62 | } 63 | else { 64 | printf("suffix %s is wrong !!!\n", suffix.c_str()); 65 | std::abort(); 66 | } 67 | } 68 | else if (IsFolder(path)) { 69 | cv::glob(path + "/*.jpg", imagePathList); 70 | } 71 | 72 | cv::Mat res, image; 73 | cv::Size size = cv::Size{640, 640}; 74 | std::vector objs; 75 | 76 | cv::namedWindow("result", cv::WINDOW_AUTOSIZE); 77 | 78 | if (isVideo) { 79 | cv::VideoCapture cap(path); 80 | 81 | if (!cap.isOpened()) { 82 | printf("can not open %s\n", path.c_str()); 83 | return -1; 84 | } 85 | while (cap.read(image)) { 86 | objs.clear(); 87 | yolov8->copy_from_Mat(image, size); 88 | auto start = std::chrono::system_clock::now(); 89 | yolov8->infer(); 90 | auto end = std::chrono::system_clock::now(); 91 | yolov8->postprocess(objs); 92 | yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS); 93 | auto tc = (double)std::chrono::duration_cast(end - start).count() / 1000.; 94 | printf("cost %2.4lf ms\n", tc); 95 | cv::imshow("result", res); 96 | if (cv::waitKey(10) == 'q') { 97 | break; 98 | } 99 | } 100 | } 101 | else { 102 | for (auto& path : imagePathList) { 103 | objs.clear(); 104 | image = cv::imread(path); 105 | yolov8->copy_from_Mat(image, size); 106 | auto start = std::chrono::system_clock::now(); 107 | yolov8->infer(); 108 | auto end = std::chrono::system_clock::now(); 109 | yolov8->postprocess(objs); 110 | yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS); 111 | auto tc = (double)std::chrono::duration_cast(end - start).count() / 1000.; 112 | printf("cost %2.4lf ms\n", tc); 113 | cv::imshow("result", res); 114 | cv::waitKey(0); 115 | } 116 | } 117 | cv::destroyAllWindows(); 118 | delete yolov8; 119 | return 0; 120 | } 121 | -------------------------------------------------------------------------------- /src/detect/normal/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.1) 2 | 3 | set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86 89 90) 4 | set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc) 5 | 6 | project(yolov8 LANGUAGES CXX CUDA) 7 | 8 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -O3") 9 | set(CMAKE_CXX_STANDARD 14) 10 | set(CMAKE_BUILD_TYPE Release) 11 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 12 | 13 | # CUDA 14 | find_package(CUDA REQUIRED) 15 | message(STATUS "CUDA Libs: \n${CUDA_LIBRARIES}\n") 16 | get_filename_component(CUDA_LIB_DIR ${CUDA_LIBRARIES} DIRECTORY) 17 | message(STATUS "CUDA Headers: \n${CUDA_INCLUDE_DIRS}\n") 18 | 19 | # OpenCV 20 | find_package(OpenCV REQUIRED) 21 | message(STATUS "OpenCV Libs: \n${OpenCV_LIBS}\n") 22 | message(STATUS "OpenCV Libraries: \n${OpenCV_LIBRARIES}\n") 23 | message(STATUS "OpenCV Headers: \n${OpenCV_INCLUDE_DIRS}\n") 24 | 25 | # TensorRT 26 | set(TensorRT_INCLUDE_DIRS /usr/include/x86_64-linux-gnu) 27 | set(TensorRT_LIBRARIES /usr/lib/x86_64-linux-gnu) 28 | 29 | 30 | message(STATUS "TensorRT Libs: \n${TensorRT_LIBRARIES}\n") 31 | message(STATUS "TensorRT Headers: \n${TensorRT_INCLUDE_DIRS}\n") 32 | 33 | list(APPEND INCLUDE_DIRS 34 | ${CUDA_INCLUDE_DIRS} 35 | ${OpenCV_INCLUDE_DIRS} 36 | ${TensorRT_INCLUDE_DIRS} 37 | include 38 | ) 39 | 40 | list(APPEND ALL_LIBS 41 | ${CUDA_LIBRARIES} 42 | ${CUDA_LIB_DIR} 43 | ${OpenCV_LIBRARIES} 44 | ${TensorRT_LIBRARIES} 45 | ) 46 | 47 | include_directories(${INCLUDE_DIRS}) 48 | 49 | add_executable(${PROJECT_NAME} 50 | main.cpp 51 | include/yolov8.hpp 52 | include/common.hpp 53 | ) 54 | 55 | target_link_directories(${PROJECT_NAME} PUBLIC ${ALL_LIBS}) 56 | target_link_libraries(${PROJECT_NAME} PRIVATE nvinfer nvinfer_plugin cudart ${OpenCV_LIBS}) 57 | 58 | if (${OpenCV_VERSION} VERSION_GREATER_EQUAL 4.7.0) 59 | message(STATUS "Build with -DBATCHED_NMS") 60 | add_definitions(-DBATCHED_NMS) 61 | endif () 62 | -------------------------------------------------------------------------------- /src/detect/normal/include/common.hpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 1/24/23. 3 | // 4 | 5 | #ifndef DETECT_NORMAL_COMMON_HPP 6 | #define DETECT_NORMAL_COMMON_HPP 7 | #include "NvInfer.h" 8 | #include "opencv2/opencv.hpp" 9 | #include 10 | #include 11 | 12 | #define CHECK(call) \ 13 | do { \ 14 | const cudaError_t error_code = call; \ 15 | if (error_code != cudaSuccess) { \ 16 | printf("CUDA Error:\n"); \ 17 | printf(" File: %s\n", __FILE__); \ 18 | printf(" Line: %d\n", __LINE__); \ 19 | printf(" Error code: %d\n", error_code); \ 20 | printf(" Error text: %s\n", cudaGetErrorString(error_code)); \ 21 | exit(1); \ 22 | } \ 23 | } while (0) 24 | 25 | class Logger: public nvinfer1::ILogger { 26 | public: 27 | nvinfer1::ILogger::Severity reportableSeverity; 28 | 29 | explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO): 30 | reportableSeverity(severity) 31 | { 32 | } 33 | 34 | void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override 35 | { 36 | if (severity > reportableSeverity) { 37 | return; 38 | } 39 | switch (severity) { 40 | case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: 41 | std::cerr << "INTERNAL_ERROR: "; 42 | break; 43 | case nvinfer1::ILogger::Severity::kERROR: 44 | std::cerr << "ERROR: "; 45 | break; 46 | case nvinfer1::ILogger::Severity::kWARNING: 47 | std::cerr << "WARNING: "; 48 | break; 49 | case nvinfer1::ILogger::Severity::kINFO: 50 | std::cerr << "INFO: "; 51 | break; 52 | default: 53 | std::cerr << "VERBOSE: "; 54 | break; 55 | } 56 | std::cerr << msg << std::endl; 57 | } 58 | }; 59 | 60 | inline int get_size_by_dims(const nvinfer1::Dims& dims) 61 | { 62 | int size = 1; 63 | for (int i = 0; i < dims.nbDims; i++) { 64 | size *= dims.d[i]; 65 | } 66 | return size; 67 | } 68 | 69 | inline int type_to_size(const nvinfer1::DataType& dataType) 70 | { 71 | switch (dataType) { 72 | case nvinfer1::DataType::kFLOAT: 73 | return 4; 74 | case nvinfer1::DataType::kHALF: 75 | return 2; 76 | case nvinfer1::DataType::kINT32: 77 | return 4; 78 | case nvinfer1::DataType::kINT8: 79 | return 1; 80 | case nvinfer1::DataType::kBOOL: 81 | return 1; 82 | default: 83 | return 4; 84 | } 85 | } 86 | 87 | inline static float clamp(float val, float min, float max) 88 | { 89 | return val > min ? (val < max ? val : max) : min; 90 | } 91 | 92 | inline bool IsPathExist(const std::string& path) 93 | { 94 | if (access(path.c_str(), 0) == F_OK) { 95 | return true; 96 | } 97 | return false; 98 | } 99 | 100 | inline bool IsFile(const std::string& path) 101 | { 102 | if (!IsPathExist(path)) { 103 | printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str()); 104 | return false; 105 | } 106 | struct stat buffer; 107 | return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode)); 108 | } 109 | 110 | inline bool IsFolder(const std::string& path) 111 | { 112 | if (!IsPathExist(path)) { 113 | return false; 114 | } 115 | struct stat buffer; 116 | return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode)); 117 | } 118 | 119 | namespace det { 120 | struct Binding { 121 | size_t size = 1; 122 | size_t dsize = 1; 123 | nvinfer1::Dims dims; 124 | std::string name; 125 | }; 126 | 127 | struct Object { 128 | cv::Rect_ rect; 129 | int label = 0; 130 | float prob = 0.0; 131 | }; 132 | 133 | struct PreParam { 134 | float ratio = 1.0f; 135 | float dw = 0.0f; 136 | float dh = 0.0f; 137 | float height = 0; 138 | float width = 0; 139 | }; 140 | } // namespace det 141 | #endif // DETECT_NORMAL_COMMON_HPP 142 | -------------------------------------------------------------------------------- /src/detect/normal/main.cpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 1/20/23. 3 | // 4 | #include "chrono" 5 | #include "opencv2/opencv.hpp" 6 | #include "yolov8.hpp" 7 | 8 | const std::vector CLASS_NAMES = { 9 | "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", 10 | "truck", "boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", 11 | "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", 12 | "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", 13 | "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat", 14 | "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup", 15 | "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", 16 | "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", 17 | "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", 18 | "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven", 19 | "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", 20 | "teddy bear", "hair drier", "toothbrush"}; 21 | 22 | const std::vector> COLORS = { 23 | {0, 114, 189}, {217, 83, 25}, {237, 177, 32}, {126, 47, 142}, {119, 172, 48}, {77, 190, 238}, 24 | {162, 20, 47}, {76, 76, 76}, {153, 153, 153}, {255, 0, 0}, {255, 128, 0}, {191, 191, 0}, 25 | {0, 255, 0}, {0, 0, 255}, {170, 0, 255}, {85, 85, 0}, {85, 170, 0}, {85, 255, 0}, 26 | {170, 85, 0}, {170, 170, 0}, {170, 255, 0}, {255, 85, 0}, {255, 170, 0}, {255, 255, 0}, 27 | {0, 85, 128}, {0, 170, 128}, {0, 255, 128}, {85, 0, 128}, {85, 85, 128}, {85, 170, 128}, 28 | {85, 255, 128}, {170, 0, 128}, {170, 85, 128}, {170, 170, 128}, {170, 255, 128}, {255, 0, 128}, 29 | {255, 85, 128}, {255, 170, 128}, {255, 255, 128}, {0, 85, 255}, {0, 170, 255}, {0, 255, 255}, 30 | {85, 0, 255}, {85, 85, 255}, {85, 170, 255}, {85, 255, 255}, {170, 0, 255}, {170, 85, 255}, 31 | {170, 170, 255}, {170, 255, 255}, {255, 0, 255}, {255, 85, 255}, {255, 170, 255}, {85, 0, 0}, 32 | {128, 0, 0}, {170, 0, 0}, {212, 0, 0}, {255, 0, 0}, {0, 43, 0}, {0, 85, 0}, 33 | {0, 128, 0}, {0, 170, 0}, {0, 212, 0}, {0, 255, 0}, {0, 0, 43}, {0, 0, 85}, 34 | {0, 0, 128}, {0, 0, 170}, {0, 0, 212}, {0, 0, 255}, {0, 0, 0}, {36, 36, 36}, 35 | {73, 73, 73}, {109, 109, 109}, {146, 146, 146}, {182, 182, 182}, {219, 219, 219}, {0, 114, 189}, 36 | {80, 183, 189}, {128, 128, 0}}; 37 | 38 | int main(int argc, char** argv) 39 | { 40 | // cuda:0 41 | cudaSetDevice(0); 42 | 43 | const std::string engine_file_path{argv[1]}; 44 | const std::string path{argv[2]}; 45 | 46 | std::vector imagePathList; 47 | bool isVideo{false}; 48 | 49 | assert(argc == 3); 50 | 51 | auto yolov8 = new YOLOv8(engine_file_path); 52 | yolov8->make_pipe(true); 53 | 54 | if (IsFile(path)) { 55 | std::string suffix = path.substr(path.find_last_of('.') + 1); 56 | if (suffix == "jpg" || suffix == "jpeg" || suffix == "png") { 57 | imagePathList.push_back(path); 58 | } 59 | else if (suffix == "mp4" || suffix == "avi" || suffix == "m4v" || suffix == "mpeg" || suffix == "mov" 60 | || suffix == "mkv") { 61 | isVideo = true; 62 | } 63 | else { 64 | printf("suffix %s is wrong !!!\n", suffix.c_str()); 65 | std::abort(); 66 | } 67 | } 68 | else if (IsFolder(path)) { 69 | cv::glob(path + "/*.jpg", imagePathList); 70 | } 71 | 72 | cv::Mat res, image; 73 | cv::Size size = cv::Size{640, 640}; 74 | int num_labels = 80; 75 | int topk = 100; 76 | float score_thres = 0.25f; 77 | float iou_thres = 0.65f; 78 | 79 | std::vector objs; 80 | 81 | cv::namedWindow("result", cv::WINDOW_AUTOSIZE); 82 | 83 | if (isVideo) { 84 | cv::VideoCapture cap(path); 85 | 86 | if (!cap.isOpened()) { 87 | printf("can not open %s\n", path.c_str()); 88 | return -1; 89 | } 90 | while (cap.read(image)) { 91 | objs.clear(); 92 | yolov8->copy_from_Mat(image, size); 93 | auto start = std::chrono::system_clock::now(); 94 | yolov8->infer(); 95 | auto end = std::chrono::system_clock::now(); 96 | yolov8->postprocess(objs, score_thres, iou_thres, topk, num_labels); 97 | yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS); 98 | auto tc = (double)std::chrono::duration_cast(end - start).count() / 1000.; 99 | printf("cost %2.4lf ms\n", tc); 100 | cv::imshow("result", res); 101 | if (cv::waitKey(10) == 'q') { 102 | break; 103 | } 104 | } 105 | } 106 | else { 107 | for (auto& path : imagePathList) { 108 | objs.clear(); 109 | image = cv::imread(path); 110 | yolov8->copy_from_Mat(image, size); 111 | auto start = std::chrono::system_clock::now(); 112 | yolov8->infer(); 113 | auto end = std::chrono::system_clock::now(); 114 | yolov8->postprocess(objs, score_thres, iou_thres, topk, num_labels); 115 | yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS); 116 | auto tc = (double)std::chrono::duration_cast(end - start).count() / 1000.; 117 | printf("cost %2.4lf ms\n", tc); 118 | cv::imshow("result", res); 119 | cv::waitKey(0); 120 | } 121 | } 122 | cv::destroyAllWindows(); 123 | delete yolov8; 124 | return 0; 125 | } 126 | -------------------------------------------------------------------------------- /src/jetson/detect/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.1) 2 | 3 | set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86) 4 | set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc) 5 | 6 | project(yolov8 LANGUAGES CXX CUDA) 7 | 8 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -O3") 9 | set(CMAKE_CXX_STANDARD 14) 10 | set(CMAKE_BUILD_TYPE Release) 11 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 12 | 13 | # CUDA 14 | find_package(CUDA REQUIRED) 15 | message(STATUS "CUDA Libs: \n${CUDA_LIBRARIES}\n") 16 | get_filename_component(CUDA_LIB_DIR ${CUDA_LIBRARIES} DIRECTORY) 17 | message(STATUS "CUDA Headers: \n${CUDA_INCLUDE_DIRS}\n") 18 | 19 | # OpenCV 20 | find_package(OpenCV REQUIRED) 21 | message(STATUS "OpenCV Libs: \n${OpenCV_LIBS}\n") 22 | message(STATUS "OpenCV Libraries: \n${OpenCV_LIBRARIES}\n") 23 | message(STATUS "OpenCV Headers: \n${OpenCV_INCLUDE_DIRS}\n") 24 | 25 | # TensorRT 26 | set(TensorRT_INCLUDE_DIRS /usr/include/aarch64-linux-gnu) 27 | set(TensorRT_LIBRARIES /usr/lib/aarch64-linux-gnu) 28 | 29 | 30 | message(STATUS "TensorRT Libs: \n${TensorRT_LIBRARIES}\n") 31 | message(STATUS "TensorRT Headers: \n${TensorRT_INCLUDE_DIRS}\n") 32 | 33 | list(APPEND INCLUDE_DIRS 34 | ${CUDA_INCLUDE_DIRS} 35 | ${OpenCV_INCLUDE_DIRS} 36 | ${TensorRT_INCLUDE_DIRS} 37 | include 38 | ) 39 | 40 | list(APPEND ALL_LIBS 41 | ${CUDA_LIBRARIES} 42 | ${CUDA_LIB_DIR} 43 | ${OpenCV_LIBRARIES} 44 | ${TensorRT_LIBRARIES} 45 | ) 46 | 47 | include_directories(${INCLUDE_DIRS}) 48 | 49 | add_executable(${PROJECT_NAME} 50 | main.cpp 51 | include/yolov8.hpp 52 | include/common.hpp 53 | ) 54 | 55 | link_directories(${ALL_LIBS}) 56 | target_link_libraries(${PROJECT_NAME} PRIVATE nvinfer nvinfer_plugin ${CUDA_LIBRARIES} ${OpenCV_LIBS}) 57 | -------------------------------------------------------------------------------- /src/jetson/detect/include/common.hpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 3/16/23. 3 | // 4 | 5 | #ifndef JETSON_DETECT_COMMON_HPP 6 | #define JETSON_DETECT_COMMON_HPP 7 | #include "NvInfer.h" 8 | #include "opencv2/opencv.hpp" 9 | #include 10 | #include 11 | 12 | #define CHECK(call) \ 13 | do { \ 14 | const cudaError_t error_code = call; \ 15 | if (error_code != cudaSuccess) { \ 16 | printf("CUDA Error:\n"); \ 17 | printf(" File: %s\n", __FILE__); \ 18 | printf(" Line: %d\n", __LINE__); \ 19 | printf(" Error code: %d\n", error_code); \ 20 | printf(" Error text: %s\n", cudaGetErrorString(error_code)); \ 21 | exit(1); \ 22 | } \ 23 | } while (0) 24 | 25 | class Logger: public nvinfer1::ILogger { 26 | public: 27 | nvinfer1::ILogger::Severity reportableSeverity; 28 | 29 | explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO): 30 | reportableSeverity(severity) 31 | { 32 | } 33 | 34 | void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override 35 | { 36 | if (severity > reportableSeverity) { 37 | return; 38 | } 39 | switch (severity) { 40 | case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: 41 | std::cerr << "INTERNAL_ERROR: "; 42 | break; 43 | case nvinfer1::ILogger::Severity::kERROR: 44 | std::cerr << "ERROR: "; 45 | break; 46 | case nvinfer1::ILogger::Severity::kWARNING: 47 | std::cerr << "WARNING: "; 48 | break; 49 | case nvinfer1::ILogger::Severity::kINFO: 50 | std::cerr << "INFO: "; 51 | break; 52 | default: 53 | std::cerr << "VERBOSE: "; 54 | break; 55 | } 56 | std::cerr << msg << std::endl; 57 | } 58 | }; 59 | 60 | inline int get_size_by_dims(const nvinfer1::Dims& dims) 61 | { 62 | int size = 1; 63 | for (int i = 0; i < dims.nbDims; i++) { 64 | size *= dims.d[i]; 65 | } 66 | return size; 67 | } 68 | 69 | inline int type_to_size(const nvinfer1::DataType& dataType) 70 | { 71 | switch (dataType) { 72 | case nvinfer1::DataType::kFLOAT: 73 | return 4; 74 | case nvinfer1::DataType::kHALF: 75 | return 2; 76 | case nvinfer1::DataType::kINT32: 77 | return 4; 78 | case nvinfer1::DataType::kINT8: 79 | return 1; 80 | case nvinfer1::DataType::kBOOL: 81 | return 1; 82 | default: 83 | return 4; 84 | } 85 | } 86 | 87 | inline static float clamp(float val, float min, float max) 88 | { 89 | return val > min ? (val < max ? val : max) : min; 90 | } 91 | 92 | inline bool IsPathExist(const std::string& path) 93 | { 94 | if (access(path.c_str(), 0) == F_OK) { 95 | return true; 96 | } 97 | return false; 98 | } 99 | 100 | inline bool IsFile(const std::string& path) 101 | { 102 | if (!IsPathExist(path)) { 103 | printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str()); 104 | return false; 105 | } 106 | struct stat buffer; 107 | return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode)); 108 | } 109 | 110 | inline bool IsFolder(const std::string& path) 111 | { 112 | if (!IsPathExist(path)) { 113 | return false; 114 | } 115 | struct stat buffer; 116 | return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode)); 117 | } 118 | 119 | namespace det { 120 | struct Binding { 121 | size_t size = 1; 122 | size_t dsize = 1; 123 | nvinfer1::Dims dims; 124 | std::string name; 125 | }; 126 | 127 | struct Object { 128 | cv::Rect_ rect; 129 | int label = 0; 130 | float prob = 0.0; 131 | }; 132 | 133 | struct PreParam { 134 | float ratio = 1.0f; 135 | float dw = 0.0f; 136 | float dh = 0.0f; 137 | float height = 0; 138 | float width = 0; 139 | }; 140 | } // namespace det 141 | #endif // JETSON_DETECT_COMMON_HPP 142 | -------------------------------------------------------------------------------- /src/jetson/detect/main.cpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 3/16/23. 3 | // 4 | #include "chrono" 5 | #include "opencv2/opencv.hpp" 6 | #include "yolov8.hpp" 7 | 8 | const std::vector CLASS_NAMES = { 9 | "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", 10 | "truck", "boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", 11 | "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", 12 | "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", 13 | "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat", 14 | "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup", 15 | "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", 16 | "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", 17 | "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", 18 | "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven", 19 | "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", 20 | "teddy bear", "hair drier", "toothbrush"}; 21 | 22 | const std::vector> COLORS = { 23 | {0, 114, 189}, {217, 83, 25}, {237, 177, 32}, {126, 47, 142}, {119, 172, 48}, {77, 190, 238}, 24 | {162, 20, 47}, {76, 76, 76}, {153, 153, 153}, {255, 0, 0}, {255, 128, 0}, {191, 191, 0}, 25 | {0, 255, 0}, {0, 0, 255}, {170, 0, 255}, {85, 85, 0}, {85, 170, 0}, {85, 255, 0}, 26 | {170, 85, 0}, {170, 170, 0}, {170, 255, 0}, {255, 85, 0}, {255, 170, 0}, {255, 255, 0}, 27 | {0, 85, 128}, {0, 170, 128}, {0, 255, 128}, {85, 0, 128}, {85, 85, 128}, {85, 170, 128}, 28 | {85, 255, 128}, {170, 0, 128}, {170, 85, 128}, {170, 170, 128}, {170, 255, 128}, {255, 0, 128}, 29 | {255, 85, 128}, {255, 170, 128}, {255, 255, 128}, {0, 85, 255}, {0, 170, 255}, {0, 255, 255}, 30 | {85, 0, 255}, {85, 85, 255}, {85, 170, 255}, {85, 255, 255}, {170, 0, 255}, {170, 85, 255}, 31 | {170, 170, 255}, {170, 255, 255}, {255, 0, 255}, {255, 85, 255}, {255, 170, 255}, {85, 0, 0}, 32 | {128, 0, 0}, {170, 0, 0}, {212, 0, 0}, {255, 0, 0}, {0, 43, 0}, {0, 85, 0}, 33 | {0, 128, 0}, {0, 170, 0}, {0, 212, 0}, {0, 255, 0}, {0, 0, 43}, {0, 0, 85}, 34 | {0, 0, 128}, {0, 0, 170}, {0, 0, 212}, {0, 0, 255}, {0, 0, 0}, {36, 36, 36}, 35 | {73, 73, 73}, {109, 109, 109}, {146, 146, 146}, {182, 182, 182}, {219, 219, 219}, {0, 114, 189}, 36 | {80, 183, 189}, {128, 128, 0}}; 37 | 38 | int main(int argc, char** argv) 39 | { 40 | const std::string engine_file_path{argv[1]}; 41 | const std::string path{argv[2]}; 42 | 43 | std::vector imagePathList; 44 | bool isVideo{false}; 45 | 46 | assert(argc == 3); 47 | 48 | auto yolov8 = new YOLOv8(engine_file_path); 49 | yolov8->make_pipe(true); 50 | 51 | if (IsFile(path)) { 52 | std::string suffix = path.substr(path.find_last_of('.') + 1); 53 | if (suffix == "jpg" || suffix == "jpeg" || suffix == "png") { 54 | imagePathList.push_back(path); 55 | } 56 | else if (suffix == "mp4" || suffix == "avi" || suffix == "m4v" || suffix == "mpeg" || suffix == "mov" 57 | || suffix == "mkv") { 58 | isVideo = true; 59 | } 60 | else { 61 | printf("suffix %s is wrong !!!\n", suffix.c_str()); 62 | std::abort(); 63 | } 64 | } 65 | else if (IsFolder(path)) { 66 | cv::glob(path + "/*.jpg", imagePathList); 67 | } 68 | 69 | cv::Mat res, image; 70 | cv::Size size = cv::Size{640, 640}; 71 | std::vector objs; 72 | 73 | cv::namedWindow("result", cv::WINDOW_AUTOSIZE); 74 | 75 | if (isVideo) { 76 | cv::VideoCapture cap(path); 77 | 78 | if (!cap.isOpened()) { 79 | printf("can not open %s\n", path.c_str()); 80 | return -1; 81 | } 82 | while (cap.read(image)) { 83 | objs.clear(); 84 | yolov8->copy_from_Mat(image, size); 85 | auto start = std::chrono::system_clock::now(); 86 | yolov8->infer(); 87 | auto end = std::chrono::system_clock::now(); 88 | yolov8->postprocess(objs); 89 | yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS); 90 | auto tc = (double)std::chrono::duration_cast(end - start).count() / 1000.; 91 | printf("cost %2.4lf ms\n", tc); 92 | cv::imshow("result", res); 93 | if (cv::waitKey(10) == 'q') { 94 | break; 95 | } 96 | } 97 | } 98 | else { 99 | for (auto& path : imagePathList) { 100 | objs.clear(); 101 | image = cv::imread(path); 102 | yolov8->copy_from_Mat(image, size); 103 | auto start = std::chrono::system_clock::now(); 104 | yolov8->infer(); 105 | auto end = std::chrono::system_clock::now(); 106 | yolov8->postprocess(objs); 107 | yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS); 108 | auto tc = (double)std::chrono::duration_cast(end - start).count() / 1000.; 109 | printf("cost %2.4lf ms\n", tc); 110 | cv::imshow("result", res); 111 | cv::waitKey(0); 112 | } 113 | } 114 | cv::destroyAllWindows(); 115 | delete yolov8; 116 | return 0; 117 | } 118 | -------------------------------------------------------------------------------- /src/jetson/pose/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.1) 2 | 3 | set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86) 4 | set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc) 5 | 6 | project(yolov8-pose LANGUAGES CXX CUDA) 7 | 8 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -O3") 9 | set(CMAKE_CXX_STANDARD 14) 10 | set(CMAKE_BUILD_TYPE Release) 11 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 12 | 13 | # CUDA 14 | find_package(CUDA REQUIRED) 15 | message(STATUS "CUDA Libs: \n${CUDA_LIBRARIES}\n") 16 | get_filename_component(CUDA_LIB_DIR ${CUDA_LIBRARIES} DIRECTORY) 17 | message(STATUS "CUDA Headers: \n${CUDA_INCLUDE_DIRS}\n") 18 | 19 | # OpenCV 20 | find_package(OpenCV REQUIRED) 21 | message(STATUS "OpenCV Libs: \n${OpenCV_LIBS}\n") 22 | message(STATUS "OpenCV Libraries: \n${OpenCV_LIBRARIES}\n") 23 | message(STATUS "OpenCV Headers: \n${OpenCV_INCLUDE_DIRS}\n") 24 | 25 | # TensorRT 26 | set(TensorRT_INCLUDE_DIRS /usr/include/aarch64-linux-gnu) 27 | set(TensorRT_LIBRARIES /usr/lib/aarch64-linux-gnu) 28 | 29 | 30 | message(STATUS "TensorRT Libs: \n${TensorRT_LIBRARIES}\n") 31 | message(STATUS "TensorRT Headers: \n${TensorRT_INCLUDE_DIRS}\n") 32 | 33 | list(APPEND INCLUDE_DIRS 34 | ${CUDA_INCLUDE_DIRS} 35 | ${OpenCV_INCLUDE_DIRS} 36 | ${TensorRT_INCLUDE_DIRS} 37 | include 38 | ) 39 | 40 | list(APPEND ALL_LIBS 41 | ${CUDA_LIBRARIES} 42 | ${CUDA_LIB_DIR} 43 | ${OpenCV_LIBRARIES} 44 | ${TensorRT_LIBRARIES} 45 | ) 46 | 47 | include_directories(${INCLUDE_DIRS}) 48 | 49 | add_executable(${PROJECT_NAME} 50 | main.cpp 51 | include/yolov8-pose.hpp 52 | include/common.hpp 53 | ) 54 | 55 | link_directories(${ALL_LIBS}) 56 | target_link_libraries(${PROJECT_NAME} PRIVATE nvinfer nvinfer_plugin ${CUDA_LIBRARIES} ${OpenCV_LIBS}) 57 | 58 | 59 | if(${OpenCV_VERSION} VERSION_GREATER_EQUAL 4.7.0) 60 | message(STATUS "Build with -DBATCHED_NMS") 61 | add_definitions(-DBATCHED_NMS) 62 | endif() 63 | -------------------------------------------------------------------------------- /src/jetson/pose/include/common.hpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 5/15/23. 3 | // 4 | 5 | #ifndef JETSON_POSE_COMMON_HPP 6 | #define JETSON_POSE_COMMON_HPP 7 | #include "NvInfer.h" 8 | #include "opencv2/opencv.hpp" 9 | #include 10 | #include 11 | 12 | #define CHECK(call) \ 13 | do { \ 14 | const cudaError_t error_code = call; \ 15 | if (error_code != cudaSuccess) { \ 16 | printf("CUDA Error:\n"); \ 17 | printf(" File: %s\n", __FILE__); \ 18 | printf(" Line: %d\n", __LINE__); \ 19 | printf(" Error code: %d\n", error_code); \ 20 | printf(" Error text: %s\n", cudaGetErrorString(error_code)); \ 21 | exit(1); \ 22 | } \ 23 | } while (0) 24 | 25 | class Logger: public nvinfer1::ILogger { 26 | public: 27 | nvinfer1::ILogger::Severity reportableSeverity; 28 | 29 | explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO): 30 | reportableSeverity(severity) 31 | { 32 | } 33 | 34 | void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override 35 | { 36 | if (severity > reportableSeverity) { 37 | return; 38 | } 39 | switch (severity) { 40 | case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: 41 | std::cerr << "INTERNAL_ERROR: "; 42 | break; 43 | case nvinfer1::ILogger::Severity::kERROR: 44 | std::cerr << "ERROR: "; 45 | break; 46 | case nvinfer1::ILogger::Severity::kWARNING: 47 | std::cerr << "WARNING: "; 48 | break; 49 | case nvinfer1::ILogger::Severity::kINFO: 50 | std::cerr << "INFO: "; 51 | break; 52 | default: 53 | std::cerr << "VERBOSE: "; 54 | break; 55 | } 56 | std::cerr << msg << std::endl; 57 | } 58 | }; 59 | 60 | inline int get_size_by_dims(const nvinfer1::Dims& dims) 61 | { 62 | int size = 1; 63 | for (int i = 0; i < dims.nbDims; i++) { 64 | size *= dims.d[i]; 65 | } 66 | return size; 67 | } 68 | 69 | inline int type_to_size(const nvinfer1::DataType& dataType) 70 | { 71 | switch (dataType) { 72 | case nvinfer1::DataType::kFLOAT: 73 | return 4; 74 | case nvinfer1::DataType::kHALF: 75 | return 2; 76 | case nvinfer1::DataType::kINT32: 77 | return 4; 78 | case nvinfer1::DataType::kINT8: 79 | return 1; 80 | case nvinfer1::DataType::kBOOL: 81 | return 1; 82 | default: 83 | return 4; 84 | } 85 | } 86 | 87 | inline static float clamp(float val, float min, float max) 88 | { 89 | return val > min ? (val < max ? val : max) : min; 90 | } 91 | 92 | inline bool IsPathExist(const std::string& path) 93 | { 94 | if (access(path.c_str(), 0) == F_OK) { 95 | return true; 96 | } 97 | return false; 98 | } 99 | 100 | inline bool IsFile(const std::string& path) 101 | { 102 | if (!IsPathExist(path)) { 103 | printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str()); 104 | return false; 105 | } 106 | struct stat buffer; 107 | return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode)); 108 | } 109 | 110 | inline bool IsFolder(const std::string& path) 111 | { 112 | if (!IsPathExist(path)) { 113 | return false; 114 | } 115 | struct stat buffer; 116 | return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode)); 117 | } 118 | 119 | namespace pose { 120 | struct Binding { 121 | size_t size = 1; 122 | size_t dsize = 1; 123 | nvinfer1::Dims dims; 124 | std::string name; 125 | }; 126 | 127 | struct Object { 128 | cv::Rect_ rect; 129 | int label = 0; 130 | float prob = 0.0; 131 | std::vector kps; 132 | }; 133 | 134 | struct PreParam { 135 | float ratio = 1.0f; 136 | float dw = 0.0f; 137 | float dh = 0.0f; 138 | float height = 0; 139 | float width = 0; 140 | }; 141 | } // namespace pose 142 | #endif // JETSON_POSE_COMMON_HPP 143 | -------------------------------------------------------------------------------- /src/jetson/pose/main.cpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 4/7/23. 3 | // 4 | #include "chrono" 5 | #include "opencv2/opencv.hpp" 6 | #include "yolov8-pose.hpp" 7 | 8 | const std::vector> KPS_COLORS = {{0, 255, 0}, 9 | {0, 255, 0}, 10 | {0, 255, 0}, 11 | {0, 255, 0}, 12 | {0, 255, 0}, 13 | {255, 128, 0}, 14 | {255, 128, 0}, 15 | {255, 128, 0}, 16 | {255, 128, 0}, 17 | {255, 128, 0}, 18 | {255, 128, 0}, 19 | {51, 153, 255}, 20 | {51, 153, 255}, 21 | {51, 153, 255}, 22 | {51, 153, 255}, 23 | {51, 153, 255}, 24 | {51, 153, 255}}; 25 | 26 | const std::vector> SKELETON = {{16, 14}, 27 | {14, 12}, 28 | {17, 15}, 29 | {15, 13}, 30 | {12, 13}, 31 | {6, 12}, 32 | {7, 13}, 33 | {6, 7}, 34 | {6, 8}, 35 | {7, 9}, 36 | {8, 10}, 37 | {9, 11}, 38 | {2, 3}, 39 | {1, 2}, 40 | {1, 3}, 41 | {2, 4}, 42 | {3, 5}, 43 | {4, 6}, 44 | {5, 7}}; 45 | 46 | const std::vector> LIMB_COLORS = {{51, 153, 255}, 47 | {51, 153, 255}, 48 | {51, 153, 255}, 49 | {51, 153, 255}, 50 | {255, 51, 255}, 51 | {255, 51, 255}, 52 | {255, 51, 255}, 53 | {255, 128, 0}, 54 | {255, 128, 0}, 55 | {255, 128, 0}, 56 | {255, 128, 0}, 57 | {255, 128, 0}, 58 | {0, 255, 0}, 59 | {0, 255, 0}, 60 | {0, 255, 0}, 61 | {0, 255, 0}, 62 | {0, 255, 0}, 63 | {0, 255, 0}, 64 | {0, 255, 0}}; 65 | 66 | int main(int argc, char** argv) 67 | { 68 | // cuda:0 69 | cudaSetDevice(0); 70 | 71 | const std::string engine_file_path{argv[1]}; 72 | const std::string path{argv[2]}; 73 | 74 | std::vector imagePathList; 75 | bool isVideo{false}; 76 | 77 | assert(argc == 3); 78 | 79 | auto yolov8_pose = new YOLOv8_pose(engine_file_path); 80 | yolov8_pose->make_pipe(true); 81 | 82 | if (IsFile(path)) { 83 | std::string suffix = path.substr(path.find_last_of('.') + 1); 84 | if (suffix == "jpg" || suffix == "jpeg" || suffix == "png") { 85 | imagePathList.push_back(path); 86 | } 87 | else if (suffix == "mp4" || suffix == "avi" || suffix == "m4v" || suffix == "mpeg" || suffix == "mov" 88 | || suffix == "mkv") { 89 | isVideo = true; 90 | } 91 | else { 92 | printf("suffix %s is wrong !!!\n", suffix.c_str()); 93 | std::abort(); 94 | } 95 | } 96 | else if (IsFolder(path)) { 97 | cv::glob(path + "/*.jpg", imagePathList); 98 | } 99 | 100 | cv::Mat res, image; 101 | cv::Size size = cv::Size{640, 640}; 102 | int topk = 100; 103 | float score_thres = 0.25f; 104 | float iou_thres = 0.65f; 105 | 106 | std::vector objs; 107 | 108 | cv::namedWindow("result", cv::WINDOW_AUTOSIZE); 109 | 110 | if (isVideo) { 111 | cv::VideoCapture cap(path); 112 | 113 | if (!cap.isOpened()) { 114 | printf("can not open %s\n", path.c_str()); 115 | return -1; 116 | } 117 | while (cap.read(image)) { 118 | objs.clear(); 119 | yolov8_pose->copy_from_Mat(image, size); 120 | auto start = std::chrono::system_clock::now(); 121 | yolov8_pose->infer(); 122 | auto end = std::chrono::system_clock::now(); 123 | yolov8_pose->postprocess(objs, score_thres, iou_thres, topk); 124 | yolov8_pose->draw_objects(image, res, objs, SKELETON, KPS_COLORS, LIMB_COLORS); 125 | auto tc = (double)std::chrono::duration_cast(end - start).count() / 1000.; 126 | printf("cost %2.4lf ms\n", tc); 127 | cv::imshow("result", res); 128 | if (cv::waitKey(10) == 'q') { 129 | break; 130 | } 131 | } 132 | } 133 | else { 134 | for (auto& path : imagePathList) { 135 | objs.clear(); 136 | image = cv::imread(path); 137 | yolov8_pose->copy_from_Mat(image, size); 138 | auto start = std::chrono::system_clock::now(); 139 | yolov8_pose->infer(); 140 | auto end = std::chrono::system_clock::now(); 141 | yolov8_pose->postprocess(objs, score_thres, iou_thres, topk); 142 | yolov8_pose->draw_objects(image, res, objs, SKELETON, KPS_COLORS, LIMB_COLORS); 143 | auto tc = (double)std::chrono::duration_cast(end - start).count() / 1000.; 144 | printf("cost %2.4lf ms\n", tc); 145 | cv::imshow("result", res); 146 | cv::waitKey(0); 147 | } 148 | } 149 | cv::destroyAllWindows(); 150 | delete yolov8_pose; 151 | return 0; 152 | } 153 | -------------------------------------------------------------------------------- /src/jetson/segment/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.1) 2 | 3 | set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86) 4 | set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc) 5 | 6 | project(yolov8-seg LANGUAGES CXX CUDA) 7 | 8 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -O3") 9 | set(CMAKE_CXX_STANDARD 14) 10 | set(CMAKE_BUILD_TYPE Release) 11 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 12 | 13 | # CUDA 14 | find_package(CUDA REQUIRED) 15 | message(STATUS "CUDA Libs: \n${CUDA_LIBRARIES}\n") 16 | get_filename_component(CUDA_LIB_DIR ${CUDA_LIBRARIES} DIRECTORY) 17 | message(STATUS "CUDA Headers: \n${CUDA_INCLUDE_DIRS}\n") 18 | 19 | # OpenCV 20 | find_package(OpenCV REQUIRED) 21 | message(STATUS "OpenCV Libs: \n${OpenCV_LIBS}\n") 22 | message(STATUS "OpenCV Libraries: \n${OpenCV_LIBRARIES}\n") 23 | message(STATUS "OpenCV Headers: \n${OpenCV_INCLUDE_DIRS}\n") 24 | 25 | # TensorRT 26 | set(TensorRT_INCLUDE_DIRS /usr/include/aarch64-linux-gnu) 27 | set(TensorRT_LIBRARIES /usr/lib/aarch64-linux-gnu) 28 | 29 | 30 | message(STATUS "TensorRT Libs: \n${TensorRT_LIBRARIES}\n") 31 | message(STATUS "TensorRT Headers: \n${TensorRT_INCLUDE_DIRS}\n") 32 | 33 | list(APPEND INCLUDE_DIRS 34 | ${CUDA_INCLUDE_DIRS} 35 | ${OpenCV_INCLUDE_DIRS} 36 | ${TensorRT_INCLUDE_DIRS} 37 | include 38 | ) 39 | 40 | list(APPEND ALL_LIBS 41 | ${CUDA_LIBRARIES} 42 | ${CUDA_LIB_DIR} 43 | ${OpenCV_LIBRARIES} 44 | ${TensorRT_LIBRARIES} 45 | ) 46 | 47 | include_directories(${INCLUDE_DIRS}) 48 | 49 | add_executable(${PROJECT_NAME} 50 | main.cpp 51 | include/yolov8-seg.hpp 52 | include/common.hpp 53 | ) 54 | 55 | link_directories(${ALL_LIBS}) 56 | target_link_libraries(${PROJECT_NAME} PRIVATE nvinfer nvinfer_plugin ${CUDA_LIBRARIES} ${OpenCV_LIBS}) 57 | 58 | 59 | if(${OpenCV_VERSION} VERSION_GREATER_EQUAL 4.7.0) 60 | message(STATUS "Build with -DBATCHED_NMS") 61 | add_definitions(-DBATCHED_NMS) 62 | endif() 63 | -------------------------------------------------------------------------------- /src/jetson/segment/include/common.hpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 3/16/23. 3 | // 4 | 5 | #ifndef JETSON_SEGMENT_COMMON_HPP 6 | #define JETSON_SEGMENT_COMMON_HPP 7 | #include "NvInfer.h" 8 | #include "opencv2/opencv.hpp" 9 | #include 10 | #include 11 | 12 | #define CHECK(call) \ 13 | do { \ 14 | const cudaError_t error_code = call; \ 15 | if (error_code != cudaSuccess) { \ 16 | printf("CUDA Error:\n"); \ 17 | printf(" File: %s\n", __FILE__); \ 18 | printf(" Line: %d\n", __LINE__); \ 19 | printf(" Error code: %d\n", error_code); \ 20 | printf(" Error text: %s\n", cudaGetErrorString(error_code)); \ 21 | exit(1); \ 22 | } \ 23 | } while (0) 24 | 25 | class Logger: public nvinfer1::ILogger { 26 | public: 27 | nvinfer1::ILogger::Severity reportableSeverity; 28 | 29 | explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO): 30 | reportableSeverity(severity) 31 | { 32 | } 33 | 34 | void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override 35 | { 36 | if (severity > reportableSeverity) { 37 | return; 38 | } 39 | switch (severity) { 40 | case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: 41 | std::cerr << "INTERNAL_ERROR: "; 42 | break; 43 | case nvinfer1::ILogger::Severity::kERROR: 44 | std::cerr << "ERROR: "; 45 | break; 46 | case nvinfer1::ILogger::Severity::kWARNING: 47 | std::cerr << "WARNING: "; 48 | break; 49 | case nvinfer1::ILogger::Severity::kINFO: 50 | std::cerr << "INFO: "; 51 | break; 52 | default: 53 | std::cerr << "VERBOSE: "; 54 | break; 55 | } 56 | std::cerr << msg << std::endl; 57 | } 58 | }; 59 | 60 | inline int get_size_by_dims(const nvinfer1::Dims& dims) 61 | { 62 | int size = 1; 63 | for (int i = 0; i < dims.nbDims; i++) { 64 | size *= dims.d[i]; 65 | } 66 | return size; 67 | } 68 | 69 | inline int type_to_size(const nvinfer1::DataType& dataType) 70 | { 71 | switch (dataType) { 72 | case nvinfer1::DataType::kFLOAT: 73 | return 4; 74 | case nvinfer1::DataType::kHALF: 75 | return 2; 76 | case nvinfer1::DataType::kINT32: 77 | return 4; 78 | case nvinfer1::DataType::kINT8: 79 | return 1; 80 | case nvinfer1::DataType::kBOOL: 81 | return 1; 82 | default: 83 | return 4; 84 | } 85 | } 86 | 87 | inline static float clamp(float val, float min, float max) 88 | { 89 | return val > min ? (val < max ? val : max) : min; 90 | } 91 | 92 | inline bool IsPathExist(const std::string& path) 93 | { 94 | if (access(path.c_str(), 0) == F_OK) { 95 | return true; 96 | } 97 | return false; 98 | } 99 | 100 | inline bool IsFile(const std::string& path) 101 | { 102 | if (!IsPathExist(path)) { 103 | printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str()); 104 | return false; 105 | } 106 | struct stat buffer; 107 | return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode)); 108 | } 109 | 110 | inline bool IsFolder(const std::string& path) 111 | { 112 | if (!IsPathExist(path)) { 113 | return false; 114 | } 115 | struct stat buffer; 116 | return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode)); 117 | } 118 | 119 | namespace seg { 120 | struct Binding { 121 | size_t size = 1; 122 | size_t dsize = 1; 123 | nvinfer1::Dims dims; 124 | std::string name; 125 | }; 126 | 127 | struct Object { 128 | cv::Rect_ rect; 129 | int label = 0; 130 | float prob = 0.0; 131 | cv::Mat boxMask; 132 | }; 133 | 134 | struct PreParam { 135 | float ratio = 1.0f; 136 | float dw = 0.0f; 137 | float dh = 0.0f; 138 | float height = 0; 139 | float width = 0; 140 | }; 141 | } // namespace seg 142 | #endif // JETSON_SEGMENT_COMMON_HPP 143 | -------------------------------------------------------------------------------- /src/jetson/segment/main.cpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 3/16/23. 3 | // 4 | #include "chrono" 5 | #include "opencv2/opencv.hpp" 6 | #include "yolov8-seg.hpp" 7 | 8 | const std::vector CLASS_NAMES = { 9 | "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", 10 | "truck", "boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", 11 | "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", 12 | "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", 13 | "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat", 14 | "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup", 15 | "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", 16 | "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", 17 | "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", 18 | "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven", 19 | "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", 20 | "teddy bear", "hair drier", "toothbrush"}; 21 | 22 | const std::vector> COLORS = { 23 | {0, 114, 189}, {217, 83, 25}, {237, 177, 32}, {126, 47, 142}, {119, 172, 48}, {77, 190, 238}, 24 | {162, 20, 47}, {76, 76, 76}, {153, 153, 153}, {255, 0, 0}, {255, 128, 0}, {191, 191, 0}, 25 | {0, 255, 0}, {0, 0, 255}, {170, 0, 255}, {85, 85, 0}, {85, 170, 0}, {85, 255, 0}, 26 | {170, 85, 0}, {170, 170, 0}, {170, 255, 0}, {255, 85, 0}, {255, 170, 0}, {255, 255, 0}, 27 | {0, 85, 128}, {0, 170, 128}, {0, 255, 128}, {85, 0, 128}, {85, 85, 128}, {85, 170, 128}, 28 | {85, 255, 128}, {170, 0, 128}, {170, 85, 128}, {170, 170, 128}, {170, 255, 128}, {255, 0, 128}, 29 | {255, 85, 128}, {255, 170, 128}, {255, 255, 128}, {0, 85, 255}, {0, 170, 255}, {0, 255, 255}, 30 | {85, 0, 255}, {85, 85, 255}, {85, 170, 255}, {85, 255, 255}, {170, 0, 255}, {170, 85, 255}, 31 | {170, 170, 255}, {170, 255, 255}, {255, 0, 255}, {255, 85, 255}, {255, 170, 255}, {85, 0, 0}, 32 | {128, 0, 0}, {170, 0, 0}, {212, 0, 0}, {255, 0, 0}, {0, 43, 0}, {0, 85, 0}, 33 | {0, 128, 0}, {0, 170, 0}, {0, 212, 0}, {0, 255, 0}, {0, 0, 43}, {0, 0, 85}, 34 | {0, 0, 128}, {0, 0, 170}, {0, 0, 212}, {0, 0, 255}, {0, 0, 0}, {36, 36, 36}, 35 | {73, 73, 73}, {109, 109, 109}, {146, 146, 146}, {182, 182, 182}, {219, 219, 219}, {0, 114, 189}, 36 | {80, 183, 189}, {128, 128, 0}}; 37 | 38 | const std::vector> MASK_COLORS = { 39 | {255, 56, 56}, {255, 157, 151}, {255, 112, 31}, {255, 178, 29}, {207, 210, 49}, {72, 249, 10}, {146, 204, 23}, 40 | {61, 219, 134}, {26, 147, 52}, {0, 212, 187}, {44, 153, 168}, {0, 194, 255}, {52, 69, 147}, {100, 115, 255}, 41 | {0, 24, 236}, {132, 56, 255}, {82, 0, 133}, {203, 56, 255}, {255, 149, 200}, {255, 55, 199}}; 42 | 43 | int main(int argc, char** argv) 44 | { 45 | // cuda:0 46 | cudaSetDevice(0); 47 | 48 | const std::string engine_file_path{argv[1]}; 49 | const std::string path{argv[2]}; 50 | 51 | std::vector imagePathList; 52 | bool isVideo{false}; 53 | 54 | assert(argc == 3); 55 | 56 | auto yolov8 = new YOLOv8_seg(engine_file_path); 57 | yolov8->make_pipe(true); 58 | 59 | if (IsFile(path)) { 60 | std::string suffix = path.substr(path.find_last_of('.') + 1); 61 | if (suffix == "jpg" || suffix == "jpeg" || suffix == "png") { 62 | imagePathList.push_back(path); 63 | } 64 | else if (suffix == "mp4" || suffix == "avi" || suffix == "m4v" || suffix == "mpeg" || suffix == "mov" 65 | || suffix == "mkv") { 66 | isVideo = true; 67 | } 68 | else { 69 | printf("suffix %s is wrong !!!\n", suffix.c_str()); 70 | std::abort(); 71 | } 72 | } 73 | else if (IsFolder(path)) { 74 | cv::glob(path + "/*.jpg", imagePathList); 75 | } 76 | 77 | cv::Mat res, image; 78 | cv::Size size = cv::Size{640, 640}; 79 | int topk = 100; 80 | int seg_h = 160; 81 | int seg_w = 160; 82 | int seg_channels = 32; 83 | float score_thres = 0.25f; 84 | float iou_thres = 0.65f; 85 | 86 | std::vector objs; 87 | 88 | cv::namedWindow("result", cv::WINDOW_AUTOSIZE); 89 | 90 | if (isVideo) { 91 | cv::VideoCapture cap(path); 92 | 93 | if (!cap.isOpened()) { 94 | printf("can not open %s\n", path.c_str()); 95 | return -1; 96 | } 97 | while (cap.read(image)) { 98 | objs.clear(); 99 | yolov8->copy_from_Mat(image, size); 100 | auto start = std::chrono::system_clock::now(); 101 | yolov8->infer(); 102 | auto end = std::chrono::system_clock::now(); 103 | yolov8->postprocess(objs, score_thres, iou_thres, topk, seg_channels, seg_h, seg_w); 104 | yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS, MASK_COLORS); 105 | auto tc = (double)std::chrono::duration_cast(end - start).count() / 1000.; 106 | printf("cost %2.4lf ms\n", tc); 107 | cv::imshow("result", res); 108 | if (cv::waitKey(10) == 'q') { 109 | break; 110 | } 111 | } 112 | } 113 | else { 114 | for (auto& path : imagePathList) { 115 | objs.clear(); 116 | image = cv::imread(path); 117 | yolov8->copy_from_Mat(image, size); 118 | auto start = std::chrono::system_clock::now(); 119 | yolov8->infer(); 120 | auto end = std::chrono::system_clock::now(); 121 | yolov8->postprocess(objs, score_thres, iou_thres, topk, seg_channels, seg_h, seg_w); 122 | yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS, MASK_COLORS); 123 | auto tc = (double)std::chrono::duration_cast(end - start).count() / 1000.; 124 | printf("cost %2.4lf ms\n", tc); 125 | cv::imshow("result", res); 126 | cv::waitKey(0); 127 | } 128 | } 129 | cv::destroyAllWindows(); 130 | delete yolov8; 131 | return 0; 132 | } 133 | -------------------------------------------------------------------------------- /src/pose/normal/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.1) 2 | 3 | set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86 89 90) 4 | set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc) 5 | 6 | project(yolov8-pose LANGUAGES CXX CUDA) 7 | 8 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -O3") 9 | set(CMAKE_CXX_STANDARD 14) 10 | set(CMAKE_BUILD_TYPE Release) 11 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 12 | 13 | # CUDA 14 | find_package(CUDA REQUIRED) 15 | message(STATUS "CUDA Libs: \n${CUDA_LIBRARIES}\n") 16 | get_filename_component(CUDA_LIB_DIR ${CUDA_LIBRARIES} DIRECTORY) 17 | message(STATUS "CUDA Headers: \n${CUDA_INCLUDE_DIRS}\n") 18 | 19 | # OpenCV 20 | find_package(OpenCV REQUIRED) 21 | message(STATUS "OpenCV Libs: \n${OpenCV_LIBS}\n") 22 | message(STATUS "OpenCV Libraries: \n${OpenCV_LIBRARIES}\n") 23 | message(STATUS "OpenCV Headers: \n${OpenCV_INCLUDE_DIRS}\n") 24 | 25 | # TensorRT 26 | # set(TensorRT_ROOT /root/share/TensorRT-8.5.3.1) 27 | set(TensorRT_ROOT /home/lin/software/TensorRT-8.5.3.1) 28 | set(TensorRT_INCLUDE_DIRS ${TensorRT_ROOT}/include) 29 | set(TensorRT_LIBRARIES ${TensorRT_ROOT}/lib) 30 | 31 | message(STATUS "TensorRT Libs: \n${TensorRT_LIBRARIES}\n") 32 | message(STATUS "TensorRT Headers: \n${TensorRT_INCLUDE_DIRS}\n") 33 | 34 | list(APPEND INCLUDE_DIRS 35 | ${CUDA_INCLUDE_DIRS} 36 | ${OpenCV_INCLUDE_DIRS} 37 | ${TensorRT_INCLUDE_DIRS} 38 | ./include 39 | ) 40 | 41 | list(APPEND ALL_LIBS 42 | ${CUDA_LIBRARIES} 43 | ${CUDA_LIB_DIR} 44 | ${OpenCV_LIBRARIES} 45 | ${TensorRT_LIBRARIES} 46 | ) 47 | 48 | include_directories(${INCLUDE_DIRS}) 49 | 50 | add_executable(${PROJECT_NAME} 51 | main.cpp 52 | include/yolov8-pose.hpp 53 | include/common.hpp 54 | ) 55 | 56 | target_link_directories(${PROJECT_NAME} PUBLIC ${ALL_LIBS}) 57 | target_link_libraries(${PROJECT_NAME} PRIVATE nvinfer nvinfer_plugin cudart ${OpenCV_LIBS}) 58 | 59 | if (${OpenCV_VERSION} VERSION_GREATER_EQUAL 4.7.0) 60 | message(STATUS "Build with -DBATCHED_NMS") 61 | add_definitions(-DBATCHED_NMS) 62 | endif () 63 | -------------------------------------------------------------------------------- /src/pose/normal/include/common.hpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 4/7/23. 3 | // 4 | 5 | #ifndef POSE_NORMAL_COMMON_HPP 6 | #define POSE_NORMAL_COMMON_HPP 7 | #include "NvInfer.h" 8 | #include "opencv2/opencv.hpp" 9 | #include 10 | #include 11 | 12 | #define CHECK(call) \ 13 | do { \ 14 | const cudaError_t error_code = call; \ 15 | if (error_code != cudaSuccess) { \ 16 | printf("CUDA Error:\n"); \ 17 | printf(" File: %s\n", __FILE__); \ 18 | printf(" Line: %d\n", __LINE__); \ 19 | printf(" Error code: %d\n", error_code); \ 20 | printf(" Error text: %s\n", cudaGetErrorString(error_code)); \ 21 | exit(1); \ 22 | } \ 23 | } while (0) 24 | 25 | class Logger: public nvinfer1::ILogger { 26 | public: 27 | nvinfer1::ILogger::Severity reportableSeverity; 28 | 29 | explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO): 30 | reportableSeverity(severity) 31 | { 32 | } 33 | 34 | void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override 35 | { 36 | if (severity > reportableSeverity) { 37 | return; 38 | } 39 | switch (severity) { 40 | case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: 41 | std::cerr << "INTERNAL_ERROR: "; 42 | break; 43 | case nvinfer1::ILogger::Severity::kERROR: 44 | std::cerr << "ERROR: "; 45 | break; 46 | case nvinfer1::ILogger::Severity::kWARNING: 47 | std::cerr << "WARNING: "; 48 | break; 49 | case nvinfer1::ILogger::Severity::kINFO: 50 | std::cerr << "INFO: "; 51 | break; 52 | default: 53 | std::cerr << "VERBOSE: "; 54 | break; 55 | } 56 | std::cerr << msg << std::endl; 57 | } 58 | }; 59 | 60 | inline int get_size_by_dims(const nvinfer1::Dims& dims) 61 | { 62 | int size = 1; 63 | for (int i = 0; i < dims.nbDims; i++) { 64 | size *= dims.d[i]; 65 | } 66 | return size; 67 | } 68 | 69 | inline int type_to_size(const nvinfer1::DataType& dataType) 70 | { 71 | switch (dataType) { 72 | case nvinfer1::DataType::kFLOAT: 73 | return 4; 74 | case nvinfer1::DataType::kHALF: 75 | return 2; 76 | case nvinfer1::DataType::kINT32: 77 | return 4; 78 | case nvinfer1::DataType::kINT8: 79 | return 1; 80 | case nvinfer1::DataType::kBOOL: 81 | return 1; 82 | default: 83 | return 4; 84 | } 85 | } 86 | 87 | inline static float clamp(float val, float min, float max) 88 | { 89 | return val > min ? (val < max ? val : max) : min; 90 | } 91 | 92 | inline bool IsPathExist(const std::string& path) 93 | { 94 | if (access(path.c_str(), 0) == F_OK) { 95 | return true; 96 | } 97 | return false; 98 | } 99 | 100 | inline bool IsFile(const std::string& path) 101 | { 102 | if (!IsPathExist(path)) { 103 | printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str()); 104 | return false; 105 | } 106 | struct stat buffer; 107 | return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode)); 108 | } 109 | 110 | inline bool IsFolder(const std::string& path) 111 | { 112 | if (!IsPathExist(path)) { 113 | return false; 114 | } 115 | struct stat buffer; 116 | return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode)); 117 | } 118 | 119 | namespace pose { 120 | struct Binding { 121 | size_t size = 1; 122 | size_t dsize = 1; 123 | nvinfer1::Dims dims; 124 | std::string name; 125 | }; 126 | 127 | struct Object { 128 | cv::Rect_ rect; 129 | int label = 0; 130 | float prob = 0.0; 131 | std::vector kps; 132 | }; 133 | 134 | struct PreParam { 135 | float ratio = 1.0f; 136 | float dw = 0.0f; 137 | float dh = 0.0f; 138 | float height = 0; 139 | float width = 0; 140 | }; 141 | } // namespace pose 142 | #endif // POSE_NORMAL_COMMON_HPP 143 | -------------------------------------------------------------------------------- /src/pose/normal/main.cpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 4/7/23. 3 | // 4 | #include "chrono" 5 | #include "opencv2/opencv.hpp" 6 | #include "yolov8-pose.hpp" 7 | 8 | const std::vector> KPS_COLORS = {{0, 255, 0}, 9 | {0, 255, 0}, 10 | {0, 255, 0}, 11 | {0, 255, 0}, 12 | {0, 255, 0}, 13 | {255, 128, 0}, 14 | {255, 128, 0}, 15 | {255, 128, 0}, 16 | {255, 128, 0}, 17 | {255, 128, 0}, 18 | {255, 128, 0}, 19 | {51, 153, 255}, 20 | {51, 153, 255}, 21 | {51, 153, 255}, 22 | {51, 153, 255}, 23 | {51, 153, 255}, 24 | {51, 153, 255}}; 25 | 26 | const std::vector> SKELETON = {{16, 14}, 27 | {14, 12}, 28 | {17, 15}, 29 | {15, 13}, 30 | {12, 13}, 31 | {6, 12}, 32 | {7, 13}, 33 | {6, 7}, 34 | {6, 8}, 35 | {7, 9}, 36 | {8, 10}, 37 | {9, 11}, 38 | {2, 3}, 39 | {1, 2}, 40 | {1, 3}, 41 | {2, 4}, 42 | {3, 5}, 43 | {4, 6}, 44 | {5, 7}}; 45 | 46 | const std::vector> LIMB_COLORS = {{51, 153, 255}, 47 | {51, 153, 255}, 48 | {51, 153, 255}, 49 | {51, 153, 255}, 50 | {255, 51, 255}, 51 | {255, 51, 255}, 52 | {255, 51, 255}, 53 | {255, 128, 0}, 54 | {255, 128, 0}, 55 | {255, 128, 0}, 56 | {255, 128, 0}, 57 | {255, 128, 0}, 58 | {0, 255, 0}, 59 | {0, 255, 0}, 60 | {0, 255, 0}, 61 | {0, 255, 0}, 62 | {0, 255, 0}, 63 | {0, 255, 0}, 64 | {0, 255, 0}}; 65 | 66 | int main(int argc, char** argv) 67 | { 68 | // cuda:0 69 | cudaSetDevice(0); 70 | 71 | const std::string engine_file_path{argv[1]}; 72 | const std::string path{argv[2]}; 73 | 74 | std::vector imagePathList; 75 | bool isVideo{false}; 76 | 77 | assert(argc == 3); 78 | 79 | auto yolov8_pose = new YOLOv8_pose(engine_file_path); 80 | yolov8_pose->make_pipe(true); 81 | 82 | if (IsFile(path)) { 83 | std::string suffix = path.substr(path.find_last_of('.') + 1); 84 | if (suffix == "jpg" || suffix == "jpeg" || suffix == "png") { 85 | imagePathList.push_back(path); 86 | } 87 | else if (suffix == "mp4" || suffix == "avi" || suffix == "m4v" || suffix == "mpeg" || suffix == "mov" 88 | || suffix == "mkv") { 89 | isVideo = true; 90 | } 91 | else { 92 | printf("suffix %s is wrong !!!\n", suffix.c_str()); 93 | std::abort(); 94 | } 95 | } 96 | else if (IsFolder(path)) { 97 | cv::glob(path + "/*.jpg", imagePathList); 98 | } 99 | 100 | cv::Mat res, image; 101 | cv::Size size = cv::Size{640, 640}; 102 | int topk = 100; 103 | float score_thres = 0.25f; 104 | float iou_thres = 0.65f; 105 | 106 | std::vector objs; 107 | 108 | cv::namedWindow("result", cv::WINDOW_AUTOSIZE); 109 | 110 | if (isVideo) { 111 | cv::VideoCapture cap(path); 112 | 113 | if (!cap.isOpened()) { 114 | printf("can not open %s\n", path.c_str()); 115 | return -1; 116 | } 117 | while (cap.read(image)) { 118 | objs.clear(); 119 | yolov8_pose->copy_from_Mat(image, size); 120 | auto start = std::chrono::system_clock::now(); 121 | yolov8_pose->infer(); 122 | auto end = std::chrono::system_clock::now(); 123 | yolov8_pose->postprocess(objs, score_thres, iou_thres, topk); 124 | yolov8_pose->draw_objects(image, res, objs, SKELETON, KPS_COLORS, LIMB_COLORS); 125 | auto tc = (double)std::chrono::duration_cast(end - start).count() / 1000.; 126 | printf("cost %2.4lf ms\n", tc); 127 | cv::imshow("result", res); 128 | if (cv::waitKey(10) == 'q') { 129 | break; 130 | } 131 | } 132 | } 133 | else { 134 | for (auto& path : imagePathList) { 135 | objs.clear(); 136 | image = cv::imread(path); 137 | yolov8_pose->copy_from_Mat(image, size); 138 | auto start = std::chrono::system_clock::now(); 139 | yolov8_pose->infer(); 140 | auto end = std::chrono::system_clock::now(); 141 | yolov8_pose->postprocess(objs, score_thres, iou_thres, topk); 142 | yolov8_pose->draw_objects(image, res, objs, SKELETON, KPS_COLORS, LIMB_COLORS); 143 | auto tc = (double)std::chrono::duration_cast(end - start).count() / 1000.; 144 | printf("cost %2.4lf ms\n", tc); 145 | cv::imshow("result", res); 146 | cv::waitKey(0); 147 | } 148 | } 149 | cv::destroyAllWindows(); 150 | delete yolov8_pose; 151 | return 0; 152 | } 153 | -------------------------------------------------------------------------------- /src/segment/normal/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.1) 2 | 3 | set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86 89 90) 4 | set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc) 5 | 6 | project(yolov8-seg LANGUAGES CXX CUDA) 7 | 8 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -O3") 9 | set(CMAKE_CXX_STANDARD 14) 10 | set(CMAKE_BUILD_TYPE Release) 11 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 12 | 13 | # CUDA 14 | find_package(CUDA REQUIRED) 15 | message(STATUS "CUDA Libs: \n${CUDA_LIBRARIES}\n") 16 | get_filename_component(CUDA_LIB_DIR ${CUDA_LIBRARIES} DIRECTORY) 17 | message(STATUS "CUDA Headers: \n${CUDA_INCLUDE_DIRS}\n") 18 | 19 | # OpenCV 20 | find_package(OpenCV REQUIRED) 21 | message(STATUS "OpenCV Libs: \n${OpenCV_LIBS}\n") 22 | message(STATUS "OpenCV Libraries: \n${OpenCV_LIBRARIES}\n") 23 | message(STATUS "OpenCV Headers: \n${OpenCV_INCLUDE_DIRS}\n") 24 | 25 | # TensorRT 26 | set(TensorRT_ROOT /home/lin/software/TensorRT-8.5.3.1) 27 | set(TensorRT_INCLUDE_DIRS ${TensorRT_ROOT}/include) 28 | set(TensorRT_LIBRARIES ${TensorRT_ROOT}/lib) 29 | 30 | message(STATUS "TensorRT Libs: \n${TensorRT_LIBRARIES}\n") 31 | message(STATUS "TensorRT Headers: \n${TensorRT_INCLUDE_DIRS}\n") 32 | 33 | list(APPEND INCLUDE_DIRS 34 | ${CUDA_INCLUDE_DIRS} 35 | ${OpenCV_INCLUDE_DIRS} 36 | ${TensorRT_INCLUDE_DIRS} 37 | include 38 | ) 39 | 40 | list(APPEND ALL_LIBSlll 41 | ${CUDA_LIBRARIES} 42 | ${CUDA_LIB_DIR} 43 | ${OpenCV_LIBRARIES} 44 | ${TensorRT_LIBRARIES} 45 | ) 46 | 47 | include_directories(${INCLUDE_DIRS}) 48 | 49 | add_executable(${PROJECT_NAME} 50 | main.cpp 51 | include/yolov8-seg.hpp 52 | include/common.hpp 53 | ) 54 | 55 | target_link_directories(${PROJECT_NAME} PUBLIC ${ALL_LIBS}) 56 | target_link_libraries(${PROJECT_NAME} PRIVATE nvinfer nvinfer_plugin cudart ${OpenCV_LIBS}) 57 | 58 | if (${OpenCV_VERSION} VERSION_GREATER_EQUAL 4.7.0) 59 | message(STATUS "Build with -DBATCHED_NMS") 60 | add_definitions(-DBATCHED_NMS) 61 | endif () 62 | -------------------------------------------------------------------------------- /src/segment/normal/include/common.hpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 2/8/23. 3 | // 4 | 5 | #ifndef SEGMENT_NORMAL_COMMON_HPP 6 | #define SEGMENT_NORMAL_COMMON_HPP 7 | #include "NvInfer.h" 8 | #include "opencv2/opencv.hpp" 9 | #include 10 | #include 11 | 12 | #define CHECK(call) \ 13 | do { \ 14 | const cudaError_t error_code = call; \ 15 | if (error_code != cudaSuccess) { \ 16 | printf("CUDA Error:\n"); \ 17 | printf(" File: %s\n", __FILE__); \ 18 | printf(" Line: %d\n", __LINE__); \ 19 | printf(" Error code: %d\n", error_code); \ 20 | printf(" Error text: %s\n", cudaGetErrorString(error_code)); \ 21 | exit(1); \ 22 | } \ 23 | } while (0) 24 | 25 | class Logger: public nvinfer1::ILogger { 26 | public: 27 | nvinfer1::ILogger::Severity reportableSeverity; 28 | 29 | explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO): 30 | reportableSeverity(severity) 31 | { 32 | } 33 | 34 | void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override 35 | { 36 | if (severity > reportableSeverity) { 37 | return; 38 | } 39 | switch (severity) { 40 | case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: 41 | std::cerr << "INTERNAL_ERROR: "; 42 | break; 43 | case nvinfer1::ILogger::Severity::kERROR: 44 | std::cerr << "ERROR: "; 45 | break; 46 | case nvinfer1::ILogger::Severity::kWARNING: 47 | std::cerr << "WARNING: "; 48 | break; 49 | case nvinfer1::ILogger::Severity::kINFO: 50 | std::cerr << "INFO: "; 51 | break; 52 | default: 53 | std::cerr << "VERBOSE: "; 54 | break; 55 | } 56 | std::cerr << msg << std::endl; 57 | } 58 | }; 59 | 60 | inline int get_size_by_dims(const nvinfer1::Dims& dims) 61 | { 62 | int size = 1; 63 | for (int i = 0; i < dims.nbDims; i++) { 64 | size *= dims.d[i]; 65 | } 66 | return size; 67 | } 68 | 69 | inline int type_to_size(const nvinfer1::DataType& dataType) 70 | { 71 | switch (dataType) { 72 | case nvinfer1::DataType::kFLOAT: 73 | return 4; 74 | case nvinfer1::DataType::kHALF: 75 | return 2; 76 | case nvinfer1::DataType::kINT32: 77 | return 4; 78 | case nvinfer1::DataType::kINT8: 79 | return 1; 80 | case nvinfer1::DataType::kBOOL: 81 | return 1; 82 | default: 83 | return 4; 84 | } 85 | } 86 | 87 | inline static float clamp(float val, float min, float max) 88 | { 89 | return val > min ? (val < max ? val : max) : min; 90 | } 91 | 92 | inline bool IsPathExist(const std::string& path) 93 | { 94 | if (access(path.c_str(), 0) == F_OK) { 95 | return true; 96 | } 97 | return false; 98 | } 99 | 100 | inline bool IsFile(const std::string& path) 101 | { 102 | if (!IsPathExist(path)) { 103 | printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str()); 104 | return false; 105 | } 106 | struct stat buffer; 107 | return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode)); 108 | } 109 | 110 | inline bool IsFolder(const std::string& path) 111 | { 112 | if (!IsPathExist(path)) { 113 | return false; 114 | } 115 | struct stat buffer; 116 | return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode)); 117 | } 118 | 119 | namespace seg { 120 | struct Binding { 121 | size_t size = 1; 122 | size_t dsize = 1; 123 | nvinfer1::Dims dims; 124 | std::string name; 125 | }; 126 | 127 | struct Object { 128 | cv::Rect_ rect; 129 | int label = 0; 130 | float prob = 0.0; 131 | cv::Mat boxMask; 132 | }; 133 | 134 | struct PreParam { 135 | float ratio = 1.0f; 136 | float dw = 0.0f; 137 | float dh = 0.0f; 138 | float height = 0; 139 | float width = 0; 140 | }; 141 | } // namespace seg 142 | #endif // SEGMENT_NORMAL_COMMON_HPP 143 | -------------------------------------------------------------------------------- /src/segment/normal/main.cpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 2/8/23. 3 | // 4 | #include "chrono" 5 | #include "opencv2/opencv.hpp" 6 | #include "yolov8-seg.hpp" 7 | 8 | const std::vector CLASS_NAMES = { 9 | "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", 10 | "truck", "boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", 11 | "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", 12 | "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", 13 | "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat", 14 | "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup", 15 | "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", 16 | "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", 17 | "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", 18 | "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven", 19 | "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", 20 | "teddy bear", "hair drier", "toothbrush"}; 21 | 22 | const std::vector> COLORS = { 23 | {0, 114, 189}, {217, 83, 25}, {237, 177, 32}, {126, 47, 142}, {119, 172, 48}, {77, 190, 238}, 24 | {162, 20, 47}, {76, 76, 76}, {153, 153, 153}, {255, 0, 0}, {255, 128, 0}, {191, 191, 0}, 25 | {0, 255, 0}, {0, 0, 255}, {170, 0, 255}, {85, 85, 0}, {85, 170, 0}, {85, 255, 0}, 26 | {170, 85, 0}, {170, 170, 0}, {170, 255, 0}, {255, 85, 0}, {255, 170, 0}, {255, 255, 0}, 27 | {0, 85, 128}, {0, 170, 128}, {0, 255, 128}, {85, 0, 128}, {85, 85, 128}, {85, 170, 128}, 28 | {85, 255, 128}, {170, 0, 128}, {170, 85, 128}, {170, 170, 128}, {170, 255, 128}, {255, 0, 128}, 29 | {255, 85, 128}, {255, 170, 128}, {255, 255, 128}, {0, 85, 255}, {0, 170, 255}, {0, 255, 255}, 30 | {85, 0, 255}, {85, 85, 255}, {85, 170, 255}, {85, 255, 255}, {170, 0, 255}, {170, 85, 255}, 31 | {170, 170, 255}, {170, 255, 255}, {255, 0, 255}, {255, 85, 255}, {255, 170, 255}, {85, 0, 0}, 32 | {128, 0, 0}, {170, 0, 0}, {212, 0, 0}, {255, 0, 0}, {0, 43, 0}, {0, 85, 0}, 33 | {0, 128, 0}, {0, 170, 0}, {0, 212, 0}, {0, 255, 0}, {0, 0, 43}, {0, 0, 85}, 34 | {0, 0, 128}, {0, 0, 170}, {0, 0, 212}, {0, 0, 255}, {0, 0, 0}, {36, 36, 36}, 35 | {73, 73, 73}, {109, 109, 109}, {146, 146, 146}, {182, 182, 182}, {219, 219, 219}, {0, 114, 189}, 36 | {80, 183, 189}, {128, 128, 0}}; 37 | 38 | const std::vector> MASK_COLORS = { 39 | {255, 56, 56}, {255, 157, 151}, {255, 112, 31}, {255, 178, 29}, {207, 210, 49}, {72, 249, 10}, {146, 204, 23}, 40 | {61, 219, 134}, {26, 147, 52}, {0, 212, 187}, {44, 153, 168}, {0, 194, 255}, {52, 69, 147}, {100, 115, 255}, 41 | {0, 24, 236}, {132, 56, 255}, {82, 0, 133}, {203, 56, 255}, {255, 149, 200}, {255, 55, 199}}; 42 | 43 | int main(int argc, char** argv) 44 | { 45 | // cuda:0 46 | cudaSetDevice(0); 47 | 48 | const std::string engine_file_path{argv[1]}; 49 | const std::string path{argv[2]}; 50 | 51 | std::vector imagePathList; 52 | bool isVideo{false}; 53 | 54 | assert(argc == 3); 55 | 56 | auto yolov8 = new YOLOv8_seg(engine_file_path); 57 | yolov8->make_pipe(true); 58 | 59 | if (IsFile(path)) { 60 | std::string suffix = path.substr(path.find_last_of('.') + 1); 61 | if (suffix == "jpg" || suffix == "jpeg" || suffix == "png") { 62 | imagePathList.push_back(path); 63 | } 64 | else if (suffix == "mp4" || suffix == "avi" || suffix == "m4v" || suffix == "mpeg" || suffix == "mov" 65 | || suffix == "mkv") { 66 | isVideo = true; 67 | } 68 | else { 69 | printf("suffix %s is wrong !!!\n", suffix.c_str()); 70 | std::abort(); 71 | } 72 | } 73 | else if (IsFolder(path)) { 74 | cv::glob(path + "/*.jpg", imagePathList); 75 | } 76 | 77 | cv::Mat res, image; 78 | cv::Size size = cv::Size{640, 640}; 79 | int topk = 100; 80 | int seg_h = 160; 81 | int seg_w = 160; 82 | int seg_channels = 32; 83 | float score_thres = 0.25f; 84 | float iou_thres = 0.65f; 85 | 86 | std::vector objs; 87 | 88 | cv::namedWindow("result", cv::WINDOW_AUTOSIZE); 89 | 90 | if (isVideo) { 91 | cv::VideoCapture cap(path); 92 | 93 | if (!cap.isOpened()) { 94 | printf("can not open %s\n", path.c_str()); 95 | return -1; 96 | } 97 | while (cap.read(image)) { 98 | objs.clear(); 99 | yolov8->copy_from_Mat(image, size); 100 | auto start = std::chrono::system_clock::now(); 101 | yolov8->infer(); 102 | auto end = std::chrono::system_clock::now(); 103 | yolov8->postprocess(objs, score_thres, iou_thres, topk, seg_channels, seg_h, seg_w); 104 | yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS, MASK_COLORS); 105 | auto tc = (double)std::chrono::duration_cast(end - start).count() / 1000.; 106 | printf("cost %2.4lf ms\n", tc); 107 | cv::imshow("result", res); 108 | if (cv::waitKey(10) == 'q') { 109 | break; 110 | } 111 | } 112 | } 113 | else { 114 | for (auto& path : imagePathList) { 115 | objs.clear(); 116 | image = cv::imread(path); 117 | yolov8->copy_from_Mat(image, size); 118 | auto start = std::chrono::system_clock::now(); 119 | yolov8->infer(); 120 | auto end = std::chrono::system_clock::now(); 121 | yolov8->postprocess(objs, score_thres, iou_thres, topk, seg_channels, seg_h, seg_w); 122 | yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS, MASK_COLORS); 123 | auto tc = (double)std::chrono::duration_cast(end - start).count() / 1000.; 124 | printf("cost %2.4lf ms\n", tc); 125 | cv::imshow("result", res); 126 | cv::waitKey(0); 127 | } 128 | } 129 | cv::destroyAllWindows(); 130 | delete yolov8; 131 | return 0; 132 | } 133 | -------------------------------------------------------------------------------- /src/segment/simple/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.1) 2 | 3 | set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86 89 90) 4 | set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc) 5 | 6 | project(yolov8-seg LANGUAGES CXX CUDA) 7 | 8 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -O3") 9 | set(CMAKE_CXX_STANDARD 14) 10 | set(CMAKE_BUILD_TYPE Release) 11 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 12 | 13 | # CUDA 14 | find_package(CUDA REQUIRED) 15 | message(STATUS "CUDA Libs: \n${CUDA_LIBRARIES}\n") 16 | get_filename_component(CUDA_LIB_DIR ${CUDA_LIBRARIES} DIRECTORY) 17 | message(STATUS "CUDA Headers: \n${CUDA_INCLUDE_DIRS}\n") 18 | 19 | # OpenCV 20 | find_package(OpenCV REQUIRED) 21 | message(STATUS "OpenCV Libs: \n${OpenCV_LIBS}\n") 22 | message(STATUS "OpenCV Libraries: \n${OpenCV_LIBRARIES}\n") 23 | message(STATUS "OpenCV Headers: \n${OpenCV_INCLUDE_DIRS}\n") 24 | 25 | # TensorRT 26 | set(TensorRT_ROOT /home/lin/software/TensorRT-8.5.3.1) 27 | set(TensorRT_INCLUDE_DIRS ${TensorRT_ROOT}/include) 28 | set(TensorRT_LIBRARIES ${TensorRT_ROOT}/lib) 29 | 30 | message(STATUS "TensorRT Libs: \n${TensorRT_LIBRARIES}\n") 31 | message(STATUS "TensorRT Headers: \n${TensorRT_INCLUDE_DIRS}\n") 32 | 33 | list(APPEND INCLUDE_DIRS 34 | ${CUDA_INCLUDE_DIRS} 35 | ${OpenCV_INCLUDE_DIRS} 36 | ${TensorRT_INCLUDE_DIRS} 37 | ./include 38 | ) 39 | 40 | list(APPEND ALL_LIBS 41 | ${CUDA_LIBRARIES} 42 | ${CUDA_LIB_DIR} 43 | ${OpenCV_LIBRARIES} 44 | ${TensorRT_LIBRARIES} 45 | ) 46 | 47 | include_directories(${INCLUDE_DIRS}) 48 | 49 | add_executable(${PROJECT_NAME} 50 | main.cpp 51 | include/yolov8-seg.hpp 52 | include/common.hpp 53 | ) 54 | 55 | target_link_directories(${PROJECT_NAME} PUBLIC ${ALL_LIBS}) 56 | target_link_libraries(${PROJECT_NAME} PRIVATE nvinfer nvinfer_plugin cudart ${OpenCV_LIBS}) 57 | 58 | if (${OpenCV_VERSION} VERSION_GREATER_EQUAL 4.7.0) 59 | message(STATUS "Build with -DBATCHED_NMS") 60 | add_definitions(-DBATCHED_NMS) 61 | endif () 62 | -------------------------------------------------------------------------------- /src/segment/simple/include/common.hpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 2/9/23. 3 | // 4 | 5 | #ifndef SEGMENT_SIMPLE_COMMON_HPP 6 | #define SEGMENT_SIMPLE_COMMON_HPP 7 | #include "NvInfer.h" 8 | #include "opencv2/opencv.hpp" 9 | #include 10 | #include 11 | 12 | #define CHECK(call) \ 13 | do { \ 14 | const cudaError_t error_code = call; \ 15 | if (error_code != cudaSuccess) { \ 16 | printf("CUDA Error:\n"); \ 17 | printf(" File: %s\n", __FILE__); \ 18 | printf(" Line: %d\n", __LINE__); \ 19 | printf(" Error code: %d\n", error_code); \ 20 | printf(" Error text: %s\n", cudaGetErrorString(error_code)); \ 21 | exit(1); \ 22 | } \ 23 | } while (0) 24 | 25 | class Logger: public nvinfer1::ILogger { 26 | public: 27 | nvinfer1::ILogger::Severity reportableSeverity; 28 | 29 | explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO): 30 | reportableSeverity(severity) 31 | { 32 | } 33 | 34 | void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override 35 | { 36 | if (severity > reportableSeverity) { 37 | return; 38 | } 39 | switch (severity) { 40 | case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: 41 | std::cerr << "INTERNAL_ERROR: "; 42 | break; 43 | case nvinfer1::ILogger::Severity::kERROR: 44 | std::cerr << "ERROR: "; 45 | break; 46 | case nvinfer1::ILogger::Severity::kWARNING: 47 | std::cerr << "WARNING: "; 48 | break; 49 | case nvinfer1::ILogger::Severity::kINFO: 50 | std::cerr << "INFO: "; 51 | break; 52 | default: 53 | std::cerr << "VERBOSE: "; 54 | break; 55 | } 56 | std::cerr << msg << std::endl; 57 | } 58 | }; 59 | 60 | inline int get_size_by_dims(const nvinfer1::Dims& dims) 61 | { 62 | int size = 1; 63 | for (int i = 0; i < dims.nbDims; i++) { 64 | size *= dims.d[i]; 65 | } 66 | return size; 67 | } 68 | 69 | inline int type_to_size(const nvinfer1::DataType& dataType) 70 | { 71 | switch (dataType) { 72 | case nvinfer1::DataType::kFLOAT: 73 | return 4; 74 | case nvinfer1::DataType::kHALF: 75 | return 2; 76 | case nvinfer1::DataType::kINT32: 77 | return 4; 78 | case nvinfer1::DataType::kINT8: 79 | return 1; 80 | case nvinfer1::DataType::kBOOL: 81 | return 1; 82 | default: 83 | return 4; 84 | } 85 | } 86 | 87 | inline static float clamp(float val, float min, float max) 88 | { 89 | return val > min ? (val < max ? val : max) : min; 90 | } 91 | 92 | inline bool IsPathExist(const std::string& path) 93 | { 94 | if (access(path.c_str(), 0) == F_OK) { 95 | return true; 96 | } 97 | return false; 98 | } 99 | 100 | inline bool IsFile(const std::string& path) 101 | { 102 | if (!IsPathExist(path)) { 103 | printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str()); 104 | return false; 105 | } 106 | struct stat buffer; 107 | return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode)); 108 | } 109 | 110 | inline bool IsFolder(const std::string& path) 111 | { 112 | if (!IsPathExist(path)) { 113 | return false; 114 | } 115 | struct stat buffer; 116 | return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode)); 117 | } 118 | 119 | namespace seg { 120 | struct Binding { 121 | size_t size = 1; 122 | size_t dsize = 1; 123 | nvinfer1::Dims dims; 124 | std::string name; 125 | }; 126 | 127 | struct Object { 128 | cv::Rect_ rect; 129 | int label = 0; 130 | float prob = 0.0; 131 | cv::Mat boxMask; 132 | }; 133 | 134 | struct PreParam { 135 | float ratio = 1.0f; 136 | float dw = 0.0f; 137 | float dh = 0.0f; 138 | float height = 0; 139 | float width = 0; 140 | }; 141 | } // namespace seg 142 | #endif // SEGMENT_SIMPLE_COMMON_HPP 143 | -------------------------------------------------------------------------------- /src/segment/simple/main.cpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by ubuntu on 1/20/23. 3 | // 4 | #include "chrono" 5 | #include "opencv2/opencv.hpp" 6 | #include "yolov8-seg.hpp" 7 | 8 | const std::vector CLASS_NAMES = { 9 | "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", 10 | "truck", "boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", 11 | "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", 12 | "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", 13 | "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat", 14 | "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup", 15 | "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", 16 | "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", 17 | "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", 18 | "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven", 19 | "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", 20 | "teddy bear", "hair drier", "toothbrush"}; 21 | 22 | const std::vector> COLORS = { 23 | {0, 114, 189}, {217, 83, 25}, {237, 177, 32}, {126, 47, 142}, {119, 172, 48}, {77, 190, 238}, 24 | {162, 20, 47}, {76, 76, 76}, {153, 153, 153}, {255, 0, 0}, {255, 128, 0}, {191, 191, 0}, 25 | {0, 255, 0}, {0, 0, 255}, {170, 0, 255}, {85, 85, 0}, {85, 170, 0}, {85, 255, 0}, 26 | {170, 85, 0}, {170, 170, 0}, {170, 255, 0}, {255, 85, 0}, {255, 170, 0}, {255, 255, 0}, 27 | {0, 85, 128}, {0, 170, 128}, {0, 255, 128}, {85, 0, 128}, {85, 85, 128}, {85, 170, 128}, 28 | {85, 255, 128}, {170, 0, 128}, {170, 85, 128}, {170, 170, 128}, {170, 255, 128}, {255, 0, 128}, 29 | {255, 85, 128}, {255, 170, 128}, {255, 255, 128}, {0, 85, 255}, {0, 170, 255}, {0, 255, 255}, 30 | {85, 0, 255}, {85, 85, 255}, {85, 170, 255}, {85, 255, 255}, {170, 0, 255}, {170, 85, 255}, 31 | {170, 170, 255}, {170, 255, 255}, {255, 0, 255}, {255, 85, 255}, {255, 170, 255}, {85, 0, 0}, 32 | {128, 0, 0}, {170, 0, 0}, {212, 0, 0}, {255, 0, 0}, {0, 43, 0}, {0, 85, 0}, 33 | {0, 128, 0}, {0, 170, 0}, {0, 212, 0}, {0, 255, 0}, {0, 0, 43}, {0, 0, 85}, 34 | {0, 0, 128}, {0, 0, 170}, {0, 0, 212}, {0, 0, 255}, {0, 0, 0}, {36, 36, 36}, 35 | {73, 73, 73}, {109, 109, 109}, {146, 146, 146}, {182, 182, 182}, {219, 219, 219}, {0, 114, 189}, 36 | {80, 183, 189}, {128, 128, 0}}; 37 | 38 | const std::vector> MASK_COLORS = { 39 | {255, 56, 56}, {255, 157, 151}, {255, 112, 31}, {255, 178, 29}, {207, 210, 49}, {72, 249, 10}, {146, 204, 23}, 40 | {61, 219, 134}, {26, 147, 52}, {0, 212, 187}, {44, 153, 168}, {0, 194, 255}, {52, 69, 147}, {100, 115, 255}, 41 | {0, 24, 236}, {132, 56, 255}, {82, 0, 133}, {203, 56, 255}, {255, 149, 200}, {255, 55, 199}}; 42 | 43 | int main(int argc, char** argv) 44 | { 45 | // cuda:0 46 | cudaSetDevice(0); 47 | 48 | const std::string engine_file_path{argv[1]}; 49 | const std::string path{argv[2]}; 50 | 51 | std::vector imagePathList; 52 | bool isVideo{false}; 53 | 54 | assert(argc == 3); 55 | 56 | auto yolov8 = new YOLOv8_seg(engine_file_path); 57 | yolov8->make_pipe(true); 58 | 59 | if (IsFile(path)) { 60 | std::string suffix = path.substr(path.find_last_of('.') + 1); 61 | if (suffix == "jpg" || suffix == "jpeg" || suffix == "png") { 62 | imagePathList.push_back(path); 63 | } 64 | else if (suffix == "mp4" || suffix == "avi" || suffix == "m4v" || suffix == "mpeg" || suffix == "mov" 65 | || suffix == "mkv") { 66 | isVideo = true; 67 | } 68 | else { 69 | printf("suffix %s is wrong !!!\n", suffix.c_str()); 70 | std::abort(); 71 | } 72 | } 73 | else if (IsFolder(path)) { 74 | cv::glob(path + "/*.jpg", imagePathList); 75 | } 76 | 77 | cv::Mat res, image; 78 | cv::Size size = cv::Size{640, 640}; 79 | int topk = 100; 80 | int seg_h = 160; 81 | int seg_w = 160; 82 | int seg_channels = 32; 83 | float score_thres = 0.25f; 84 | float iou_thres = 0.65f; 85 | 86 | std::vector objs; 87 | 88 | cv::namedWindow("result", cv::WINDOW_AUTOSIZE); 89 | 90 | if (isVideo) { 91 | cv::VideoCapture cap(path); 92 | 93 | if (!cap.isOpened()) { 94 | printf("can not open %s\n", path.c_str()); 95 | return -1; 96 | } 97 | while (cap.read(image)) { 98 | objs.clear(); 99 | yolov8->copy_from_Mat(image, size); 100 | auto start = std::chrono::system_clock::now(); 101 | yolov8->infer(); 102 | auto end = std::chrono::system_clock::now(); 103 | yolov8->postprocess(objs, score_thres, iou_thres, topk, seg_channels, seg_h, seg_w); 104 | yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS, MASK_COLORS); 105 | auto tc = (double)std::chrono::duration_cast(end - start).count() / 1000.; 106 | printf("cost %2.4lf ms\n", tc); 107 | cv::imshow("result", res); 108 | if (cv::waitKey(10) == 'q') { 109 | break; 110 | } 111 | } 112 | } 113 | else { 114 | for (auto& path : imagePathList) { 115 | objs.clear(); 116 | image = cv::imread(path); 117 | yolov8->copy_from_Mat(image, size); 118 | auto start = std::chrono::system_clock::now(); 119 | yolov8->infer(); 120 | auto end = std::chrono::system_clock::now(); 121 | yolov8->postprocess(objs, score_thres, iou_thres, topk, seg_channels, seg_h, seg_w); 122 | yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS, MASK_COLORS); 123 | auto tc = (double)std::chrono::duration_cast(end - start).count() / 1000.; 124 | printf("cost %2.4lf ms\n", tc); 125 | cv::imshow("result", res); 126 | cv::waitKey(0); 127 | } 128 | } 129 | cv::destroyAllWindows(); 130 | delete yolov8; 131 | return 0; 132 | } 133 | -------------------------------------------------------------------------------- /weights/.readme.md: -------------------------------------------------------------------------------- 1 | + 权重下载后放此目录下 -------------------------------------------------------------------------------- /weights/yolov8n-pose.onnx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linClubs/YOLOv8-ROS-TensorRT/0b23e9753b2cf0a2f26eda2335d867f1f2f69dda/weights/yolov8n-pose.onnx -------------------------------------------------------------------------------- /weights/yolov8n-seg.onnx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linClubs/YOLOv8-ROS-TensorRT/0b23e9753b2cf0a2f26eda2335d867f1f2f69dda/weights/yolov8n-seg.onnx -------------------------------------------------------------------------------- /weights/yolov8n.onnx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linClubs/YOLOv8-ROS-TensorRT/0b23e9753b2cf0a2f26eda2335d867f1f2f69dda/weights/yolov8n.onnx --------------------------------------------------------------------------------