feat: YoloSegPostprocess plugin + shared parallel NMS helpers

.gitignore

@@ -143,3 +143,4 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+tensorrtx/

CMakeLists.txt

@@ -28,6 +28,26 @@ endif()
 find_package(CUDAToolkit REQUIRED)  # provides CUDA::cudart (no CUDA language needed)
 find_package(OpenCV REQUIRED)
 find_package(TensorRT REQUIRED)
+
+# Optional CUDA toolchain for the GPU preprocessing kernel (opt-in --gpu-preprocess).
+# If nvcc is available we enable the CUDA language and compile preprocess.cu into the
+# core lib; otherwise the GPU path compiles to a stub that throws (CPU path unaffected).
+include(CheckLanguage)
+if(NOT CMAKE_CUDA_COMPILER AND EXISTS /usr/local/cuda/bin/nvcc)
+    set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc)
+endif()
+check_language(CUDA)
+if(CMAKE_CUDA_COMPILER)
+    # Set before enable_language(CUDA), which would otherwise default it to "52".
+    if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
+        set(CMAKE_CUDA_ARCHITECTURES 75 86 89 90)  # Turing..Hopper; override with -D for others
+    endif()
+    enable_language(CUDA)
+    set(YOLOV8_GPU_PREPROCESS ON)
+    message(STATUS "GPU preprocessing: enabled (CUDA ${CMAKE_CUDA_COMPILER_VERSION}, arch ${CMAKE_CUDA_ARCHITECTURES})")
+else()
+    message(STATUS "GPU preprocessing: disabled (no CUDA compiler; --gpu-preprocess will be unavailable)")
+endif()
 print_var(TensorRT_VERSION_STRING)
 print_var(OpenCV_VERSION)
 
@@ -71,6 +91,15 @@ endif()
 # them natively on the device (or cross-compile) with -DTensorRT_ROOT pointing at
 # the aarch64 TensorRT. No separate Jetson sources are needed.
 
+# Custom TensorRT postprocess plugins (libyolov8_plugins.so) — opt-in, needs CUDA.
+option(BUILD_PLUGINS "Build the custom TensorRT postprocess plugins" OFF)
+if(BUILD_PLUGINS)
+    if(NOT YOLOV8_GPU_PREPROCESS)
+        message(FATAL_ERROR "BUILD_PLUGINS requires a CUDA compiler (none found at configure time)")
+    endif()
+    add_subdirectory(csrc/plugins)
+endif()
+
 # DeepStream parser plugin — needs the DeepStream SDK, so it is off by default.
 option(BUILD_DEEPSTREAM "Build the DeepStream parser plugin" OFF)
 if(BUILD_DEEPSTREAM)

README.md

@@ -118,9 +118,13 @@ python infer.py --task pose --backend pycuda --engine yolov8s-pose.engine --imgs
 cmake -S . -B build -DTensorRT_ROOT=/path/to/TensorRT
 cmake --build build -j
 export LD_LIBRARY_PATH=/path/to/TensorRT/lib:$LD_LIBRARY_PATH
