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

news2024/10/7 12:18:55

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 models
import torch
import numpy as np
import os

from segment_anything.utils.transforms import ResizeLongestSide

from onnxruntime.quantization import QuantType
from onnxruntime.quantization.quantize import quantize_dynamic

output_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-anything
import 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 below
checkpoint = 'sam_vit_h_4b8939.pth'
model_type = 'vit_h'
output_path = 'models/sam_preprocess.onnx'
quantize = True


# Target image size is 1024x720
image_size = (1024, 720)

output_raw_path = output_path
if 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,
    )

生成的这一堆文件如下，