简介:本资源面向算法研发工程师、图像处理研究人员及深度学习C++部署实践者,提供Segment Anything模型在Windows平台的纯C++本地化部署方案,彻底规避Python环境依赖,适用于嵌入式集成、工业视觉系统开发及无Python运行时的生产环境部署。压缩包含2001个文件,主体为733个C++源码(.cpp)与307个头文件(.hpp),覆盖模型推理、图像预处理(如resize、color_lab、imgwarp)、后处理(connectedcomponents、colormap)等核心模块;辅以少量Python脚本用于模型转换与数据准备,整体体积676.44MB。已有357人学习下载,资源结构清晰、模块职责分明,配套完整构建流程与release版可执行程序生成指引,开箱即得可运行的图像分割掩膜生成能力,显著降低大模型C++工程化门槛。
1. 项目缘起:为什么要在C++环境中折腾Segment Anything?
最近在做一个工业视觉的POC项目,客户现场的环境限制非常严格:服务器没有外网权限,GPU资源有限(只有一张消费级的RTX 3060),并且明确要求推理服务必须通过C++接口集成到他们现有的、基于Qt框架的MES(制造执行系统)客户端里。当团队提出用Python+PyTorch的方案时,直接被否了,原因很简单——他们不想在成百上千台工控机上部署和维护一整套Python环境,担心依赖冲突和版本管理问题。目标很明确:需要一个高性能、轻量级、纯C++的视觉模型本地部署方案。
这时,Meta AI的Segment Anything Model(SAM)进入了视野。这个模型在零样本分割上的能力有目共睹,如果能把它“塞”进C++环境,很多需要快速原型验证或对实时性有要求的场景(比如交互式标注工具、边缘计算设备)就能直接受益。然而,翻遍官方仓库和社区,你会发现主流的玩法清一色是Python。想用C++直接调用?官方没提供,社区资料也零零散散。这恰恰是我想分享这个踩坑过程的原因——把前沿的AI模型从Python的舒适区“移植”到更贴近底层、更追求性能与集成的C++世界,本身就是一件充满挑战但极具价值的事。
这个过程涉及的不是简单的API调用,而是一整套技术栈的打通:从模型格式的转换、推理引擎的选型,到前后处理的C++实现,再到内存管理与性能优化。如果你也在面临类似的困境,比如需要将AI能力嵌入到C++桌面应用、移动端App,或者为嵌入式设备开发智能视觉模块,那么这篇从零到一的实践记录,或许能帮你避开我走过的弯路。
2. 核心工具链选型:ONNX Runtime为何是当前的最优解?
要实现SAM的C++部署,第一步也是最重要的一步,就是选择一个合适的推理运行时(Runtime)。这不是一个可以随意拍脑袋的决定,它直接决定了后续开发的复杂度、性能上限和系统兼容性。我评估了当时几个主流选项:
方案一:LibTorch (PyTorch C++)这是最“原教旨”的思路,因为SAM本身就是PyTorch模型。使用LibTorch意味着你可以几乎无损地加载原始.pth模型文件,理论上兼容性最好。但它的缺点同样致命:
- 体积庞大:一个完整的LibTorch库(带CUDA支持)轻松超过1GB,这对于需要分发给终端用户的应用程序来说是难以接受的。
- 依赖复杂:它依然依赖一套复杂的动态链接库,环境配置并不比Python简单多少。
- 优化局限:虽然LibTorch支持JIT,但其对模型图级的优化能力(如算子融合、常量折叠)不如专门的推理框架深入。
方案二:TensorRTNVIDIA的专属推理优化引擎,性能无疑是顶级的。但它锁死了NVIDIA GPU平台,失去了CPU部署的灵活性。并且,将PyTorch模型转换到TensorRT(通常通过ONNX中转)的过程,对于SAM这样包含动态控制流(如ViT中的注意力机制)和复杂后处理的模型,极易出错,调试成本极高。
方案三:ONNX Runtime这是一个由微软维护的开源跨平台推理引擎,支持多种硬件后端(CPU, CUDA, TensorRT, OpenVINO等)。它的核心优势在于:
- 格式统一:以ONNX(Open Neural Network Exchange)作为中间表示。只要模型能成功导出为ONNX格式,就可以被ONNX Runtime加载和推理。
- 轻量高效:C++库体积相对较小,且提供了丰富的图优化和硬件加速能力。
- 社区活跃:对Transformer类模型的支持较好,遇到问题容易找到解决方案。
综合来看,ONNX Runtime提供了最佳的平衡点:它既保持了跨平台和硬件的灵活性,又通过ONNX格式实现了与训练框架的解耦,同时其性能经过优化后完全可以满足生产要求。因此,我最终确定了“PyTorch -> ONNX -> ONNX Runtime (C++)”的技术路径。
注意:SAM模型包含两个部分:图像编码器(Image Encoder)和提示解码器(Prompt-guided Mask Decoder)。在部署时,一个常见的优化策略是将编码器和解码器分开导出和部署。因为一张图像的编码向量可以缓存并重复用于不同的提示输入(如多个点、框),这能极大提升交互式应用的响应速度。我们的部署也将采用这种分离式结构。
3. 从PyTorch到ONNX:模型导出的核心陷阱与解决之道
确定了ONNX Runtime这条路,接下来就是打通第一关:把PyTorch模型转换成ONNX格式。这听起来像是一行torch.onnx.export()的事,但实际操作中坑点密布。
3.1 环境准备与模型获取
首先,你需要一个配置好的Python环境,用于执行导出脚本。关键依赖如下:
pip install torch torchvision onnx onnxruntime # 以及SAM的官方库 pip install git+https://github.com/facebookresearch/segment-anything.git建议使用PyTorch 1.12+版本,其对ONNX导出的支持更完善。
SAM提供了多个预训练模型(ViT-H, ViT-L, ViT-B)。考虑到我们后续要在C++端和可能有限的GPU上运行,我选择了体积和计算量折中的vit_l模型。你可以从官方仓库提供的链接下载对应的.pth检查点文件。
3.2 分离式导出:编码器与解码器
如前所述,分别导出编码器和解码器是高效部署的关键。这里分享我编写的导出脚本核心部分,以及遇到的典型问题。
图像编码器导出:图像编码器是一个标准的Vision Transformer。导出它相对直接,但需要注意输入尺寸的动态性。
import torch import numpy as np from segment_anything import sam_model_registry from segment_anything.utils.onnx import SamOnnxModel # 加载模型 sam_checkpoint = "./models/sam_vit_l_0b3195.pth" model_type = "vit_l" sam = sam_model_registry[model_type](checkpoint=sam_checkpoint) # 使用SAM官方提供的ONNX转换工具类 onnx_model = SamOnnxModel(sam, return_single_mask=True) # 动态轴设置:批次(batch_size)和图像尺寸(高度、宽度)设为动态 dynamic_axes = { "image": {0: "batch_size", 2: "height", 3: "width"}, "image_embeddings": {0: "batch_size"}, } # 构造伪输入 image_size = 1024 # SAM编码器训练尺寸 dummy_image = torch.randn(1, 3, image_size, image_size, dtype=torch.float32) embedding_size = sam.prompt_encoder.image_embedding_size # 通常是64x64 dummy_embeddings = torch.randn(1, 256, embedding_size, embedding_size, dtype=torch.float32) # 导出编码器 torch.onnx.export( onnx_model.image_encoder, dummy_image, "./models/sam_image_encoder.onnx", input_names=["image"], output_names=["image_embeddings"], dynamic_axes=dynamic_axes, opset_version=17, # 使用较高的opset以获得更好的算子支持 do_constant_folding=True, )这段代码利用了SAM仓库中提供的SamOnnxModel工具类,它已经处理了模型内部的一些包装逻辑。关键点在于dynamic_axes的设置,它允许导出的ONNX模型接受不同尺寸的输入图像,这对于实际应用至关重要。
提示解码器导出:解码器的导出更为复杂,因为它涉及多模态输入:图像嵌入、点坐标、点标签、框坐标,以及可选的掩模输入。
# 准备解码器的伪输入 input_points = torch.tensor([[[500, 400]]], dtype=torch.float32) # 一个前景点 input_labels = torch.tensor([[1]], dtype=torch.int64) # 标签1表示前景 input_box = torch.tensor([[100, 100, 900, 900]], dtype=torch.float32) # 可选框 [x1, y1, x2, y2] has_mask_input = torch.tensor([[0]], dtype=torch.float32) # 初始掩模输入(全零) orig_im_size = torch.tensor([image_size, image_size], dtype=torch.float32) # 原始图像尺寸 # 解码器输入字典 decoder_inputs = { "image_embeddings": dummy_embeddings, "point_coords": input_points, "point_labels": input_labels, "mask_input": has_mask_input, "has_mask_input": has_mask_input, "orig_im_size": orig_im_size, } # 如果想支持框输入,需要额外处理 # decoder_inputs["box"] = input_box # 动态轴设置(解码器) decoder_dynamic_axes = { "image_embeddings": {0: "batch_size"}, "point_coords": {0: "batch_size", 1: "num_points"}, "point_labels": {0: "batch_size", 1: "num_points"}, "mask_input": {0: "batch_size"}, "has_mask_input": {0: "batch_size"}, "orig_im_size": {0: "batch_size"}, } # 导出解码器(注意:这里导出的是onnx_model.mask_decoder) torch.onnx.export( onnx_model.mask_decoder, (decoder_inputs,), "./models/sam_mask_decoder.onnx", input_names=list(decoder_inputs.keys()), output_names=["masks", "iou_predictions", "low_res_masks"], dynamic_axes=decoder_dynamic_axes, opset_version=17, do_constant_folding=True, )3.3 导出过程中的“坑”与填坑实录
- opset_version 不兼容:最初使用opset 11,导出时警告某些算子不支持。解决方案:升级到opset 16或17。ONNX opset版本定义了可用的算子集,版本越高,对PyTorch新特性的支持越好。
- 动态尺寸与
orig_im_size:解码器需要orig_im_size参数来将输出的掩模上采样回原始图像尺寸。这个参数必须与输入图像的原始高宽一致,且需要以(H, W)的格式传入。在C++端预处理时就必须计算好。 mask_input与has_mask_input的混淆:这是最大的一个坑。SAM解码器有两个相关输入:一个是实际的掩模张量mask_input(形状为[1, 1, 256, 256]),另一个是表示是否有掩模输入的标量has_mask_input(形状为[1])。在**首次推理(没有先前掩模)**时,mask_input应传入全零张量,has_mask_input应传入0。如果后续想基于前一个掩模进行优化(迭代式细化),则需要将前一次输出的low_res_masks经过sigmoid和阈值处理后,作为新的mask_input传入,并将has_mask_input设为1。导出和后续C++调用时必须严格区分这两个输入。- 验证导出结果:导出后,务必在Python端用ONNX Runtime验证一下导出的模型是否能正确运行,输出是否与原始PyTorch模型一致。这能提前发现严重的导出错误。
import onnxruntime as ort import numpy as np # 在Python中用ONNX Runtime验证编码器 ort_session = ort.InferenceSession("./models/sam_image_encoder.onnx") ort_inputs = {"image": dummy_image.numpy()} ort_outs = ort_session.run(None, ort_inputs) # 对比 ort_outs[0] 和 sam.image_encoder(dummy_image) 的输出是否接近 print(np.allclose(ort_outs[0], sam.image_encoder(dummy_image).detach().numpy(), rtol=1e-3, atol=1e-5))4. C++工程搭建:从零构建ONNX Runtime推理环境
模型准备好了,现在进入核心的C++部分。我们的目标是创建一个纯净的、可移植的C++项目,能够加载并运行这两个ONNX模型。
4.1 项目结构与依赖管理
我推荐使用CMake来管理项目,它能很好地处理跨平台编译和依赖查找。一个典型的项目结构如下:
sam_cpp_deploy/ ├── CMakeLists.txt ├── src/ │ ├── main.cpp │ ├── SamPipeline.cpp │ └── SamPipeline.h ├── include/ │ └── (第三方头文件,如果需要) ├── lib/ │ └── (预编译的ONNX Runtime库,可选) ├── models/ │ ├── sam_image_encoder.onnx │ └── sam_mask_decoder.onnx └── build/ └── (编译输出目录)最关键的依赖是ONNX Runtime的C++库。你有两种获取方式:
- 下载预编译包:从ONNX Runtime的GitHub Release页面下载对应平台(Windows/Linux/macOS)和硬件(CPU/GPU)的预编译包。这是最快的方式。
- 从源码编译:如果你想进行深度定制或链接特定版本的依赖,可以从源码编译。过程稍复杂,但可控性更强。
这里以使用预编译的Linux CPU版本为例。假设你将解压后的文件夹命名为onnxruntime-linux-x64-1.15.1,并将其放在项目根目录下。
4.2 CMakeLists.txt 配置详解
CMakeLists.txt是项目的构建蓝图,需要正确配置以找到ONNX Runtime的头文件和库。
cmake_minimum_required(VERSION 3.16) project(SAM_CPP_Deploy) set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) # 假设ONNX Runtime预编译包放在项目根目录下 set(ONNXRUNTIME_ROOT_DIR ${CMAKE_SOURCE_DIR}/onnxruntime-linux-x64-1.15.1) # 查找头文件 find_path(ONNXRUNTIME_INCLUDE_DIR NAMES onnxruntime_cxx_api.h PATHS ${ONNXRUNTIME_ROOT_DIR}/include NO_DEFAULT_PATH ) # 查找库文件(注意库名和路径) find_library(ONNXRUNTIME_LIB NAMES onnxruntime PATHS ${ONNXRUNTIME_ROOT_DIR}/lib NO_DEFAULT_PATH ) if (NOT ONNXRUNTIME_INCLUDE_DIR OR NOT ONNXRUNTIME_LIB) message(FATAL_ERROR "Failed to find ONNX Runtime. Please check ONNXRUNTIME_ROOT_DIR.") endif() include_directories(${ONNXRUNTIME_INCLUDE_DIR} ${CMAKE_SOURCE_DIR}/src) # 添加你的可执行文件 add_executable(sam_demo src/main.cpp src/SamPipeline.cpp) # 链接ONNX Runtime库 target_link_libraries(sam_demo ${ONNXRUNTIME_LIB}) # 在Linux/macOS上可能需要链接其他系统库,如pthread if(UNIX AND NOT APPLE) target_link_libraries(sam_demo pthread) endif()4.3 核心推理类 SamPipeline 的设计与实现
我们将所有与模型交互的逻辑封装到一个SamPipeline类中。这个类负责:
- 初始化ONNX Runtime环境(
Ort::Env)和会话选项(Ort::SessionOptions)。 - 加载编码器和解码器模型,创建会话(
Ort::Session)。 - 管理输入/输出张量的内存分配和数据类型转换。
- 提供
encodeImage和decodeMask两个核心方法。
以下是头文件SamPipeline.h的概要:
// SamPipeline.h #pragma once #include <onnxruntime_cxx_api.h> #include <opencv2/opencv.hpp> // 使用OpenCV进行图像处理 #include <vector> #include <memory> class SamPipeline { public: // 构造函数,传入两个ONNX模型路径 SamPipeline(const std::string& encoder_model_path, const std::string& decoder_model_path); ~SamPipeline(); // 编码图像,返回图像嵌入向量 std::vector<float> encodeImage(const cv::Mat& image); // 根据提示生成掩模 // image_embedding: encodeImage的输出 // points: 点坐标向量,格式为[x1, y1, x2, y2, ...] // point_labels: 对应的标签向量,1为前景,0为背景 // original_size: 原始图像尺寸 (height, width) // box: 可选框坐标 [x1, y1, x2, y2],如果为空则不使用 cv::Mat decodeMask(const std::vector<float>& image_embedding, const std::vector<float>& points, const std::vector<int64_t>& point_labels, const std::pair<int, int>& original_size, const std::vector<float>& box = {}); private: // ONNX Runtime环境,全局一个即可 Ort::Env env_; // 编码器和解码器会话 std::unique_ptr<Ort::Session> encoder_session_; std::unique_ptr<Ort::Session> decoder_session_; // 会话选项,可用于配置线程数、优化级别等 Ort::SessionOptions session_options_; // 输入/输出节点名称缓存 std::vector<const char*> encoder_input_names_; std::vector<const char*> encoder_output_names_; std::vector<const char*> decoder_input_names_; std::vector<const char*> decoder_output_names_; // 辅助函数:将cv::Mat转换为模型需要的张量格式 Ort::Value preprocessImage(const cv::Mat& image); // 辅助函数:创建张量 template <typename T> Ort::Value createTensor(const std::vector<int64_t>& shape, const std::vector<T>& data); };在SamPipeline.cpp的实现中,有几个细节至关重要:
1. 图像预处理必须与训练对齐SAM训练时,图像经过ResizeLongestSide变换(长边缩放到1024,保持比例),然后归一化到[0,1],再使用ImageNet的均值和标准差进行标准化。C++端的预处理必须完全复现这个过程。
cv::Mat SamPipeline::preprocess(const cv::Mat& rgb_image) { cv::Mat processed; // 1. ResizeLongestSide int long_side = 1024; int h = rgb_image.rows, w = rgb_image.cols; float scale = (float)long_side / std::max(h, w); int new_h = (int)(h * scale + 0.5f); int new_w = (int)(w * scale + 0.5f); cv::resize(rgb_image, processed, cv::Size(new_w, new_h), 0, 0, cv::INTER_LINEAR); // 2. 填充到1024x1024 int top = (1024 - new_h) / 2; int bottom = 1024 - new_h - top; int left = (1024 - new_w) / 2; int right = 1024 - new_w - left; cv::copyMakeBorder(processed, processed, top, bottom, left, right, cv::BORDER_CONSTANT, cv::Scalar(0)); // 3. 转换为float并归一化到[0,1] processed.convertTo(processed, CV_32FC3, 1.0 / 255.0); // 4. ImageNet标准化 (mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) std::vector<cv::Mat> channels; cv::split(processed, channels); for (int i = 0; i < 3; ++i) { channels[i] = (channels[i] - mean_[i]) / std_[i]; } cv::merge(channels, processed); // 5. 调整维度顺序为 NCHW (1, 3, 1024, 1024) cv::dnn::blobFromImage(processed, processed, 1.0, cv::Size(), cv::Scalar(), false, false); return processed; // 此时processed是1x3x1024x1024的cv::Mat }2. 输入张量的创建与内存管理ONNX Runtime的Ort::Value需要精确的形状和数据。必须注意数据在内存中的连续性以及数据类型的匹配。
template <typename T> Ort::Value SamPipeline::createTensor(const std::vector<int64_t>& shape, const std::vector<T>& data) { auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault); // 注意:data.data()返回的指针在vector生命周期内必须有效 return Ort::Value::CreateTensor<T>(memory_info, const_cast<T*>(data.data()), data.size(), shape.data(), shape.size()); } std::vector<float> SamPipeline::encodeImage(const cv::Mat& image) { cv::Mat input_blob = preprocess(image); // 将cv::Mat数据复制到vector,确保内存连续 std::vector<float> input_tensor_values(input_blob.total()); std::memcpy(input_tensor_values.data(), input_blob.data, input_blob.total() * sizeof(float)); std::vector<int64_t> input_shape = {1, 3, 1024, 1024}; auto input_tensor = createTensor(input_shape, input_tensor_values); // 运行推理 auto output_tensors = encoder_session_->Run(Ort::RunOptions{nullptr}, encoder_input_names_.data(), &input_tensor, 1, encoder_output_names_.data(), 1); // 获取输出张量(图像嵌入) float* output_data = output_tensors[0].GetTensorMutableData<float>(); auto output_shape = output_tensors[0].GetTensorTypeAndShapeInfo().GetShape(); size_t num_elements = 1; for (auto dim : output_shape) num_elements *= dim; return std::vector<float>(output_data, output_data + num_elements); }3. 解码器输入的复杂组装解码器的输入有多个,且部分输入(如点坐标)需要根据预处理时的变换进行归一化。这是C++实现中最容易出错的部分。
cv::Mat SamPipeline::decodeMask(const std::vector<float>& image_embedding, const std::vector<float>& points, const std::vector<int64_t>& point_labels, const std::pair<int, int>& original_size, const std::vector<float>& box) { // 1. 准备图像嵌入张量 std::vector<int64_t> embedding_shape = {1, 256, 64, 64}; // vit_l模型的嵌入形状 auto embedding_tensor = createTensor(embedding_shape, image_embedding); // 2. 处理点坐标:需要从原始图像坐标转换到1024x1024预处理空间的坐标 std::vector<float> transformed_points; for (size_t i = 0; i < points.size(); i += 2) { float x = points[i], y = points[i+1]; // 应用与预处理时相同的变换(缩放+填充偏移) float scale = 1024.0f / std::max(original_size.first, original_size.second); int pad_h = (1024 - int(original_size.first * scale + 0.5f)) / 2; int pad_w = (1024 - int(original_size.second * scale + 0.5f)) / 2; float tx = x * scale + pad_w; float ty = y * scale + pad_h; // 进一步归一化到[0, 1](SAM解码器要求) transformed_points.push_back(tx / 1024.0f); transformed_points.push_back(ty / 1024.0f); } std::vector<int64_t> points_shape = {1, static_cast<int64_t>(transformed_points.size()/2), 2}; auto points_tensor = createTensor(points_shape, transformed_points); // 3. 点标签张量 std::vector<int64_t> labels_shape = {1, static_cast<int64_t>(point_labels.size())}; auto labels_tensor = createTensor(labels_shape, point_labels); // 4. 掩模输入(全零)和 has_mask_input (0) std::vector<float> zero_mask(1 * 1 * 256 * 256, 0.0f); std::vector<int64_t> mask_shape = {1, 1, 256, 256}; auto mask_tensor = createTensor(mask_shape, zero_mask); std::vector<float> has_mask_vec = {0.0f}; std::vector<int64_t> has_mask_shape = {1}; auto has_mask_tensor = createTensor(has_mask_shape, has_mask_vec); // 5. 原始图像尺寸张量 std::vector<float> orig_size_vec = {static_cast<float>(original_size.first), static_cast<float>(original_size.second)}; std::vector<int64_t> orig_size_shape = {2}; auto orig_size_tensor = createTensor(orig_size_shape, orig_size_vec); // 6. 组装所有输入 std::vector<Ort::Value> decoder_inputs; decoder_inputs.push_back(std::move(embedding_tensor)); decoder_inputs.push_back(std::move(points_tensor)); decoder_inputs.push_back(std::move(labels_tensor)); decoder_inputs.push_back(std::move(mask_tensor)); decoder_inputs.push_back(std::move(has_mask_tensor)); decoder_inputs.push_back(std::move(orig_size_tensor)); // 如果有框输入,在此处添加 // 7. 运行解码器推理 auto output_tensors = decoder_session_->Run(Ort::RunOptions{nullptr}, decoder_input_names_.data(), decoder_inputs.data(), decoder_inputs.size(), decoder_output_names_.data(), 3); // 输出:masks, iou_predictions, low_res_masks // 8. 处理输出:第一个输出是最高分辨率的掩模 float* mask_data = output_tensors[0].GetTensorMutableData<float>(); auto mask_output_shape = output_tensors[0].GetTensorTypeAndShapeInfo().GetShape(); // mask_output_shape 通常是 [1, 1, H, W],其中H,W是原始图像尺寸 int out_h = mask_output_shape[2]; int out_w = mask_output_shape[3]; cv::Mat mask(out_h, out_w, CV_32FC1, mask_data); cv::Mat mask_8u; mask.convertTo(mask_8u, CV_8UC1, 255.0); // 转换为0-255的灰度图 return mask_8u; }5. 性能调优与实战踩坑记录
将流程跑通只是第一步,要让它在实际应用中可用,性能优化和稳定性处理必不可少。
5.1 会话配置与性能优化
在创建Ort::Session时,可以通过SessionOptions进行大量优化:
// 在SamPipeline构造函数中 session_options_.SetIntraOpNumThreads(4); // 设置并行计算线程数,通常设为CPU核心数 session_options_.SetInterOpNumThreads(2); // 对于多流执行,如果模型有多个子图 session_options_.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL); // 如果使用GPU(CUDA) Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_CUDA(session_options_, 0)); // 启用内存模式优化,减少内存碎片 session_options_.EnableCpuMemArena(); // session_options_.EnableMemPattern(); // 对于固定形状输入,可以启用此选项提升性能 encoder_session_ = std::make_unique<Ort::Session>(env_, encoder_model_path.c_str(), session_options_);5.2 图像编码缓存策略
这是提升交互式应用体验的关键。在真实场景中,用户会对同一张图像点击多个点或框。我们不应该每次都将图像重新编码。
class SamPipeline { private: std::vector<float> cached_embedding_; cv::Mat cached_original_image_; bool is_cache_valid_ = false; // ... 其他成员 public: cv::Mat predict(const cv::Mat& image, const std::vector<cv::Point>& points, const std::vector<int>& labels) { std::vector<float> embedding; if (!is_cache_valid_ || image.data != cached_original_image_.data) { // 简单比较,实际应用需更健壮的缓存键 embedding = encodeImage(image); cached_embedding_ = embedding; cached_original_image_ = image.clone(); is_cache_valid_ = true; } else { embedding = cached_embedding_; } // ... 转换points格式,调用decodeMask return decodeMask(embedding, transformed_points, point_labels, {image.rows, image.cols}); } };5.3 多提示输入与后处理
SAM支持同时输入多个点(前景和背景)以及框。在C++端,我们需要将这些输入正确地组装成模型期望的格式。一个常见的需求是结合一个框和框内的一个正点来获得更精确的分割。
// 假设同时有一个框和一个点 std::vector<float> box = {x1, y1, x2, y2}; // 原始图像坐标 std::vector<float> points = {point_x, point_y}; std::vector<int64_t> labels = {1}; // 前景点 // 在decodeMask函数内部,需要将框坐标也进行同样的变换和归一化 if (!box.empty()) { // 变换框坐标到1024x1024空间并归一化 float scale = 1024.0f / std::max(orig_h, orig_w); int pad_h = (1024 - int(orig_h * scale + 0.5f)) / 2; int pad_w = (1024 - int(orig_w * scale + 0.5f)) / 2; std::vector<float> transformed_box = { box[0] * scale + pad_w, box[1] * scale + pad_h, box[2] * scale + pad_w, box[3] * scale + pad_h }; for (auto& val : transformed_box) val /= 1024.0f; // 创建框输入张量,形状为 [1, 4] auto box_tensor = createTensor(std::vector<int64_t>{1, 4}, transformed_box); // 需要将box_tensor添加到decoder_inputs中,并确保decoder_input_names_包含"box" }后处理方面,模型直接输出的掩模是float类型,值在[0, 1]之间,表示每个像素是前景的概率。通常我们需要一个二值化的掩模。
cv::Mat probability_mask = decodeMask(...); // CV_32FC1 cv::Mat binary_mask; // 方法1:固定阈值 cv::threshold(probability_mask, binary_mask, 0.5, 1.0, cv::THRESH_BINARY); // 方法2:使用模型同时输出的iou_predictions作为置信度参考,动态调整阈值 binary_mask.convertTo(binary_mask, CV_8UC1, 255); // 转换为0-255的uchar图像5.4 内存与异常处理
C++环境下,内存管理和错误处理必须格外小心。ONNX Runtime的API会抛出Ort::Exception类型的异常。
try { auto output_tensors = session_->Run(...); } catch (const Ort::Exception& e) { std::cerr << "ONNX Runtime inference failed: " << e.what() << std::endl; // 检查常见问题:输入形状不匹配、数据类型错误、节点找不到等 // 通常e.what()会给出比较详细的错误信息 }对于图像编码器等大张量,要注意std::vector的生命周期。在createTensor函数中,我们使用了const_cast,这意味着Ort::Value内部直接引用了vector的数据指针,而不是复制。因此,必须确保在Ort::Value被使用(即Run方法调用)期间,原始的vector对象不能被销毁或重新分配内存。一种更安全的方法是直接为Ort::Value分配内存并拷贝数据,但这会增加一次内存拷贝的开销。在实际应用中,如果张量是在推理函数栈上创建的局部变量,由于其生命周期覆盖了整个推理过程,直接引用是安全的。
6. 集成与测试:构建一个简单的交互式演示
最后,我们可以创建一个简单的main.cpp来测试整个流程,并模拟一个交互式应用。这里使用OpenCV的GUI功能来显示图像并捕获鼠标点击作为提示点。
// main.cpp #include "SamPipeline.h" #include <opencv2/highgui.hpp> #include <iostream> std::vector<cv::Point> points; std::vector<int> labels; // 1 for foreground, 0 for background void onMouse(int event, int x, int y, int flags, void* userdata) { if (event == cv::EVENT_LBUTTONDOWN) { points.push_back(cv::Point(x, y)); labels.push_back(1); // 左键前景点 std::cout << "Added foreground point at (" << x << ", " << y << ")" << std::endl; } else if (event == cv::EVENT_RBUTTONDOWN) { points.push_back(cv::Point(x, y)); labels.push_back(0); // 右键背景点 std::cout << "Added background point at (" << x << ", " << y << ")" << std::endl; } } int main() { // 初始化Pipeline SamPipeline sam("./models/sam_image_encoder.onnx", "./models/sam_mask_decoder.onnx"); // 加载测试图像 cv::Mat image = cv::imread("./test_image.jpg"); if (image.empty()) { std::cerr << "Could not open image!" << std::endl; return -1; } cv::Mat image_rgb; cv::cvtColor(image, image_rgb, cv::COLOR_BGR2RGB); // SAM期望RGB cv::namedWindow("SAM C++ Demo"); cv::setMouseCallback("SAM C++ Demo", onMouse, nullptr); cv::Mat display_image = image.clone(); while (true) { cv::imshow("SAM C++ Demo", display_image); int key = cv::waitKey(1); if (key == 'q') { break; } else if (key == 'c') { // 清空点 points.clear(); labels.clear(); display_image = image.clone(); std::cout << "Cleared points." << std::endl; } else if (key == ' ') { // 空格键触发预测 if (!points.empty()) { std::cout << "Predicting with " << points.size() << " points..." << std::endl; // 转换点格式 std::vector<float> point_coords; for (auto& p : points) { point_coords.push_back(static_cast<float>(p.x)); point_coords.push_back(static_cast<float>(p.y)); } std::vector<int64_t> point_labels_int64(labels.begin(), labels.end()); cv::Mat mask = sam.decodeMask(sam.encodeImage(image_rgb), point_coords, point_labels_int64, {image.rows, image.cols}); // 将掩模叠加到原图显示 cv::Mat colored_mask; cv::applyColorMap(mask, colored_mask, cv::COLORMAP_JET); cv::addWeighted(image, 0.7, colored_mask, 0.3, 0, display_image); // 在图像上绘制点 for (size_t i = 0; i < points.size(); ++i) { cv::Scalar color = (labels[i] == 1) ? cv::Scalar(0, 255, 0) : cv::Scalar(0, 0, 255); // 绿前景,红背景 cv::circle(display_image, points[i], 5, color, -1); } std::cout << "Prediction done." << std::endl; } } } return 0; }编译并运行这个程序,你就能看到一个简单的交互式窗口。在图像上点击左键添加前景点,右键添加背景点,按空格键进行分割,按‘c’键清除所有点,按‘q’键退出。这验证了从图像加载、预处理、编码、提示输入到解码、后处理显示的完整链路。
整个项目从模型导出到C++集成,涉及了深度学习部署中的多个关键环节:模型格式转换、推理引擎选型、精确的预处理/后处理对齐、C++内存管理、性能优化以及简单的应用集成。虽然SAM模型本身较大(vit_l的编码器ONNX文件约350MB),但在现代PC上,CPU推理单次编码也能在数秒内完成,解码更是毫秒级,对于很多离线或对实时性要求不极致的交互应用来说,已经具备了实用性。更重要的是,这套纯C++的部署方案,为将最前沿的视觉AI能力无缝集成到现有的、复杂的C++工业软件体系中,铺平了道路。
本文还有配套的精品资源,点击获取