Add Support for YOLO OBB Export and Inference with EfficientRotatedNM…

…S Plugin (#16)

.gitignore

@@ -10,4 +10,27 @@ __pycache__/
 
 
 # output Cache
-output/
+output/
+
+# demo images
+demo/**/*.jpg
+demo/**/*.jpeg
+demo/**/*.png
+demo/**/*.gif
+demo/**/*.bmp
+demo/**/*.tiff
+demo/**/*.webp
+
+# demo models
+demo/**/*.pt
+demo/**/*.onnx
+demo/**/*.engine
+
+# python package
+dist/
+tensorrt_yolo.egg-info/
+tensorrt_yolo/c_lib_wrap.py
+
+# build libs
+lib/
+tensorrt_yolo/libs/

demo/detect/README.en.md

@@ -1,77 +1,120 @@
-English | [简体中文](README.md)
+[简体中文](README.md) | English
 
 # Model Inference Examples
 
-In this example, we demonstrate how to perform model inference using YOLOv8s with CLI, Python, and C++.
+This example demonstrates how to perform model inference using CLI, Python, and C++ with the YOLOv8s model as an example.
+
+> [!IMPORTANT]  
+> If you want to use the [EfficientRotatedNMS](../../plugin/efficientRotatedNMSPlugin) plugin to infer an OBB model, please refer to [Building TensorRT Custom Plugins](../../docs/en/build_trt_custom_plugin.md) for guidance.
 
 ## Model Export
 
-First, download the YOLOv8s model from [YOLOv8s](https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8s.pt) and save it to the `models` folder.
+### Detection Model
 
-Then, export the model to ONNX format with the [EfficientNMS](https://github.com/NVIDIA/TensorRT/tree/main/plugin/efficientNMSPlugin) plugin using the following command:
+1. Download the [YOLOv8s](https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8s.pt) model and save it to the `models` folder.
+2. Use the following command to export the model to ONNX format with the [EfficientNMS](https://github.com/NVIDIA/TensorRT/tree/main/plugin/efficientNMSPlugin) plugin:
 
-```bash
-trtyolo export -w models/yolov8s.pt -v yolov8 -o models
-```
+    ```bash
+    trtyolo export -w models/yolov8s.pt -v yolov8 -o models
+    ```
 
-After executing the above command, a file named `yolov8s.onnx` will be generated in the `models` folder. Next, convert the ONNX file to a TensorRT engine using the `trtexec` tool:
+    After running the above command, a `yolov8s.onnx` file will be generated in the `models` folder. Next, use the `trtexec` tool to convert the ONNX file to a TensorRT engine:
 
-```bash
-trtexec --onnx=models/yolov8s.onnx --saveEngine=models/yolov8s.engine --fp16
-```
+    ```bash
+    trtexec --onnx=models/yolov8s.onnx --saveEngine=models/yolov8s.engine --fp16
+    ```
 
-Now, we will perform model inference using different methods.
+### Oriented Bounding Box Model (OBB)
 
-## Model Inference
+1. Download the [YOLOv8s-obb](https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8s-obb.pt) model and save it to the `models` folder.
+2. Use the following command to export the model to ONNX format with the [EfficientRotatedNMS](../../plugin/efficientRotatedNMSPlugin/) plugin:
+
+    ```bash
+    trtyolo export -w models/yolov8s-obb.pt -v yolov8 -o models
+    ```
+
+    After running the above command, a `yolov8s-obb.onnx` file will be generated in the `models` folder. Next, use the `trtexec` tool to convert the ONNX file to a TensorRT engine:
+
+    ```bash
+    trtexec --onnx=models/yolov8s-obb.onnx --saveEngine=models/yolov8s-obb.engine --fp16
+    ```
+
+## Dataset Preparation
 
-Download the [coco128](https://ultralytics.com/assets/coco128.zip) dataset, unzip it, and move the images from the `coco128/images/train2017` folder to the `images` folder for inference.
+### Detection Model
+
+1. Download the [coco128](https://ultralytics.com/assets/coco128.zip) dataset.
+2. After extraction, move the images from the `coco128/images/train2017` folder to the `images` folder for inference.
+
+### Oriented Bounding Box Model (OBB)
+
+1. Download the [DOTA-v1.0](https://drive.google.com/file/d/1fwiTNqRRen09E-O9VSpcMV2e6_d4GGVK/view) dataset.
+2. After extraction, move the images from the `part1/images` folder to the `images` folder for inference.
+
+## Model Inference
 
-### Inference using CLI
+### Inference Using CLI
 
-You can use the `trtyolo` command-line tool provided by `tensorrt_yolo` for inference. Run the following command to see the help information related to inference:
+1. Use the `trtyolo` command-line tool for inference. Run the following command to view the help information:
 
-```bash
-trtyolo infer --help
-```
+    ```bash
+    trtyolo infer --help
+    ```
 
-Then, perform inference using the following command:
+2. Run the following commands for inference:
 
-> To further speed up the inference process, use the `--cudaGraph` option, but this feature only supports static models, not dynamic models. (not supported before version 4.0)
+    > [!NOTE] 
+    > The `--cudaGraph` option, introduced in version 4.0, can further accelerate the inference process, but this feature only supports static models.
+    > 
+    > From version 4.2 onwards, OBB model inference is supported, with the new `-m, --mode` option for selecting Detection or OBB models.
 
-```bash
-trtyolo infer -e models/yolov8s.engine -i images -o output -l labels.txt --cudaGraph
-```
+    ```bash
+    # Detection Model
+    trtyolo infer -e models/yolov8s.engine -m 0 -i images -o output -l labels_det.txt --cudaGraph
+    # Oriented Bounding Box Model
+    trtyolo infer -e models/yolov8s-obb.engine -m 1 -i images -o output -l labels_obb.txt --cudaGraph
+    ```
 
-This command will generate visualized inference results in the `output` folder.
+    The inference results will be saved to the `output` folder and generate visualized results.
 
-### Inference using Python
+### Inference Using Python
 
-You can also write a script using the `tensorrt_yolo` library for inference. The script `detect.py` is already written for this purpose.
+1. Use the `tensorrt_yolo` library for Python inference. The sample script `detect.py` is ready for use.
+2. Run the following commands for inference:
 
-> To further speed up the inference process, use the `--cudaGraph` option, but this feature only supports static models, not dynamic models.
+    > [!NOTE] 
+    > The `--cudaGraph` option can further accelerate the inference process, but this feature only supports static models.
 
-```bash
-python detect.py -e models/yolov8s.engine -i images -o output -l labels.txt --cudaGraph
-```
+    ```bash
+    # Detection Model
+    python detect.py -e models/yolov8s.engine -m 0 -i images -o output -l labels_det.txt --cudaGraph
+    # Oriented Bounding Box Model
+    python detect.py -e models/yolov8s-obb.engine -m 1 -i images -o output -l labels_obb.txt --cudaGraph
+    ```
 
-### Inference using C++
+### Inference Using C++
 
-Before performing inference using C++, ensure that you have compiled Deploy as per the [Deploy Build Guide](../../docs/en/build_and_install.md#deploy-Build).
+1. Ensure that the project has been compiled according to the [Deploy Compilation Guide](../../docs/en/build_and_install.md#deploy-compilation).
+2. Use `xmake` to compile `detect.cpp` into an executable file:
 
-Next, compile `detect.cpp` into an executable using xmake:
+    ```bash
+    xmake f -P . --tensorrt="/path/to/your/TensorRT" --deploy="/path/to/your/TensorRT-YOLO"
 
-```bash
-xmake f -P . --tensorrt="/path/to/your/TensorRT" --deploy=/path/to/your/TensorRT-YOLO
+    xmake -P . -r
+    ```
 
-xmake -P . -r
-```
+    After compilation, the executable file will be generated in the `build` folder at the project root.
 
-After executing the above commands, an executable named `detect` will be generated in the `build` directory at the root. Finally, you can run the executable directly or use the `xmake run` command for inference. Use `--help` to see detailed command options:
+3. Run the following commands for inference:
 
-> To further speed up the inference process, use the `--cudaGraph` option, but this feature only supports static models, not dynamic models.
+    > [!NOTE] 
+    > The `--cudaGraph` option can further accelerate the inference process, but this feature only supports static models.
 
-```bash
-xmake run -P . detect -e models/yolov8s.engine -i images -o output -l labels.txt --cudaGraph
-```
+    ```bash
+    # Detection Model
+    xmake run -P . detect -e models/yolov8s.engine -m 0 -i images -o output -l labels_det.txt --cudaGraph
+    # Oriented Bounding Box Model
+    xmake run -P . detect -e models/yolov8s-obb.engine -m 1 -i images -o output -l labels_obb.txt --cudaGraph
+    ```
 
-The above are examples of how to perform model inference.
+By following the steps above, you can successfully complete model inference.

demo/detect/README.md

@@ -2,76 +2,120 @@
 
 # 模型推理示例
 
-在这个示例中，我们以 YOLOv8s 模型为例，演示了如何使用 CLI、Python 和 C++ 三种方式进行模型推理。
+本示例以 YOLOv8s 模型为例，展示如何使用 CLI、Python 和 C++ 三种方式进行模型推理。
+
+> [!IMPORTANT]  
+> 如果要使用 [EfficientRotatedNMS](../../plugin/efficientRotatedNMSPlugin)  插件推理 OBB 模型，请参考 [构建 TensorRT 自定义插件](../../docs/cn/build_trt_custom_plugin.md) 进行构建。
 
 ## 模型导出
 
-首先，从 [YOLOv8s](https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8s.pt) 下载 YOLOv8s 模型并保存到 `models` 文件夹中。
+### 检测模型 (Detection)
 
-然后，使用以下指令将模型导出为带有 [EfficientNMS](https://github.com/NVIDIA/TensorRT/tree/main/plugin/efficientNMSPlugin) 插件的 ONNX 格式：
+1. 下载 [YOLOv8s](https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8s.pt) 模型，并将其保存至 `models` 文件夹。
+2. 使用以下命令将模型导出为带有 [EfficientNMS](https://github.com/NVIDIA/TensorRT/tree/main/plugin/efficientNMSPlugin) 插件的 ONNX 格式：
 
-```bash
-trtyolo export -w models/yolov8s.pt -v yolov8 -o models
-```
+    ```bash
+    trtyolo export -w models/yolov8s.pt -v yolov8 -o models
+    ```
 
-执行以上命令后，将在 `models` 文件夹下生成名为 `yolov8s.onnx` 的文件。然后，使用 `trtexec` 工具将 ONNX 文件转换为 TensorRT engine：
+    运行上述命令后，`models` 文件夹中将生成 `yolov8s.onnx` 文件。接下来，使用 `trtexec` 工具将 ONNX 文件转换为 TensorRT 引擎：
 
-```bash
-trtexec --onnx=models/yolov8s.onnx --saveEngine=models/yolov8s.engine --fp16
-```
+    ```bash
+    trtexec --onnx=models/yolov8s.onnx --saveEngine=models/yolov8s.engine --fp16
+    ```
 
-接下来，我们将使用不同的方式进行模型推理。
+### 旋转边界框模型 (OBB)
 
-## 模型推理
+1. 下载 [YOLOv8s-obb](https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8s-obb.pt) 模型，并将其保存至 `models` 文件夹。
+2. 使用以下命令将模型导出为带有 [EfficientRotatedNMS](../../plugin/efficientRotatedNMSPlugin/) 插件的 ONNX 格式：
+
+    ```bash
+    trtyolo export -w models/yolov8s-obb.pt -v yolov8 -o models
+    ```
+
+    运行上述命令后，`models` 文件夹中将生成 `yolov8s-obb.onnx` 文件。接下来，使用 `trtexec` 工具将 ONNX 文件转换为 TensorRT 引擎：
+
+    ```bash
+    trtexec --onnx=models/yolov8s-obb.onnx --saveEngine=models/yolov8s-obb.engine --fp16
+    ```
+
+## 数据集准备
 
-下载 [coco128](https://ultralytics.com/assets/coco128.zip) 数据集，解压缩后，将 `coco128/images/train2017` 文件夹中的图像移动到 `images` 文件夹中以供推理使用。
+### 检测模型 (Detection)
+
+1. 下载 [coco128](https://ultralytics.com/assets/coco128.zip) 数据集。
+2. 解压后，将 `coco128/images/train2017` 文件夹中的图像移动到 `images` 文件夹，以供推理使用。
+
+### 旋转边界框模型 (OBB)
+
+1. 下载 [DOTA-v1.0](https://drive.google.com/file/d/1fwiTNqRRen09E-O9VSpcMV2e6_d4GGVK/view) 数据集。
+2. 解压后，将 `part1/images` 文件夹中的图像移动到 `images` 文件夹，以供推理使用。
+
+## 模型推理
 
 ### 使用 CLI 进行推理
 
-您可以使用 `tensorrt_yolo` 自带的命令行工具 `trtyolo` 进行推理。运行以下命令查看推理相关的帮助信息：
+1. 使用 `trtyolo` 命令行工具进行推理。运行以下命令查看帮助信息：
 
-```bash
-trtyolo infer --help
-```
+    ```bash
+    trtyolo infer --help
+    ```
 
-然后，执行以下命令进行推理：
+2. 运行以下命令进行推理：
 
-> 要进一步加速推理过程，请使用 `--cudaGraph` 指令，但此功能仅支持静态模型，不支持动态模型。（4.0之前不支持）
+    > [!NOTE] 
+    > 从 4.0 版本开始新增的 `--cudaGraph` 指令可以进一步加速推理过程，但该功能仅支持静态模型。
+    > 
+    > 从 4.2 版本开始，支持 OBB 模型推理，并新增 `-m, --mode` 指令，用于选择 Detection 还是 OBB 模型。
 
-```bash
-trtyolo infer -e models/yolov8s.engine -i images -o output -l labels.txt --cudaGraph
-```
+    ```bash
+    # 检测模型
+    trtyolo infer -e models/yolov8s.engine -m 0 -i images -o output -l labels_det.txt --cudaGraph
+    # 旋转边界框模型
+    trtyolo infer -e models/yolov8s-obb.engine -m 1 -i images -o output -l labels_obb.txt --cudaGraph
+    ```
 
-此命令将在 `output` 文件夹中生成可视化的推理结果。
+    推理结果将保存至 `output` 文件夹，并生成可视化结果。
 
 ### 使用 Python 进行推理
 
-也可以使用 `tensorrt_yolo` 库编写脚本进行推理，`detect.py` 是已经写好的脚本。
+1. 使用 `tensorrt_yolo` 库进行 Python 推理。示例脚本 `detect.py` 已准备好。
+2. 运行以下命令进行推理：
 
-> 要进一步加速推理过程，请使用 `--cudaGraph` 指令，但此功能仅支持静态模型，不支持动态模型。
+    > [!NOTE] 
+    > `--cudaGraph` 指令可以进一步加速推理过程，但该功能仅支持静态模型。
 
-```bash
-python detect.py -e models/yolov8s.engine -i images -o output -l labels.txt --cudaGraph
-```
+    ```bash
+    # 检测模型
+    python detect.py -e models/yolov8s.engine -m 0 -i images -o output -l labels_det.txt --cudaGraph
+    # 旋转边界框模型
+    python detect.py -e models/yolov8s-obb.engine -m 1 -i images -o output -l labels_obb.txt --cudaGraph
+    ```
 
 ### 使用 C++ 进行推理
 
-使用 C++ 进行推理前，请确保您已按照 [Deploy 编译指南](../../docs/cn/build_and_install.md#deploy-编译) 对 Deploy 进行了编译。
+1. 确保已按照 [Deploy 编译指南](../../docs/cn/build_and_install.md#deploy-编译) 对项目进行编译。
+2. 使用 xmake 将 `detect.cpp` 编译为可执行文件：
+
+    ```bash
+    xmake f -P . --tensorrt="/path/to/your/TensorRT" --deploy="/path/to/your/TensorRT-YOLO"
 
-接着，使用 xmake 将 `detect.cpp` 编译为可执行文件：
+    xmake -P . -r
+    ```
 
-```bash
-xmake f -P . --tensorrt="/path/to/your/TensorRT" --deploy=/path/to/your/TensorRT-YOLO
+    编译完成后，可执行文件将生成在项目根目录的 `build` 文件夹中。
 
-xmake -P . -r
-```
+3. 使用以下命令运行推理：
 
-在执行上述命令后，将在根目录的 `build` 目录下生成名为 `detect` 的可执行文件。最后，您可以直接运行可执行文件或使用 `xmake run` 命令进行推理。使用 `--help` 查看详细指令选项：
+    > [!NOTE] 
+    > `--cudaGraph` 指令可以进一步加速推理过程，但该功能仅支持静态模型。
 
-> 要进一步加速推理过程，请使用 `--cudaGraph` 指令，但此功能仅支持静态模型，不支持动态模型。
+    ```bash
+    # 检测模型
+    xmake run -P . detect -e models/yolov8s.engine -m 0 -i images -o output -l labels_det.txt --cudaGraph
+    # 旋转边界框模型
+    xmake run -P . detect -e models/yolov8s-obb.engine -m 1 -i images -o output -l labels_obb.txt --cudaGraph
+    ```
 
-```bash
-xmake run -P . detect -e models/yolov8s.engine -i images -o output -l labels.txt --cudaGraph
-```
+通过以上方式，您可以顺利完成模型推理。
 
-以上是进行模型推理的方法示例。