基于丹摩DAMODEL部署Segment Anything 模型，智能分割一切图片

最近在研究图像处理的过程中，接触到了Mate开源的图像分割项目，看了他们官网的demo演示，感觉很神奇，也有有趣，上传一张图片，点击一下就能将点击的物体切割出来，很类似自动物体识别，智能抠图。于是决定深入研究需学习一番。要研究学习一个模型，必不可少地要用到GPU计算，这个时候我突然想起来之前注册的丹摩（DAMODEL）送了不少代金劵，于是租借一台GPU服务器来部署学习这个模型。文末有关于丹摩平台的

拿我格子衫来

10342人浏览 · 2024-09-25 09:20:16

拿我格子衫来 · 2024-09-25 09:20:16 发布

前言

最近在研究图像处理的过程中，接触到了Mate开源的图像分割项目segment-anything，看了他们官网的demo演示，感觉很神奇，也有有趣，上传一张图片，点击一下就能将点击的物体切割出来，很类似自动物体识别，智能抠图。于是决定深入研究需学习一番。要研究学习一个模型，必不可少地要用到GPU计算，这个时候我突然想起来之前注册的丹摩（DAMODEL）送了不少代金劵，于是租借一台GPU服务器来部署学习这个模型。文末有关于丹摩平台的介绍。

租借GPU服务器

点击此处，丹摩登录入口，登录丹摩
在这里插入图片描述
点击左侧菜单的GPU云实例，然后点击创建实例，
规格我们选择如下：

镜像选择
PyTorch 2.3.0
Python 3.10.14
系统 Ubuntu-22.04
CUDA版本：12.1

在这里插入图片描述
内置了CUDA，可自由选择PyTorch，PythonCUDA版本，这是我选择丹摩的重要原因。这些基础设置要是每次人为去安装配置，将会浪费大量时间。

注意：segment-anything 模型的要求 python>=3.8 , pytorch>=1.7, torchvision>=0.8

选择配置后，点击右下角的立即创建。
稍等片刻实例就会创建成功。
在这里插入图片描述

点击操作栏的JupyterLab 就能够打开实例中已经安的Jupyter。

你也可以使用访问链接和密码来访问服务器的终端。

访问Jupyter进行编码
在这里插入图片描述

终端登录
在这里插入图片描述

Segment Anything 模型介绍

在这里插入图片描述
Segment Anything Model（SAM）可根据输入提示（如点或方框）生成高质量的对象遮罩，并可用于为图像中的所有对象生成遮罩。该模型已在一个包含 1100 万张图像和 11 亿个遮罩的数据集上进行过训练，在各种分割任务中都具有很强的零误差性能。
在这里插入图片描述

配置环境

通过上面的步骤我们已经进入到了GPU实例的终端，接下来我们要在这个实例上部署Segment Anything 模型。

要在实例中安装Segment Anything 模型，有两种方式

第一种：在实例的终端中执行下面的命令

pip install git+https://github.com/facebookresearch/segment-anything.git

第二种：你也可以执行下面的命令来安装

git clone git@github.com:facebookresearch/segment-anything.git
cd segment-anything; pip install -e .

安装还需要安装相关的依赖库

pip install opencv-python pycocotools matplotlib onnxruntime onnx

环境安装过程
在这里插入图片描述

通过上面的步骤，我么现在已经安装好了模型运行的环境。

下载模型，运行模型

安装好模型的运行环境后，我们需要下载官方训练好的分隔模型，一共有三个分别是

大杯 ViT-H SAM model 模型类型为vit_h
中杯 ViT-L SAM model 模型类型为vit_l
小杯 ViT-B SAM model 模型类型为vit_b

这里以大杯ViT-H SAM model 模型为例，使用浏览器下载，大概有2.4G。然后我们将其上传到丹摩平台的文件存储。
在这里插入图片描述
如果直接在实例中下载，会比较漫长。通过上传到文件存储就会自动挂载到实例中，这里存储的根文件对应到实例中目录 /root/shared-storage

文件上传成功后，我们可以在实例的终端中看到这些文件。如下图：
在这里插入图片描述
文件和服务都准备好后，我们下面开始测试一下模型，处理分割一些文件。

