Merge pull request #168 from laugh12321/dev

docs: update documentation for changes

Author: laugh12321

README.en.md

@@ -38,14 +38,16 @@ English | [简体中文](README.md)
 
 ## <div align="center">🌠 Recent updates</div>
 
+- 2025-04-19: Added support for [YOLO-World](https://docs.ultralytics.com/zh/models/yolo-world/) and [YOLOE](https://docs.ultralytics.com/zh/models/yoloe/), including classification, oriented bounding boxes, pose estimation, and instance segmentation. See [Bilibili](https://www.bilibili.com/video/BV12N5bzkENV) for details. 🌟 NEW
+
 - 2025-03-29: Added support for [YOLO12](https://github.com/sunsmarterjie/yolov12), including classification, oriented bounding boxes, pose estimation, and instance segmentation. See [issues](https://github.com/sunsmarterjie/yolov12/issues/22) for details. 🌟 NEW
 
 - [Performance Leap! TensorRT-YOLO 6.0: Comprehensive Upgrade Analysis and Practical Guide](https://medium.com/@laugh12321/performance-leap-tensorrt-yolo-6-0-comprehensive-upgrade-analysis-and-practical-guide-9d19ad3b53f9) 🌟 NEW
 
 ## <div align="center">✨ Key Features</div>
 
 ### 🎯 Diverse YOLO Support
-- **Comprehensive Compatibility**: Supports YOLOv3 to YOLOv11 series models, as well as PP-YOLOE and PP-YOLOE+, meeting diverse needs.
+- **Comprehensive Compatibility**: Supports YOLOv3 to YOLO12 series models, as well as PP-YOLOE, PP-YOLOE+, YOLO-World, and YOLOE, meeting diverse needs. See [🖥️ Supported Models List](#support-models) for details.
 - **Flexible Switching**: Provides simple and easy-to-use interfaces for quick switching between different YOLO versions. 🌟 NEW
 - **Multi-Scenario Applications**: Offers rich example codes covering [Detect](examples/detect/), [Segment](examples/segment/), [Classify](examples/classify/), [Pose](examples/pose/), [OBB](examples/obb/), and more.
 
@@ -143,13 +145,13 @@ English | [简体中文](README.md)
 
   ```python
   import cv2
-  from tensorrtyolo.infer import InferOption, DetectModel, generatelabels, visualize
+  from tensorrt_yolo.infer import InferOption, DetectModel, generate_labels, visualize
 
   def main():
       # -------------------- Initialization --------------------
       # Configure inference settings
       option = InferOption()
-      option.enableswaprb()  # Convert OpenCV's default BGR format to RGB
+      option.enable_swap_rb()  # Convert OpenCV's default BGR format to RGB
       # Special model configuration example (uncomment for PP-YOLOE series)
       # option.setnormalizeparams([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
 
@@ -161,34 +163,34 @@ English | [简体中文](README.md)
 
       # -------------------- Data Preprocessing --------------------
       # Load test image (add file existence check)
-      inputimg = cv2.imread("testimage.jpg")
+      input_img = cv2.imread("test_image.jpg")
       if input_img is None:
           raise FileNotFoundError("Failed to load test image. Check the file path.")
 
       # -------------------- Inference Execution --------------------
       # Perform object detection (returns bounding boxes, confidence scores, and class labels)
-      detectionresult = model.predict(inputimg)
-      print(f"==> Detection Result: {detection_result}")
+      detection_result = model.predict(input_img)
+      print(f"==> detection_result: {detection_result}")
 
       # -------------------- Result Visualization --------------------
       # Load class labels (ensure labels.txt matches the model)
-      classlabels = generate_labels(labelsfile="labels.txt")
+      class_labels = generate_labels(labels_file="labels.txt")
       # Generate visualized result
       visualized_img = visualize(
           image=input_img,
           result=detection_result,
           labels=class_labels,
       )
-      cv2.imwrite("visimage.jpg", visualizedimg)
+      cv2.imwrite("vis_image.jpg", visualized_img)
 
       # -------------------- Model Cloning Demo --------------------
       # Clone model instance (for multi-threaded scenarios)
       cloned_model = model.clone()  # Create an independent copy to avoid resource contention
       # Verify cloned model inference consistency
-      clonedresult = cloned_model.predict(inputimg)
-      print(f"==> Cloned Result: {cloned_result}")
+      cloned_result = cloned_model.predict(input_img)
+      print(f"==> cloned_result: {cloned_result}")
 
-  if _name__ == "__main_":
+  if __name__ == "__main__":
       main()
   ```
 
@@ -221,7 +223,7 @@ English | [简体中文](README.md)
           );
 
           // -------------------- Data Loading --------------------
-          cv::Mat cvimage = cv::imread("testimage.jpg");
+          cv::Mat cv_image = cv::imread("test_image.jpg");
           if (cv_image.empty()) {
               throw std::runtime_error("Failed to load test image.");
           }
@@ -350,6 +352,18 @@ Symbol legend: (1)  ✅ : Supported; (2) ❔: In progress; (3) ❎ : Not support
       <td>✅</td>
       <td>✅</td>
     </tr>
+    <tr>
+      <td>Detect</td>
+      <td><a href="https://github.com/ultralytics/ultralytics">YOLO-World V2 (ultralytics)</a></td>
+      <td>✅</td>
+      <td>✅</td>
+    </tr>
+    <tr>
+      <td>Detect</td>
+      <td><a href="https://github.com/THU-MIG/yoloe">THU-MIG/yoloe</a></td>
+      <td>✅</td>
+      <td>✅</td>
+    </tr>
     <tr>
       <td>Detect</td>
       <td><a href="https://github.com/ultralytics/ultralytics">ultralytics/ultralytics</a></td>
@@ -392,6 +406,12 @@ Symbol legend: (1)  ✅ : Supported; (2) ❔: In progress; (3) ❎ : Not support
       <td>❎ Implement yourself referring to <a href="https://github.com/laugh12321/TensorRT-YOLO/blob/main/tensorrt_yolo/export/head.py">tensorrt_yolo/export/head.py</a></td>
       <td>🟢</td>
     </tr>
+    <tr>
+      <td>Segment</td>
+      <td><a href="https://github.com/THU-MIG/yoloe">THU-MIG/yoloe</a></td>
+      <td>✅</td>
+      <td>✅</td>
+    </tr>
     <tr>
       <td>Segment</td>
       <td><a href="https://github.com/ultralytics/ultralytics">ultralytics/ultralytics</a></td>

README.md

@@ -45,6 +45,7 @@
     ▸ **CUDA Graph 原理与工程实践**  
     ▸ **Triton Inference Server 部署技巧**  
 
+- 2025-04-19: 添加对 [YOLO-World](https://docs.ultralytics.com/zh/models/yolo-world/),  [YOLOE](https://docs.ultralytics.com/zh/models/yoloe/) 的支持，包括分类、定向边界框、姿态估计以及实例分割，详见 [B站](https://www.bilibili.com/video/BV12N5bzkENV)。🌟 NEW
 
 - 2025-03-29: 添加对 [YOLO12](https://github.com/sunsmarterjie/yolov12) 的支持，包括分类、定向边界框、姿态估计以及实例分割，详见 [issues](https://github.com/sunsmarterjie/yolov12/issues/22)。🌟 NEW
 
@@ -54,7 +55,7 @@
 ## <div align="center">✨ 主要特性</div>
 
 ### 🎯 多样化的 YOLO 支持
-- **全面兼容**：支持 YOLOv3 至 YOLOv11 全系列模型，以及 PP-YOLOE 和 PP-YOLOE+，满足多样化需求。
+- **全面兼容**：支持 YOLOv3 至 YOLO12 全系列模型，以及 PP-YOLOE、PP-YOLOE+、YOLO-World 和 YOLOE 等多种变体，满足多样化需求，详见 [🖥️ 模型支持列表](#support-models)。
 - **灵活切换**：提供简洁易用的接口，支持不同版本 YOLO 模型的快速切换。🌟 NEW
 - **多场景应用**：提供丰富的示例代码，涵盖[Detect](examples/detect/)、[Segment](examples/segment/)、[Classify](examples/classify/)、[Pose](examples/pose/)、[OBB](examples/obb/)等多种应用场景。
 
@@ -360,6 +361,18 @@
       <td>✅</td>
       <td>✅</td>
     </tr>
+    <tr>
+      <td>Detect</td>
+      <td><a href="https://github.com/ultralytics/ultralytics">YOLO-World V2 (ultralytics)</a></td>
+      <td>✅</td>
+      <td>✅</td>
+    </tr>
+    <tr>
+      <td>Detect</td>
+      <td><a href="https://github.com/THU-MIG/yoloe">THU-MIG/yoloe</a></td>
+      <td>✅</td>
+      <td>✅</td>
+    </tr>
     <tr>
       <td>Detect</td>
       <td><a href="https://github.com/ultralytics/ultralytics">ultralytics/ultralytics</a></td>
@@ -402,6 +415,12 @@
       <td>❎ 参考<a href="https://github.com/laugh12321/TensorRT-YOLO/blob/main/tensorrt_yolo/export/head.py">tensorrt_yolo/export/head.py</a> 自行实现</td>
       <td>🟢</td>
     </tr>
+    <tr>
+      <td>Segment</td>
+      <td><a href="https://github.com/THU-MIG/yoloe">THU-MIG/yoloe</a></td>
+      <td>✅</td>
+      <td>✅</td>
+    </tr>
     <tr>
       <td>Segment</td>
       <td><a href="https://github.com/ultralytics/ultralytics">ultralytics/ultralytics</a></td>

docs/cn/model_export.md

@@ -4,9 +4,26 @@
 
 ## 使用 `trtyolo` CLI 导出模型
 
-`tensorrt_yolo` 提供了一个便捷的命令行界面（CLI）工具 `trtyolo`，用于将 ONNX 模型导出为适用于项目推理的格式，并集成了 TensorRT 插件。您可以通过运行 `trtyolo export --help` 查看详细的导出命令帮助信息。
-
-> [!NOTE]  
+`tensorrt_yolo` 提供了一个便捷的命令行界面（CLI）工具 `trtyolo`，用于将模型导出为适用于项目推理的格式，并集成了 TensorRT 插件。您可以通过运行 `trtyolo export --help` 查看详细的导出命令帮助信息。
+
+### 参数说明
+- `-v, --version`: 模型版本。支持的选项包括 `yolov3`, `yolov5`, `yolov8`, `yolov10`, `yolo11`, `yolo12`, `yolo-world`, `yoloe`, `pp-yoloe`, `ultralytics`。
+- `-o, --output`: 导出模型保存的目录路径。
+- `-w, --weights`: PyTorch YOLO 权重文件路径（非 PP-YOLOE 模型必需）。
+- `--model_dir`: 包含 PaddleDetection PP-YOLOE 模型的目录路径（PP-YOLOE 模型必需）。
+- `--model_filename`: PaddleDetection PP-YOLOE 模型文件名（PP-YOLOE 模型必需）。
+- `--params_filename`: PaddleDetection PP-YOLOE 参数文件名（PP-YOLOE 模型必需）。
+- `-b, --batch`: 模型的总批量大小。使用 `-1` 表示动态批量大小，默认为 `1`。
+- `--max_boxes`: 每张图像的最大检测框数量，默认为 `100`。
+- `--iou_thres`: NMS 的 IoU 阈值，默认为 `0.45`。
+- `--conf_thres`: 目标检测的置信度阈值，默认为 `0.25`。
+- `--imgsz`: 图像大小（单个值表示正方形，或 "height,width"）。默认为 `640`（非 PP-YOLOE 模型）。
+- `--names`: 自定义类别名称（仅适用于 YOLO-World 和 YOLOE 模型，以逗号分隔，例如 "person,car,dog"）。
+- `--repo_dir`: 包含本地仓库的目录路径（仅适用于 YOLOv3 和 YOLOv5 模型，使用 `torch.hub.load` 时适用）。
+- `--opset`: ONNX opset 版本，默认为 `12`。
+- `-s, --simplify`: 是否简化导出的 ONNX 模型，默认为 `False`。
+
+> [!NOTE]
 > 在导出 PP-YOLOE 和 PP-YOLOE+ 的 ONNX 模型时，仅会调整 batch 维度，而 height 和 width 维度保持不变。您需要在 [PaddleDetection](https://github.com/PaddlePaddle/PaddleDetection) 中进行相关设置，默认值通常为 640。
 >
 > 官方仓库如 [YOLOv6](https://github.com/meituan/YOLOv6/tree/main/deploy/ONNX#tensorrt-backend-tensorrt-version-800)、[YOLOv7](https://github.com/WongKinYiu/yolov7#export)、[YOLOv9](https://github.com/WongKinYiu/yolov9/issues/130#issue-2162045461) 已经提供了带有 EfficientNMS 插件的 ONNX 模型导出功能，因此此处不再重复提供。
@@ -15,25 +32,31 @@
 
 ```bash
 # 导出远程仓库中的 YOLOv3 模型
-trtyolo export -w yolov3.pt -v yolov3 -o output
+trtyolo export -v yolov3 -w yolov3.pt -o output
+
+# 导出本地仓库中的 YOLOv5 Classify 模型
+trtyolo export -v yolov5 -w yolov5s-cls.pt -o output --repo_dir your_local_yolovs_repository
 
-# 导出本地仓库中的 YOLOv5 模型
-trtyolo export -w yolov5s.pt -v yolov5 -o output --repo_dir your_local_yolovs_repository
+# 使用 Ultralytics 训练的 YOLO 系列模型（YOLOv3、YOLOv5、YOLOv6、YOLOv8、YOLOv9、YOLOv10、YOLO11 等），并指定插件参数，以动态 batch 导出
+trtyolo export -v ultralytics -w yolov8s.pt -o output --max_boxes 100 --iou_thres 0.45 --conf_thres 0.25 -b -1
 
-# 使用 Ultralytics 训练的 YOLO 系列模型（YOLOv3、YOLOv5、YOLOv6、YOLOv8、YOLOv9、YOLOv10、YOLO11），并指定插件参数，以动态 batch 导出
-trtyolo export -w yolov8s.pt -v ultralytics -o output --max_boxes 100 --iou_thres 0.45 --conf_thres 0.25 -b -1
+# 导出 PP-YOLOE 和 PP-YOLOE+ 模型
+trtyolo export -v pp-yoloe --model_dir modeldir --model_filename model.pdmodel --params_filename model.pdiparams -o output
 
 # 导出 YOLOv10 模型，高度 1080，宽度 1920
-trtyolo export -w yolov10s.pt -v yolov10 -o output --imgsz 1080 1920
+trtyolo export -v yolov10 -w yolov10s.pt -o output --imgsz 1080,1920
 
 # 导出 YOLO11 OBB 模型
-trtyolo export -w yolo11n-obb.pt -v yolo11 -o output
+trtyolo export -v yolo11 -w yolo11n-obb.pt -o output
 
 # 导出 YOLO12 Segment 模型
-trtyolo export -w yolo12n-seg.pt -v yolo12 -o output
+trtyolo export -v yolo12 -w yolo12n-seg.pt -o output
 
-# 导出 PP-YOLOE 和 PP-YOLOE+ 模型
-trtyolo export --model_dir modeldir --model_filename model.pdmodel --params_filename model.pdiparams -o output
+# 导出 YOLO-World 模型，并自定义类别
+trtyolo export -v yolo-world -w yoloworld.pt -o output --names "person,car,dog"
+
+# 导出 YOLOE 模型
+trtyolo export -v yoloe -w yoloe.pt -o output
 ```
 
 ## 使用 `trtexec` 构建 TensorRT 引擎
@@ -54,5 +77,5 @@ trtexec --onnx=yolov8n-obb.onnx --saveEngine=yolov8n-obb.engine --fp16 --staticP
 trtexec --onnx=yolo11n-obb.onnx --saveEngine=yolo11n-obb.engine --fp16 --minShapes=images:1x3x640x640 --optShapes=images:4x3x640x640 --maxShapes=images:8x3x640x640 --staticPlugins=./lib/plugin/custom_plugins.dll --setPluginsToSerialize=./lib/plugin/custom_plugins.dll
 ```
 
-> [!NOTE]  
+> [!NOTE]
 > 在使用 `--staticPlugins` 和 `--setPluginsToSerialize` 参数构建包含自定义插件的动态模型时，如果遇到错误 `[E] Error[4]: IRuntime::deserializeCudaEngine: Error Code 4: API Usage Error (Cannot register the library as plugin creator of EfficientRotatedNMS_TRT exists already.)`，这通常意味着引擎构建已成功，但在加载并反序列化引擎时检测到插件重复注册。这种情况下，可以忽略该错误。

docs/en/model_export.md

@@ -6,34 +6,57 @@ English | [中文](../cn/model_export.md)
 
 `tensorrt_yolo` provides a convenient Command Line Interface (CLI) tool, `trtyolo`, for exporting ONNX models into formats suitable for inference in this project, with integrated TensorRT plugins. You can view detailed export command instructions by running `trtyolo export --help`.
 
-> [!NOTE]  
+### Parameter Descriptions
+- `-v, --version`: Model version. Options include `yolov3`, `yolov5`, `yolov8`, `yolov10`, `yolo11`, `yolo12`, `yolo-world`, `yoloe`, `pp-yoloe`, `ultralytics`.
+- `-o, --output`: Directory path to save the exported model.
+- `-w, --weights`: Path to PyTorch YOLO weights (required for non PP-YOLOE models).
+- `--model_dir`: Directory path containing the PaddleDetection PP-YOLOE model (required for PP-YOLOE).
+- `--model_filename`: Filename of the PaddleDetection PP-YOLOE model (required for PP-YOLOE).
+- `--params_filename`: Filename of the PaddleDetection PP-YOLOE parameters (required for PP-YOLOE).
+- `-b, --batch`: Total batch size for the model. Use `-1` for dynamic batch size. Defaults to `1`.
+- `--max_boxes`: Maximum number of detections per image. Defaults to `100`.
+- `--iou_thres`: NMS IoU threshold for post-processing. Defaults to `0.45`.
+- `--conf_thres`: Confidence threshold for object detection. Defaults to `0.25`.
+- `--imgsz`: Image size (single value for square or "height,width"). Defaults to `640` (for non PP-YOLOE models).
+- `--names`: Custom class names for YOLO-World and YOLOE (comma-separated, e.g., "person,car,dog"). Only applicable for YOLO-World and YOLOE models.
+- `--repo_dir`: Directory containing the local repository (if using `torch.hub.load`). Only applicable for YOLOv3 and YOLOv5 models.
+- `--opset`: ONNX opset version. Defaults to `12`.
+- `-s, --simplify`: Whether to simplify the exported ONNX model. Defaults to `False`.
+
+> [!NOTE]
 > When exporting ONNX models for PP-YOLOE and PP-YOLOE+, only the batch dimension is adjusted, while the height and width dimensions remain unchanged. You need to configure this in [PaddleDetection](https://github.com/PaddlePaddle/PaddleDetection), with the default value typically set to 640.
 >
 > Official repositories such as [YOLOv6](https://github.com/meituan/YOLOv6/tree/main/deploy/ONNX#tensorrt-backend-tensorrt-version-800), [YOLOv7](https://github.com/WongKinYiu/yolov7#export), and [YOLOv9](https://github.com/WongKinYiu/yolov9/issues/130#issue-2162045461) already provide ONNX model exports with the EfficientNMS plugin, so this functionality is not duplicated here.
 
 ### Export Command Examples
 
 ```bash
-# Export a YOLOv3 model from a remote repository
-trtyolo export -w yolov3.pt -v yolov3 -o output
+# Exporting a YOLOv3 model from a remote repository
+trtyolo export -v yolov3 -w yolov3.pt -o output
 
-# Export a YOLOv5 model from a local repository
-trtyolo export -w yolov5s.pt -v yolov5 -o output --repo_dir your_local_yolovs_repository
+# Exporting a YOLOv5 Classify model from a local repository
+trtyolo export -v yolov5 -w yolov5s-cls.pt -o output --repo_dir your_local_yolovs_repository
 
-# Export models trained with Ultralytics (YOLOv3, YOLOv5, YOLOv6, YOLOv8, YOLOv9, YOLOv10, YOLO11) with plugin parameters, using dynamic batch export
-trtyolo export -w yolov8s.pt -v ultralytics -o output --max_boxes 100 --iou_thres 0.45 --conf_thres 0.25 -b -1
+# Exporting Ultralytics-trained YOLO series models (YOLOv3, YOLOv5, YOLOv6, YOLOv8, YOLOv9, YOLOv10, YOLO11, etc.) with plugin parameters for dynamic batch export
+trtyolo export -v ultralytics -w yolov8s.pt -o output --max_boxes 100 --iou_thres 0.45 --conf_thres 0.25 -b -1
 
-# Export a YOLOv10 model with height 1080 and width 1920
-trtyolo export -w yolov10s.pt -v yolov10 -o output --imgsz 1080 1920
+# Exporting PP-YOLOE and PP-YOLOE+ models
+trtyolo export -v pp-yoloe --model_dir modeldir --model_filename model.pdmodel --params_filename model.pdiparams -o output
 
-# Export a YOLO11 OBB model
-trtyolo export -w yolo11n-obb.pt -v yolo11 -o output
+# Exporting a YOLOv10 model with height 1080 and width 1920
+trtyolo export -v yolov10 -w yolov10s.pt -o output --imgsz 1080,1920
 
-# Export a YOLO12 Segment model
-trtyolo export -w yolo12n-seg.pt -v yolo12 -o output
+# Exporting a YOLO11 OBB model
+trtyolo export -v yolo11 -w yolo11n-obb.pt -o output
 
-# Export PP-YOLOE and PP-YOLOE+ models
-trtyolo export --model_dir modeldir --model_filename model.pdmodel --params_filename model.pdiparams -o output
+# Exporting a YOLO12 Segment model
+trtyolo export -v yolo12 -w yolo12n-seg.pt -o output
+
+# Exporting a YOLO-World model with custom classes
+trtyolo export -v yolo-world -w yoloworld.pt -o output --names "person,car,dog"
+
+# Exporting a YOLOE model
+trtyolo export -v yoloe -w yoloe.pt -o output
 ```
 
 ## Building TensorRT Engines Using `trtexec`
@@ -54,5 +77,5 @@ trtexec --onnx=yolov8n-obb.onnx --saveEngine=yolov8n-obb.engine --fp16 --staticP
 trtexec --onnx=yolo11n-obb.onnx --saveEngine=yolo11n-obb.engine --fp16 --minShapes=images:1x3x640x640 --optShapes=images:4x3x640x640 --maxShapes=images:8x3x640x640 --staticPlugins=./lib/plugin/custom_plugins.dll --setPluginsToSerialize=./lib/plugin/custom_plugins.dll
 ```
 
-> [!NOTE]  
+> [!NOTE]
 > When building dynamic models with custom plugins using the `--staticPlugins` and `--setPluginsToSerialize` parameters, if you encounter the error `[E] Error[4]: IRuntime::deserializeCudaEngine: Error Code 4: API Usage Error (Cannot register the library as plugin creator of EfficientRotatedNMS_TRT exists already.)`, this typically indicates that the engine build was successful, but a plugin duplicate registration was detected during engine loading and deserialization. In such cases, this error can be safely ignored.