当前位置：首页 > 卖家故事 > AI分割一切模型SAM(Segment Anything Model)的C++部署

AI分割一切模型SAM(Segment Anything Model)的C++部署

时间：2024-04-09 20:51:10 来源：网络cs 作者：胡椒栏目：卖家故事阅读：

标签：部署分割模型

阅读本书更多章节>>>>

2023年最火爆的分割模型莫过于SAM，截止今天2024年1月19日，github上的star已经达到了41.7k的惊人数量。下面我们来体会一下如何运行这个模型，以及如何用C++部署这个模型。

检查cuda环境

我的Cuda版本是12.0.1，如下，

Cudnn是在官网下载的版本，

cudnn-windows-x86_64-8.9.2.26_cuda12-archive.zip

如果没有cuda也没关系，测试嘛，运行cpu版本就行。有cuda后面要注意版本问题。

下载SAM

SAM的官网地址是：https://github.com/facebookresearch/segment-anything

(base) D:\aiSeg> git clone --recursive https://github.com/facebookresearch/segment-anything

(base) D:\aiSeg>cd segment-anything;

(base) D:\aiSeg\segment-anything>

配置anaconda环境

(base) D:\aiSeg\segment-anything> conda ctreate -n seg

(base) D:\aiSeg\segment-anything> conda activate seg

(seg) D:\aiSeg\segment-anything>

然后，按照pytorch上的说明安装，详情可以参考：

PyTorch

(seg) D:\aiSeg\segment-anything> conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

这里选pytorch-cuda=12.1，反正后面也没碰到啥问题，就酱。

再然后安装cudatools，默认安装了11.8，

(seg) D:\aiSeg\segment-anything> conda install cudatoolkit

The following NEW packages will be INSTALLED:

cudatoolkit pkgs/main/win-64::cudatoolkit-11.8.0-hd77b12b_0

再安装segment-anything,

(seg) D:\aiSeg\segment-anything>pip install -e .

Obtaining file:///D:/aiSeg/segment-anything

Preparing metadata (setup.py) ... done

Installing collected packages: segment-anything

Running setup.py develop for segment-anything

Successfully installed segment-anything-1.0

还有OpenCV等等等等要安装。。。。

(seg) D:\aiSeg\segment-anything>conda install opencv

The following NEW packages will be INSTALLED:

opencv pkgs/main/win-64::opencv-4.6.0-py311h5d08a89_5

(seg) D:\aiSeg\segment-anything>

测试SAM

安装完毕后，用vscode直接打开segment-anything文件夹，选择要测试运行的amg.py文件，

然后配置launch.json中的参数，内容如下，

{    // Use IntelliSense to learn about possible attributes.    // Hover to view descriptions of existing attributes.    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387    "version": "0.2.0",    "configurations": [        {            "name": "Python: Current File",            "type": "python",            "request": "launch",            "program": "${file}",            "console": "integratedTerminal",            "justMyCode": false,            "args": [                "--input=D:/aiImages/SoderbergLenna/Lenna.png",                "--output=D:/aiImages/SoderbergLenna/",                "--model-type=“vit_b",                "--checkpoint=D:/aiDataset/sam/sam_vit_b_01ec64.pth"            ]        }    ]}

这时候，就可以点击Starting Debug(F5)直接开跑了，一路顺利的话会得到下面的结果，

原图，

生成的mask，

导出ONNX文件

然后我们导出onnx文件，这一步是为后面的c++部署做准备的。

python scripts/export_onnx_model.py --checkpoint D:/aiDataset/sam/sam_vit_h_4b8939.pth --model-type vit_h --output sam_vit_h_4b8939.onnx

这样，我们就生成了sam_vit_h_4b8939.onnx文件。

下载cpp部署

项目地址在：https://github.com/dinglufe/segment-anything-cpp-wrapper

老办法，

> git clone --recursive https://github.com/dinglufe/segment-anything-cpp-wrapper

然后发现作者用的onnxruntime，因此我们也安装一下，

> conda install onnx

> conda install onnxruntime

另外，还要到 https://github.com/microsoft/onnxruntime

去下载windows上的动态库，我这里的本机地址是，

D:\aiEngine\onnxruntime-win-x64-gpu-1.16.2

生成所需要的辅助文件

项目中有一个export_pre_model.py文件，我略改了一下，因为我没有用hq-sam，直接用的segment-anything，改后内容如下，

# Code formatter may move "import segment_anything as SAM" and "import mobile_sam as SAM" to the top# But this may bring errors after switching modelsimport torchimport numpy as npimport osfrom segment_anything.utils.transforms import ResizeLongestSidefrom onnxruntime.quantization import QuantTypefrom onnxruntime.quantization.quantize import quantize_dynamicoutput_names = ['output']# Generate preprocessing model of Segment-anything in onnx format# Use original segment-anything, Mobile-SAM or HQ-SAM model# Each in a separate Python virtual environment# Uncomment the following lines to generate preprocessing model of Segment-anythingimport segment_anything as SAM# # Download Segment-anything model "sam_vit_h_4b8939.pth" from https://github.com/facebookresearch/segment-anything#model-checkpoints# # and change the path belowcheckpoint = 'sam_vit_h_4b8939.pth'model_type = 'vit_h'output_path = 'models/sam_preprocess.onnx'quantize = True# Target image size is 1024x720image_size = (1024, 720)output_raw_path = output_pathif quantize:    # The raw directory can be deleted after the quantization is done    output_name = os.path.basename(output_path).split('.')[0]    output_raw_path = '{}/{}_raw/{}.onnx'.format(        os.path.dirname(output_path), output_name, output_name)os.makedirs(os.path.dirname(output_raw_path), exist_ok=True)sam = SAM.sam_model_registry[model_type](checkpoint=checkpoint)sam.to(device='cpu')transform = ResizeLongestSide(sam.image_encoder.img_size)image = np.zeros((image_size[1], image_size[0], 3), dtype=np.uint8)input_image = transform.apply_image(image)input_image_torch = torch.as_tensor(input_image, device='cpu')input_image_torch = input_image_torch.permute(    2, 0, 1).contiguous()[None, :, :, :]class Model(torch.nn.Module):    def __init__(self, image_size, checkpoint, model_type):        super().__init__()        self.sam = SAM.sam_model_registry[model_type](checkpoint=checkpoint)        self.sam.to(device='cpu')        self.predictor = SAM.SamPredictor(self.sam)        self.image_size = image_size    def forward(self, x):        self.predictor.set_torch_image(x, (self.image_size))        if 'interm_embeddings' not in output_names:            return self.predictor.get_image_embedding()        else:            return self.predictor.get_image_embedding(), torch.stack(self.predictor.interm_features, dim=0)model = Model(image_size, checkpoint, model_type)model_trace = torch.jit.trace(model, input_image_torch)torch.onnx.export(model_trace, input_image_torch, output_raw_path,                  input_names=['input'], output_names=output_names)if quantize:    quantize_dynamic(        model_input=output_raw_path,        model_output=output_path,        per_channel=False,        reduce_range=False,        weight_type=QuantType.QUInt8,    )

生成的这一堆文件如下，