BiRefNet C++ TensorRT （二分类图像分割）
利用TensorRT和CUDA的双边参考网络（BiRefNet）的高性能c++实现，针对实时高分辨率二分类图像分割进行了优化。
BiRefNet c++ TENSORRT旨在有效地在GPU上运行双边参考分割任务。通过利用TensorRT的优化和CUDA内核，它旨在以最小的延迟提供最先进的性能。
主要特性
- TensorRT 加速：使用序列化的 TRT 引擎加速分割任务的推理过程。
- CUDA 集成：全面基于 GPU 的预处理、后处理及内存管理。
- 高分辨率支持：开箱即用，能够处理高分辨率图像（例如 1024x1024）。
- 易于集成：采用 C++17 编码，便于集成到现有的流程中。
📢 最新动态
- 增强的双边参考：通过利用双重参考指导，提升二分法分割输出的质量。
- 改进的内存占用：针对大批次或高分辨率工作负载优化了 GPU 分配。
- 可配置精度：支持 FP16 或 FP32 模式（需要支持半精度的 GPU）。
📂 项目结构
BiRefNet/
├── include
│ └── birefnet.h # Main BiRefNet class definition
├── src
│ └── birefnet.cpp # Implementation of the BiRefNet class
├── CMakeLists.txt # CMake configuration
└── main.cpp #
-
include/birefnet.h
头文件，定义了BiRefNet类，该类负责 TensorRT 引擎的创建、执行和内存缓冲区管理。 -
src/birefnet.cpp
源代码实现，包括加载序列化引擎、运行推理及处理输出后处理。 -
CMakeLists.txt
使用 CMake 构建项目的配置文件。根据需要调整 TensorRT、CUDA 和 OpenCV 的路径。 -
main.cpp
一个最小示例，演示了如何加载模型、对图像或视频进行推理并将结果保存。
🚀 安装步骤
-
克隆仓库
git clone https://github.com/hamdiboukamcha/BiRefNet-Cpp-TensorRT.git cd BiRefNet-Cpp-TensorRT mkdir build && cd build cmake .. make -j$(nproc)
📦 依赖项
-
CUDA
用于 GPU 加速和内核启动（例如，CUDA 11.x 或更新版本）。 -
TensorRT
高性能深度学习推理库（推荐使用 v8.x 或更新版本）。 -
OpenCV
用于图像加载、预处理和基本可视化。 -
C++17
本项目使用现代 C++ 特性。请确保您的编译器支持 C++17 或更高版本。
🔍 代码概览
主要组件
-
BiRefNet 类
- 从给定的引擎/模型路径初始化 TensorRT 引擎。
- 处理预处理（图像大小调整、均值/标准差归一化等）。
- 运行推理并将输出后处理为分割图。
- 管理 CUDA 资源和流。
-
Logger 类（在 main.cpp 中）
- 实现 TensorRT 的 ILogger 接口以进行自定义日志记录。
主要函数
-
BiRefNet::BiRefNet(…)
构造函数，将 .trt（序列化的 TensorRT）引擎加载到内存中。 -
BiRefNet::predict(cv::Mat& image)
主要推理函数：接受一个 OpenCV cv::Mat 作为输入，返回分割结果作为 cv::Mat。 -
BiRefNet::preprocess(…)
将图像转换为归一化的浮点数（均值减法、除以标准差等）。 -
BiRefNet::postprocess(…)
将原始输出重塑为有意义的图像数据，通常是一个 8 位或 32 位矩阵用于分割。
🎬 使用方法
准备TensorRT推理引擎：
-
将模型转换为 ONNX 格式，然后构建一个 TensorRT 引擎（例如，使用 trtexec 或自定义构建工具）。选择使用BiRefNet的预训练模型。
trtexec --onnx=BiRefNet-general-resolution_512x512-fp16-epoch_216.onnx --saveEngine=BiRefNet-tiny.engine --fp16
-
推理结果
./BiRefNetTRT Loading model from BiRefNet-tiny.engine... The model has been successfully loaded! Folder already exists! [Info] Processing time: 2452 ms for 1693479941020.jpg [Info] Saving result to: results/1693479941020.jpg [Info] Processing finished.