-./build/bin/yolov8_detect yolov8s.engine data/bus.jpg --out-dir output   # --show / --profile / --labels
+./build/bin/yolov8_detect yolov8s.engine data/bus.jpg --out-dir output   # --show / --profile / --labels / --gpu-preprocess
 ```
 
+`--gpu-preprocess` runs letterbox + normalize as a CUDA kernel (raw uint8 image straight
+to the network input, no CPU/OpenCV step); it is opt-in and built only when a CUDA
+compiler is found. Results match the CPU path within tolerance (GPU bilinear ≠ `cv::resize`).
+
 Build details, multiple TensorRT/OpenCV versions, cuDNN for TensorRT 8 and the C++14 fallback are in **[docs/Build.md](docs/Build.md)**. Class names live in `data/labels/*.txt` (override with `--labels`).
 
 ## Performance

README.zh-CN.md

@@ -115,9 +115,11 @@ python infer.py --task pose --backend pycuda --engine yolov8s-pose.engine --imgs
 cmake -S . -B build -DTensorRT_ROOT=/path/to/TensorRT
 cmake --build build -j
 export LD_LIBRARY_PATH=/path/to/TensorRT/lib:$LD_LIBRARY_PATH
-./build/bin/yolov8_detect yolov8s.engine data/bus.jpg --out-dir output   # --show / --profile / --labels
+./build/bin/yolov8_detect yolov8s.engine data/bus.jpg --out-dir output   # --show / --profile / --labels / --gpu-preprocess
 ```
 
+`--gpu-preprocess` 用 CUDA kernel 做 letterbox + 归一化（原始 uint8 图直接变成网络输入，无 CPU/OpenCV 步骤）；opt-in，仅在找到 CUDA 编译器时构建。结果与 CPU 路径在容差内一致（GPU 双线性 ≠ `cv::resize`）。
+
 构建细节、多版本 TensorRT/OpenCV、TensorRT 8 的 cuDNN 依赖与 C++14 回退见 **[docs/Build.md](docs/Build.md)**。类别名在 `data/labels/*.txt`（用 `--labels` 覆盖）。
 
 ## 性能

csrc/apps/cls.cpp

@@ -59,6 +59,11 @@ class ClsEngine: public Engine {
         resize_blob(image, blob, config_.input_size, pparam_);
     }
 
+    bool letterbox_preproc() const override
+    {
+        return false;  // cls uses plain resize on the GPU path too
+    }
+
 private:
     std::vector<std::string> class_names_;
 };

csrc/apps/detect.cpp

@@ -11,7 +11,10 @@
 
 using namespace yolov8;
 
-// YOLOv8 detection: a single output [1, 4+num_classes, num_anchors].
+// YOLOv8 detection. Auto-detects the engine kind from its outputs: a single
+// [1, 4+num_classes, num_anchors] tensor → raw decode + cv::dnn NMS; four tensors
+// (num_dets, bboxes, scores, labels) → NMS already baked in (EfficientNMS_TRT or the
+// YoloDetPostprocess plugin), so we only rescale. One binary handles all three.
 class DetectEngine: public Engine {
 public:
     DetectEngine(const std::string& engine_path, const InferConfig& config): Engine(engine_path, config)
@@ -22,6 +25,48 @@ class DetectEngine: public Engine {
     void postprocess(std::vector<Object>& objs) override
     {
         objs.clear();
+        if (output_bindings_.size() >= 4) {
+            postprocess_end2end(objs);
+        }
+        else {
+            postprocess_raw(objs);
+        }
+    }
+
+    void draw(const cv::Mat& image, cv::Mat& res, const std::vector<Object>& objs) const override
+    {
+        draw_detections(image, res, objs, class_names_);
+    }
+
+private:
+    // NMS baked into the engine: outputs are num_dets, bboxes, scores, labels.
+    void postprocess_end2end(std::vector<Object>& objs)
+    {
+        const int*   num_dets = static_cast<const int*>(host_ptrs_[0]);
+        const float* boxes    = static_cast<const float*>(host_ptrs_[1]);
+        const float* scores   = static_cast<const float*>(host_ptrs_[2]);
+        const int*   labels   = static_cast<const int*>(host_ptrs_[3]);
+        const float  dw = pparam_.dw, dh = pparam_.dh;
+        const float  width = pparam_.width, height = pparam_.height, ratio = pparam_.ratio;
+
+        for (int i = 0; i < num_dets[0]; ++i) {
+            const float* ptr = boxes + i * 4;
+            const float  x0  = clamp((ptr[0] - dw) * ratio, 0.f, width);
+            const float  y0  = clamp((ptr[1] - dh) * ratio, 0.f, height);
+            const float  x1  = clamp((ptr[2] - dw) * ratio, 0.f, width);
+            const float  y1  = clamp((ptr[3] - dh) * ratio, 0.f, height);
+
+            Object obj;
+            obj.rect  = cv::Rect_<float>(x0, y0, x1 - x0, y1 - y0);
+            obj.prob  = scores[i];
+            obj.label = labels[i];
+            objs.push_back(obj);
+        }
+    }
+
+    // Raw head output [1, 4+num_classes, num_anchors]: decode + NMS on the host.
+    void postprocess_raw(std::vector<Object>& objs)
+    {
         const int   num_channels = output_bindings_[0].dims.d[1];
         const int   num_anchors  = output_bindings_[0].dims.d[2];
         const int   num_labels   = num_channels - 4;
@@ -78,12 +123,6 @@ class DetectEngine: public Engine {
         }
     }
 
-    void draw(const cv::Mat& image, cv::Mat& res, const std::vector<Object>& objs) const override
-    {
-        draw_detections(image, res, objs, class_names_);
-    }
-
-private:
     std::vector<std::string> class_names_;
 };
 

csrc/apps/segment.cpp

@@ -11,8 +11,10 @@
 
 using namespace yolov8;
 
-// YOLOv8 instance segmentation. Two outputs: detections [1, 4+nc+seg_channels, anchors]
-// and mask prototypes [1, seg_channels, seg_h, seg_w]. Binding order is detected at runtime.
+// YOLOv8 instance segmentation. Auto-detects the engine kind: raw outputs
+// (detections [1, 4+nc+seg_channels, anchors] + prototypes) → decode + NMS on the host;
+// YoloSegPostprocess plugin outputs (num_dets/bboxes/scores/labels/mask_coeffs + proto)
+// → NMS already done, gather coeffs only. Mask assembly is shared by both.
 class SegEngine: public Engine {
 public:
     SegEngine(const std::string& engine_path, const InferConfig& config): Engine(engine_path, config)
@@ -23,9 +25,98 @@ class SegEngine: public Engine {
     void postprocess(std::vector<Object>& objs) override
     {
         objs.clear();
+        if (binding_index("num_dets") >= 0) {
+            postprocess_plugin(objs);
+        }
+        else {
+            postprocess_raw(objs);
+        }
+    }
+
+    void draw(const cv::Mat& image, cv::Mat& res, const std::vector<Object>& objs) const override
+    {
+        draw_segments(image, res, objs, class_names_);
+    }
+
+private:
+    // Index of an output binding by name, or -1.
+    int binding_index(const std::string& name) const
+    {
+        for (size_t i = 0; i < output_bindings_.size(); ++i) {
+            if (output_bindings_[i].name == name) {
+                return static_cast<int>(i);
+            }
+        }
+        return -1;
+    }
+
+    // matmul kept coefficients with the prototypes, sigmoid, crop the padded region,
+    // resize to the original image and threshold. Shared by both postprocess paths.
+    void assemble_masks(std::vector<Object>& objs, const cv::Mat& masks, const cv::Mat& protos)
+    {
+        if (masks.empty()) {
+            return;
+        }
+        const int   seg_h = config_.seg_h, seg_w = config_.seg_w;
+        const int   input_h = input_bindings_[0].dims.d[2];
+        const int   input_w = input_bindings_[0].dims.d[3];
+        const float dw = pparam_.dw, dh = pparam_.dh;
+        const float width = pparam_.width, height = pparam_.height;
+
+        cv::Mat matmul   = (masks * protos).t();
+        cv::Mat mask_mat = matmul.reshape(static_cast<int>(objs.size()), {seg_h, seg_w});
+
+        std::vector<cv::Mat> mask_channels;
+        cv::split(mask_mat, mask_channels);
+        const int      scale_dw = dw / input_w * seg_w;
+        const int      scale_dh = dh / input_h * seg_h;
+        const cv::Rect roi(scale_dw, scale_dh, seg_w - 2 * scale_dw, seg_h - 2 * scale_dh);
+
+        for (size_t i = 0; i < objs.size(); ++i) {
+            cv::Mat dest, mask;
+            cv::exp(-mask_channels[i], dest);
+            dest = 1.0 / (1.0 + dest);
+            dest = dest(roi);
+            cv::resize(dest, mask, cv::Size(static_cast<int>(width), static_cast<int>(height)), cv::INTER_LINEAR);
+            objs[i].boxMask = mask(objs[i].rect) > 0.5f;
+        }
+    }
+
+    // Plugin engine: NMS + coeff gather done in-engine; only rescale + assemble masks.
+    void postprocess_plugin(std::vector<Object>& objs)
+    {
+        const int    seg_channels = config_.seg_channels;
+        const int*   num_dets     = static_cast<const int*>(host_ptrs_[binding_index("num_dets")]);
+        const float* boxes        = static_cast<const float*>(host_ptrs_[binding_index("bboxes")]);
+        const float* scores       = static_cast<const float*>(host_ptrs_[binding_index("scores")]);
+        const int*   labels       = static_cast<const int*>(host_ptrs_[binding_index("labels")]);
+        const float* coeffs       = static_cast<const float*>(host_ptrs_[binding_index("mask_coeffs")]);
+        cv::Mat      protos(seg_channels, config_.seg_h * config_.seg_w, CV_32F, host_ptrs_[binding_index("proto")]);
+        const float  dw = pparam_.dw, dh = pparam_.dh;
+        const float  width = pparam_.width, height = pparam_.height, ratio = pparam_.ratio;
+
+        cv::Mat masks;
+        for (int i = 0; i < num_dets[0]; ++i) {
+            const float* ptr = boxes + i * 4;
+            const float  x0  = clamp((ptr[0] - dw) * ratio, 0.f, width);
+            const float  y0  = clamp((ptr[1] - dh) * ratio, 0.f, height);
+            const float  x1  = clamp((ptr[2] - dw) * ratio, 0.f, width);
+            const float  y1  = clamp((ptr[3] - dh) * ratio, 0.f, height);
+
+            Object obj;
+            obj.rect  = cv::Rect_<float>(x0, y0, x1 - x0, y1 - y0);
+            obj.prob  = scores[i];
+            obj.label = labels[i];
+            objs.push_back(obj);
+            masks.push_back(cv::Mat(1, seg_channels, CV_32F, const_cast<float*>(coeffs + i * seg_channels)));
+        }
+        assemble_masks(objs, masks, protos);
+    }
+
+    // Raw engine: detections [1, 4+nc+seg_channels, anchors] + prototypes; decode + NMS.
+    void postprocess_raw(std::vector<Object>& objs)
+    {
         const int seg_channels = config_.seg_channels, seg_h = config_.seg_h, seg_w = config_.seg_w;
-        const int input_h = input_bindings_[0].dims.d[2];
-        const int input_w = input_bindings_[0].dims.d[3];
 
         // Find the 3-D detection output; the other output holds the prototypes.
         int det_idx      = -1;
@@ -97,35 +188,9 @@ class SegEngine: public Engine {
             objs.push_back(obj);
             ++cnt;
         }
-
-        if (masks.empty()) {
-            return;
-        }
-        cv::Mat matmul   = (masks * protos).t();
-        cv::Mat mask_mat = matmul.reshape(static_cast<int>(objs.size()), {seg_h, seg_w});
-
-        std::vector<cv::Mat> mask_channels;
-        cv::split(mask_mat, mask_channels);
-        const int      scale_dw = dw / input_w * seg_w;
-        const int      scale_dh = dh / input_h * seg_h;
-        const cv::Rect roi(scale_dw, scale_dh, seg_w - 2 * scale_dw, seg_h - 2 * scale_dh);
-
-        for (size_t i = 0; i < objs.size(); ++i) {
-            cv::Mat dest, mask;
-            cv::exp(-mask_channels[i], dest);
-            dest = 1.0 / (1.0 + dest);
-            dest = dest(roi);
-            cv::resize(dest, mask, cv::Size(static_cast<int>(width), static_cast<int>(height)), cv::INTER_LINEAR);
-            objs[i].boxMask = mask(objs[i].rect) > 0.5f;
-        }
+        assemble_masks(objs, masks, protos);
     }
 
-    void draw(const cv::Mat& image, cv::Mat& res, const std::vector<Object>& objs) const override
-    {
-        draw_segments(image, res, objs, class_names_);
-    }
-
-private:
     std::vector<std::string> class_names_;
 };
 

csrc/core/CMakeLists.txt

@@ -24,12 +24,22 @@ target_link_libraries(yolov8_core PUBLIC
     CUDA::cudart
     ${OpenCV_LIBS}
     TensorRT::TensorRT
+    ${CMAKE_DL_LIBS}  # dlopen for custom plugin .so
 )
 
+# GPU preprocessing kernel (opt-in). Compiled into the core lib only when the root
+# build found a CUDA compiler; otherwise the engine's --gpu-preprocess branch throws.
+if(YOLOV8_GPU_PREPROCESS)
+    target_sources(yolov8_core PRIVATE src/preprocess.cu)
+    target_compile_definitions(yolov8_core PUBLIC YOLOV8_GPU_PREPROCESS)
+    set_target_properties(yolov8_core PROPERTIES CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES}")
+endif()
+
 # Require at least C++14 (the ghc::filesystem floor); the project default is C++17
 # (std::filesystem). Build with -DCMAKE_CXX_STANDARD=14 to exercise the ghc fallback.
 target_compile_features(yolov8_core PUBLIC cxx_std_14)
-target_compile_options(yolov8_core PRIVATE -Wall -Wextra)
+# Warning flags only for C++ TUs (nvcc rejects -Wall/-Wextra passed directly).
+target_compile_options(yolov8_core PRIVATE $<$<COMPILE_LANGUAGE:CXX>:-Wall;-Wextra>)
 
 # TRT_10 / BATCHED_NMS are PUBLIC so task executables compile the same code paths.
 yolov8_apply_compile_defs(yolov8_core)

csrc/core/include/yolov8/config.hpp

@@ -15,10 +15,12 @@ struct InferConfig {
     int         seg_h        = 160;
     int         seg_w        = 160;
     std::string labels_path;  // empty -> built-in COCO names
-    bool        warmup  = true;
-    bool        show    = false;  // off by default (works headless); otherwise save
-    bool        profile = false;  // attach a per-layer IProfiler and print a report
-    std::string out_dir = "output";
+    std::string plugin_lib;   // custom plugin .so to dlopen (else $YOLOV8_PLUGIN_LIB)
+    bool        warmup         = true;
+    bool        show           = false;  // off by default (works headless); otherwise save
+    bool        profile        = false;  // attach a per-layer IProfiler and print a report
+    bool        gpu_preprocess = false;  // preprocess on GPU (CUDA kernel) instead of CPU/OpenCV
+    std::string out_dir        = "output";
 };
 
 struct CliArgs {

csrc/core/include/yolov8/engine.hpp

@@ -38,8 +38,13 @@ class Engine {
     void print_profile() const;  // prints the per-layer report if --profile was set
 
 protected:
-    // Preprocess one image into an NCHW blob. Default is letterbox; cls overrides it.
+    // Preprocess one image into an NCHW blob (CPU path). Default is letterbox; cls overrides it.
     virtual void preprocess(const cv::Mat& image, cv::Mat& blob);
+    // Whether the GPU preprocess path should letterbox (true) or plain-resize (false, cls).
+    virtual bool letterbox_preproc() const
+    {
+        return true;
+    }
 
     InferConfig          config_;
     std::vector<Binding> input_bindings_;
@@ -53,7 +58,9 @@ class Engine {
     // Declared so destruction runs context -> engine -> runtime -> stream -> buffers.
     std::vector<DeviceBuffer>                 device_buffers_;
     std::vector<HostPinnedBuffer>             host_buffers_;
-    std::vector<void*>                        device_ptrs_;  // [inputs..., outputs...]
+    std::vector<void*>                        device_ptrs_;     // [inputs..., outputs...]
+    DeviceBuffer                              raw_input_;       // uint8 device buffer for --gpu-preprocess
+    HostPinnedBuffer                          raw_input_host_;  // pinned staging for a fast async H2D upload
     CudaStream                                stream_;
     TrtUniquePtr<nvinfer1::IRuntime>          runtime_;
     TrtUniquePtr<nvinfer1::ICudaEngine>       engine_;

csrc/core/include/yolov8/preprocess.hpp

@@ -11,4 +11,14 @@ void letterbox(const cv::Mat& image, cv::Mat& out, const cv::Size& size, PrePara
 // Plain resize (no padding) into an NCHW float blob (RGB, /255). Used for classification.
 void resize_blob(const cv::Mat& image, cv::Mat& out, const cv::Size& size, PreParam& pparam);
 
+// GPU counterparts (defined in preprocess.cu, compiled only when CUDA is available).
+// `src` is a device-side uint8 HWC BGR image; `dst` is the device NCHW float input
+// buffer (RGB, /255) written directly. `stream` is a cudaStream_t passed as void* to
+// keep this header free of CUDA headers. Records scale/pad in `pparam` exactly like
+// the CPU versions above.
+void letterbox_cuda(
+    const unsigned char* src, int src_w, int src_h, float* dst, int dst_w, int dst_h, PreParam& pparam, void* stream);
+void resize_cuda(
+    const unsigned char* src, int src_w, int src_h, float* dst, int dst_w, int dst_h, PreParam& pparam, void* stream);
+
 }  // namespace yolov8