使用 Segment Anything 模型进行图像分割教程

2,539 阅读4分钟

Segment Anything (SAM) 是一个可以根据点或框等输入提示生成高质量的图像分割的机器视觉模型,可用于为图像中的所有对象生成对应蒙版。该模型在包含 1100 万张图像和 11 亿个掩模的数据集上进行训练,在各种分割任务上具有强大的零样本性能。实现了真正意义上的分割万物。

使用平台:

OpenBayes贝式计算:openbayes.com/console/sig…

该教程提供三种方式,第一种是在线推理网站,方便大家对模型进行操作和更好的可视化,操作更简单方便;第二种是 GitHub 项目的源代码的实现,操作更复杂一些;最后一种是自动分割,只需上传图片并在 notebook 中修改相关代码即可实现自动分割。

教程地址:openbayes.com/console/pub…

方式一:在线推理网站

首先登录「OpenBayes」平台,打开「公共教程」,找到「Segment Anything 源代码实现与在线推理」教程。

点击「克隆」。

选择一块 RTX 4090 的算力和 PyTorch 的镜像。

第一次运行需要等待 3-5 分钟,等待它分配好资源后,打开工作空间。

在运行之前,先安装 Segment Anything 模型。新建一个终端界面后运行 pip install -e ./segment-anything 命令。
然后运行 pip install opencv-python pycocotools matplotlib onnxruntime onnx gradio 命令安装所需依赖。

最后运行 npm install --g yarn 安装模型。

在使用 npm 安装 yarn 时可能遇见网络问题,此时可以运行以下两种代码(选其中一个运行即可)更换镜像源。

npm config set registry registry.npmmirror.com

npm config set registry npm.aliyun.com

运行 python webui.py 后,打开右侧 API 地址,启动 Gradio 页面。

SAM Model 下可以选择模型的大小 ,h 版本的可以处理高清晰度的图像,需要更高的显存和算力,l 版本可以处理低清晰度图片,需要的显存和算力更低,b 版本的则介于两者之间。 RTX 4090 的算力可以使用 vit_h 版本的模型。

我们用一张熊猫的图片作为例子,上传后,点击「预览输入」,即可看见图片对应的坐标。

如果我们想要把熊猫给分割出来,可以在下图位置输入点坐标,比如(400,550),(400,250)这两个点分别选中了熊猫的头和身体,我们用这两个点去分割图像,点击「运行」即可。

运行完成后,可以看到熊猫已经被分割了出来。

我们还可以使用方框来将竹子分割出来。使用过方框输入需要方框左上角和右下角的点,这个竹子所在方框左上角的点位于(150,100),右下角的点位于(325,475),输入好之后我们点击「添加方框」,再点击预览输入,即可看见有一颗方框框住了这个柱子。

点击「运行」,即可看见它把这个竹子分割了出来。

方法二:可交互网站

如已经使用在线推理网站,需要在已运行的终端输入 ctrl + c 终止当前网站的运行。首先要对需要分割的图片进行预处理,在左侧文件管理器中打开 segment-anything/demo/src/assets/data 目录,将想要处理的图片命名好并添加进去,这里用小狗的图片举例。

然后我们将 README 里的 panda.jpeg 改为 dog.jpeg,图片的名称也改为 dog。使用 Shift+Enter 运行这段代码。

然后我们返回到 segment-anything/demo/src 目录,找到 App.tsx, IMAGE_PATH 为复制进目录的图片,IMAGE_EMBEDDING 为上一行生成的 <image_name>.npy,把名字改为输入的名字即可,MODEL_DIR 切换为使用的模型名称即可。

之后切换工作目录,新建一个终端,使用 cd segment-anything/demo 命令将目录切换到我们的 demo 目录下。

然后运行 yarn && yarn start 来启动我们的网站。如果提示网络问题无法访问,可以把它的源切换为国内的镜像源。

如果运行过程中遇到网络问题无法安装,运行其中一个命令切换 yarn 的源,选择其中一个能用的即可,如没有问题可忽略此步

淘宝源:yarn config set registry registry.npmmirror.com

阿里源:yarn config set registry npm.aliyun.com

运行起来后,我们通过 API 访问,就可以进入这个可交互网站了,它会根据鼠标的位置分割对应的图片。

方法三:自动分割

在 notebook 中提供了一个自动分割的方法,它可以自动处理图片。

把需要分割的图片放到 home/imgs 目录中,将 README 中的路径改为对应的路径,然后 Shift+Enter 运行就可以得到结果了。

可以看到它用不同的颜色标注出来了分割出来的结果。