测试模型

我们使用上传的模型，对一张图片进行分割，然后分割后的图片，放到一个目录里。
原始图片de.png
在这里插入图片描述

然后我们进入 segment-anything的目录，执行下面的指令。

python scripts/amg.py --checkpoint /root/shared-storage/sam_vit_h_4b8939.pth --model-type vit_h --input /root/shared-storage/de.png --output /home/imgoutput

使用sam_vit_h_4b8939.pth模型将 de.png 图片进行分割，分割后的图片存放在/home/imgoutput目录中。
执行过程如下：
在这里插入图片描述

进入/home/imgoutput目录，可以看到有一个de目录，de就是图片的名称。
在这里插入图片描述
一张图片分割了39个图片，我们看一下这些分割的图片是什么样的。分割后的图片如下图：
文件 matedata.csv的内容如下：

到这里我们已经成功地将模型部署到GPU的实例中了。
运行一次该模型，消耗的资源还是比较小的，在丹摩的实例监控图上，基本看不到资源的负载太大的起伏如下图：
在这里插入图片描述

运行Demo

上面我们已经使用模型将图片分割出来了，那么下面看一下，该如何使用官方提供的demo，查看一个图片的分割块。

这里是官方提供的一个demo，这个demo简单展示了如何使用Segment Anything 模型切割出来的数据，展示效果如下：

下面我教大家将该demo运行起来。
首先在jupyter中运行下面的代码，

import numpy as np
import cv2
from segment_anything import SamPredictor, sam_model_registry
checkpoint = "/root/shared-storage/sam_vit_h_4b8939.pth"
model_type = "vit_h"
sam = sam_model_registry[model_type](checkpoint=checkpoint)
sam.to(device='cuda')
predictor = SamPredictor(sam)
image = cv2.imread('/root/shared-storage/de.png')
predictor.set_image(image)
image_embedding = predictor.get_image_embedding().cpu().numpy()
np.save("de_embedding.npy", image_embedding)

运行后会生成一个de_embedding.npy文件将该文件下载到本地，并放在segment-anythingdemo/src/assets/data/目录中
在这里插入图片描述

然后再运行以下指令导出web使用的onnx模型

python scripts/export_onnx_model.py --checkpoint /root/shared-storage/sam_vit_h_4b8939.pth --model-type vit_h --output /root/shared-storage/sam_vit_h_onnx_example.onnx

正常运行完成后，将/root/shared-storage/sam_vit_h_onnx_example.onnx文件放置在项目目录segment-anything/demo/model/sam_vit_h_onnx_example.onnx

然后在 demo目录修改App.tsx的引用文件路径
在这里插入图片描述

最后在demo目录运行

yarn && yarn start

浏览器打开启动的服务，一般是 http://localhost:8081/

看一下动态效果图：
在这里插入图片描述
鼠标放置在某一色块上，该色块就会被增加一层阴影，标识被识别，分割出来啦。如果要做成官网那种交互效果还需要花点前端的功夫。
如果有人需要，我也可以努力尝试做出官方的效果。

丹摩平台介绍

DAMODEL（丹摩智算）是专为 AI 打造的智算云，致力于提供丰富的算力资源与基础设施助力 AI 应用的开发、训练、部署。

平台特性:

💡 超友好！

配备 124G 大内存和 100G 大空间系统盘，一键部署，三秒启动，让 AI 开发从未如此简单！
💡 资源多！

从入门级到专业级 GPU 全覆盖，无论初级开发还是高阶应用，你的需求，我们统统 Cover！
💡 性能强！

自建 IDC，全新 GPU，每一位开发者都能体验到顶级的计算性能和专属服务，大平台值得信赖！
💡 超实惠！

超低价格体验优质算力服务，注册即送优惠券！还有各类社区优惠活动，羊毛薅不停！

总结

通过使用丹摩DAMODEL部署Segment Anything 模型，我们了解了平台的使用方法，平台的各种资源使用，实例规格，也知道了在部署Segment Anything 模型需要注意哪些事项。再部署模型后，后续我们可以集成一个后端开发框架，让用户上传图片，然后进行图片切割，做成一个智能抠图的产品或服务。