# YOLOv8-seg Model with TensorRT The yolov8-seg model conversion route is : YOLOv8 PyTorch model -> ONNX -> TensorRT Engine ***Notice !!!*** We don't support TensorRT API building !!! # Export Modified ONNX model You can export your onnx model by `ultralytics` API and the onnx is also modify by this repo. ``` shell python3 export-seg.py \ --weights yolov8s-seg.pt \ --opset 11 \ --sim \ --input-shape 1 3 640 640 \ --device cuda:0 ``` #### Description of all arguments - `--weights` : The PyTorch model you trained. - `--opset` : ONNX opset version, default is 11. - `--sim` : Whether to simplify your onnx model. - `--input-shape` : Input shape for you model, should be 4 dimensions. - `--device` : The CUDA deivce you export engine . You will get an onnx model whose prefix is the same as input weights. This onnx model doesn't contain postprocessing. # Export Engine by TensorRT Python api You can export TensorRT engine from ONNX by [`build.py` ](../build.py). Usage: ``` shell python3 build.py \ --weights yolov8s-seg.onnx \ --fp16 \ --device cuda:0 \ --seg ``` #### Description of all arguments - `--weights` : The ONNX model you download. - `--fp16` : Whether to export half-precision engine. - `--device` : The CUDA deivce you export engine. - `--seg` : Whether to export seg engine. # Export Engine by Trtexec Tools You can export TensorRT engine by [`trtexec`](https://github.com/NVIDIA/TensorRT/tree/main/samples/trtexec) tools. Usage: ``` shell /usr/src/tensorrt/bin/trtexec \ --onnx=yolov8s-seg.onnx \ --saveEngine=yolov8s-seg.engine \ --fp16 ``` # Inference ## Infer with python script You can infer images with the engine by [`infer-seg.py`](../infer-seg.py) . Usage: ``` shell python3 infer-seg.py \ --engine yolov8s-seg.engine \ --imgs data \ --show \ --out-dir outputs \ --device cuda:0 ``` #### Description of all arguments - `--engine` : The Engine you export. - `--imgs` : The images path you want to detect. - `--show` : Whether to show detection results. - `--out-dir` : Where to save detection results images. It will not work when use `--show` flag. - `--device` : The CUDA deivce you use. - `--profile` : Profile the TensorRT engine. ## Infer with C++ You can infer segment engine with c++ in [`csrc/segment/simple`](../csrc/segment/simple) . ### Build: Please set you own librarys in [`CMakeLists.txt`](../csrc/segment/simple/CMakeLists.txt) and modify you own config in [`main.cpp`](../csrc/segment/simple/main.cpp) such as `CLASS_NAMES`, `COLORS`, `MASK_COLORS` and postprocess parameters . ```c++ int topk = 100; int seg_h = 160; // yolov8 model proto height int seg_w = 160; // yolov8 model proto width int seg_channels = 32; // yolov8 model proto channels float score_thres = 0.25f; float iou_thres = 0.65f; ``` ``` shell export root=${PWD} cd src/segment/simple mkdir build cmake .. make mv yolov8-seg ${root} cd ${root} ``` ***Notice !!!*** If you have build OpenCV(>=4.7.0) by yourself, it provides a new api [`cv::dnn::NMSBoxesBatched`](https://docs.opencv.org/4.x/d6/d0f/group__dnn.html#ga977aae09fbf7c804e003cfea1d4e928c) . It is a gread api about efficient in-class nms . It will be used by default! ***!!!*** Usage: ``` shell # infer image ./yolov8-seg yolov8s-seg.engine data/bus.jpg # infer images ./yolov8-seg yolov8s-seg.engine data # infer video ./yolov8-seg yolov8s-seg.engine data/test.mp4 # the video path ``` # Export Orin ONNX model by ultralytics You can leave this repo and use the original `ultralytics` repo for onnx export. ### 1. Python script Usage: ```python from ultralytics import YOLO # Load a model model = YOLO("yolov8s-seg.pt") # load a pretrained model (recommended for training) success = model.export(format="engine", device=0) # export the model to engine format assert success ``` After executing the above script, you will get an engine named `yolov8s-seg.engine` . ### 2. CLI tools Usage: ```shell yolo export model=yolov8s.pt format=engine device=0 ``` After executing the above command, you will get an engine named `yolov8s-seg.engine` too. ## Inference with c++ You can infer with c++ in [`csrc/segment/normal`](../csrc/segment/normal) . ### Build: Please set you own librarys in [`CMakeLists.txt`](../csrc/segment/normal/CMakeLists.txt) and modify `CLASS_NAMES` and `COLORS` in [`main.cpp`](../csrc/segment/normal/main.cpp). Besides, you can modify the postprocess parameters such as `num_labels` and `score_thres` and `iou_thres` and `topk` in [`main.cpp`](../csrc/segment/normal/main.cpp). ```c++ int topk = 100; int seg_h = 160; // yolov8 model proto height int seg_w = 160; // yolov8 model proto width int seg_channels = 32; // yolov8 model proto channels float score_thres = 0.25f; float iou_thres = 0.65f; ``` And build: ``` shell export root=${PWD} cd src/segment/normal mkdir build cmake .. make mv yolov8-seg ${root} cd ${root} ``` Usage: ``` shell # infer image ./yolov8-seg yolov8s-seg.engine data/bus.jpg # infer images ./yolov8-seg yolov8s-seg.engine data # infer video ./yolov8-seg yolov8s-seg.engine data/test.mp4 # the video path ``` # Refuse To Use PyTorch for segment Model Inference !!! It is the same as detection model. you can get more information in [`infer-seg-without-torch.py`](../infer-seg-without-torch.py), Usage: ``` shell python3 infer-seg-without-torch.py \ --engine yolov8s-seg.engine \ --imgs data \ --show \ --out-dir outputs \ --method cudart ``` #### Description of all arguments - `--engine` : The Engine you export. - `--imgs` : The images path you want to detect. - `--show` : Whether to show detection results. - `--out-dir` : Where to save detection results images. It will not work when use `--show` flag. - `--method` : Choose `cudart` or `pycuda`, default is `cudart`. - `--profile` : Profile the TensorRT engine.