feat: YoloObbPostprocess plugin (GPU rotated NMS via ProbIoU) - #322
Open
triple-mu wants to merge 5 commits into
Open
feat: YoloObbPostprocess plugin (GPU rotated NMS via ProbIoU)#322triple-mu wants to merge 5 commits into
triple-mu wants to merge 5 commits into
Conversation
Raw uint8 image -> NCHW float input via a CUDA kernel (letterbox + RGB + /255), staged through a pinned host buffer for a fast async H2D. Opt-in, default-off; the CPU path is unchanged. CMake enables CUDA only when nvcc is found. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Opt-in custom plugin (libyolov8_plugins.so, BUILD_PLUGINS=ON) fusing YOLOv8 detection decode + NMS + top-k as an alternative to EfficientNMS_TRT; consumes the decoded boxes+scores and emits num_dets/bboxes/scores/labels. Plugin .so loaded via --plugin-lib / $YOLOV8_PLUGIN_LIB (Python ctypes + C++ dlopen); detect.cpp auto-detects raw vs End2End/plugin engines. export-det.py --plugin. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Adds the seg plugin (decode+NMS, gathers 32 mask coeffs; proto matmul stays on host) and extracts shared helpers into nms_common.cuh. NMS is parallel (EfficientNMS-style: one thread per candidate marks suppression vs higher-scoring overlaps; single-thread compaction gathers top-k) — det refactored onto the shared parallel suppression too. PostSeg.plugin + export-seg.py --plugin; engine.py finds the NMS node by attribute name anywhere in the graph; segment.cpp auto-detects. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Pose plugin attached to the raw ultralytics export via onnx_graphsurgeon (export-pose.py): the transposed head [1,A,C] is fed to YoloPosePostprocess, which does parallel NMS + top-k and gathers 17 keypoints in-engine. pose.cpp auto-detects plugin outputs by binding name. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
OBB plugin attached via onnx_graphsurgeon (export-obb.py); parallel NMS + top-k in-engine using ProbIoU (the closed-form Gaussian overlap ultralytics uses for OBB), emitting center boxes + angles. obb.cpp auto-detects plugin outputs. Completes the det/seg/pose/obb GPU postprocess plugin family. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
triple-mu
force-pushed
the
feat/obb-plugin
branch
from
June 10, 2026 09:15
e3fe651 to
d5cc25f
Compare
triple-mu
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
YoloObbPostprocess 插件（GPU 旋转 NMS via ProbIoU，自包含）
后处理插件家族最后一个。obb 经
export-obb.py（onnx_graphsurgeon）插桩：转置后的 raw[1,A,C]（C=4+nc+1：xc,yc,w,h | scores | angle）喂给YoloObbPostprocess，在引擎内做并行旋转 NMS + top-k，输出中心框 + 角度。obb.cpp按 binding 名自动识别并重建cv::RotatedRect。旋转重叠用 ProbIoU（协方差高斯闭式）—— ultralytics OBB 本身用的度量，比多边形相交更省、更贴近其语义。与
ultralytics.utils.metrics.batch_probiou逐例 5 位小数吻合（独立 CUDA 自测验证）。NMS 同样是 EfficientNMS 式并行（每候选一线程 + 单线程 compact）。验证：bus.jpg 上插件与 raw（cv 旋转 NMS）路径在分离目标一致；build.py 阈值覆盖生效；TRT 8.6/10.16/11.0 编译。完成 det/seg/pose/obb 四类 GPU 后处理插件 + GPU 前处理 kernel。
🤖 Generated with Claude Code