# SAM ## 论文 Segment Anything - https://arxiv.org/abs/2304.02643 ## 模型结构 ![](./assets/model_diagram.png) 如图,该模型的网络结构主要分三个部分:Image encoder、Prompt encoder和Lightweight mask decoder。 ### Image encoder 通常,Image encoder可以是输出C×H×W图像嵌入的任何网络。受可扩展性和强大的预训练的启发,使用MAE预训练的视觉转换器(ViT),具有最小的适应性来处理高分辨率输入,特别是具有14×14窗口注意力和四个等距全局注意力块的ViT-H/16。图像编码器的输出是输入图像的16倍缩小的嵌入。使用1024×1024的输入分辨率,该分辨率是通过重新缩放图像并填充短边获得的。因此,图像嵌入是64×64。为了降低通道维度,使用1×1卷积来获得256个通道,然后使用3×3卷积来获得同样具有256个通道的通道。 ### Prompt encoder 按如下方式映射到256维向量嵌入中。一个点由其位置的位置编码的总和以及两个学习的嵌入之一表示,指示该点是在前景还是背景中。一个框由一个嵌入对表示:(1)其左上角的位置编码与表示“左上角”的学习嵌入相加。(2)相同的结构,但使用表示“右下角”的学习嵌入。最后,为了表示自由形式的文本,使用CLIP中的文本编码器(通常任何文本编码器都可以)。密集提示(即掩码)与图像具有空间对应关系。我们以比输入图像低4倍的分辨率输入掩码,然后使用两个2×2、步幅2(convolutions)进行降采样,输出通道分别为4和16。最后,1×1卷积将通道维度映射到256个维度。每层之间通过GELU激活和层归一化分开。然后将掩码嵌入和图像嵌入逐元素相加。如果没有掩码提示,则添加表示“无掩码”的学习嵌入到每个图像嵌入位置。 ### Lightweight mask decoder ![](./assets/mask_decoder.PNG) 解码器设计如图所示。使用两个转置卷积层将更新的图像嵌入放大4倍(现在相对于输入图像缩小了4倍)。将更新的输出token嵌入传递给一个小的3层MLP,它输出一个与扩展的图像嵌入的通道维度匹配的向量。最后,使用扩展的图像嵌入和MLP的输出进行空间点乘,预测出一个掩码。Transformer使用256的嵌入维度。在64×64图像嵌入的交叉注视层中,查询、键和值通道维度为128。所有注意力层都使用8个头。用于放大输出图像嵌入的转置卷积是2×2,步幅2,输出通道维度为64和32,并具有GELU激活。它们通过层归一化来分隔。 ## 算法原理 SAM分为图像编码器和快速提示编码器/掩码解码器,可以重用相同的image embedding图像嵌入(并摊销其成本)与不同的提示。给定image embedding图像嵌入,提示编码器和掩码解码器可以在web浏览器中预测掩码。为了使SAM实现模糊感知,设计它来预测单个提示的多个掩码,从而使SAM能够自然地处理模糊性。 ## 环境配置 ### Docker(方法一) 从[光源](https://www.sourcefind.cn/#/service-list)拉取镜像 ``` docker pull image.sourcefind.cn:5000/dcu/admin/base/pytorch:1.10.0-centos7.6-dtk-22.10-py39-latest docker run -it --network=host --name=SAM_pytorch --privileged --device=/dev/kfd --device=/dev/dri --ipc=host --shm-size=16G --group-add video --cap-add=SYS_PTRACE <上面获得的镜像ID> /bin/bash ``` 安装其他依赖: ``` pip install opencv-python pycocotools matplotlib onnxruntime onnx ``` ### Dockerfile(方法二) ``` docker build --no-cache -t image.sourcefind.cn:5000/dcu/admin/base/pytorch:1.10.0-centos7.6-dtk-22.10-py39-latest . docker run -it --network=host --name=SAM_pytorch --privileged --device=/dev/kfd --device=/dev/dri --ipc=host --shm-size=16G --group-add video --cap-add=SYS_PTRACE <上面获得的镜像ID> /bin/bash ``` 安装其他依赖: ``` pip install opencv-python pycocotools matplotlib onnxruntime onnx ``` ### Anaconda(方法三) 直接使用pip install的方式安装 ``` 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 . ``` ## 数据集 数据集名称:SA-1B Dataset 完整数据集可在[这里](https://ai.facebook.com/datasets/segment-anything-downloads/)进行下载 项目中用于试验训练的迷你数据集结构如下 ``` ── notebooks │   ├── images │ │   ├── dog.jpg │ │   ├── groceries.jpg │ │ └── trunk.jpg ``` ## 微调 官网提供了生成掩码的预训练权重和生成掩码的脚本,没有提供训练脚本,但可使用第三方提供的示例脚本微调 如果您有兴趣,参考[这里](https://github.com/luca-medeiros/lightning-sam/blob/main/lightning_sam/train.py). ### 单机单卡 ``` git clone https://github.com/luca-medeiros/lightning-sam.git cd lightning-sam pip install . cd lightning-sam python train.py ``` ### 推理 ``` python scripts/amg.py --checkpoint --model-type --input --output ``` 注:checkpoint预训练模型在[这里](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth)下载 指令中: 代表选择权重的路径 代表不同的模型,可选择'vit_h'、'vit_l'、 'vit_b' 代表输入图片或者文件夹的路径 代表分割结果保存路径 ### result ![](./ouputs/000000524456/0.png) 掩码生成的部分结果在同级目录outputs中可以查看,结果示例如上图,官方提供demo可在[这里](https://segment-anything.com/demo)试用 ## 应用场景 ### 算法类别 图像分割 ### 热点应用行业 SAM算法能智能抠图,可以用于识别图像和视频中的物体,具有广泛的通用性。SAM作为开源且更通用AI系统的强大组件,赋能工业、煤矿、医学影像、安防监控等多场景。 热点应用行业:能源,医疗,网安 ## 源码仓库及问题反馈 https://developer.hpccube.com/codes/modelzoo/sam_pytorch ## 参考资料 https://github.com/facebookresearch/segment-anything https://github.com/luca-medeiros/lightning-sam/blob/main/lightning_sam(第三方SAM微调)