Segment Anything 模型SAM推理链路实战跟着一次点击看掩码如何诞生【免费下载链接】segment-anythingThe repository provides code for running inference with the SegmentAnything Model (SAM), links for downloading the trained model checkpoints, and example notebooks that show how to use the model.项目地址: https://gitcode.com/GitHub_Trending/se/segment-anythingSegment AnythingSAM是一个可提示的图像分割模型给它一张图加上几个点击或一个框它就能立刻输出一张高质量的目标掩码。这篇文章不拆模块而是沿着真实推理链路走一遍——从图片进模型前的 1024 规范化到MaskDecoder吐出 256×256 的低分辨率掩码最后再裁剪回原图尺寸看每一步在代码里长什么样。仓库里的 notebooks/predictor_example.ipynb 演示的就是这类交互先框住一只狗再点进狗身内、点出狗外掩码随之收缩。但你可能会好奇为什么第二下点击几乎是瞬间返回的答案就在推理链路的第一刀上。先算贵的SamPredictor 把链路劈成两半交互式场景有个绕不开的成本问题ImageEncoderViT是一个大 ViT每点一次鼠标都重跑一遍完全浪费。segment_anything/predictor.py里的SamPredictor为此把模型劈成两段——set_image只跑一次昂贵的图像编码predict可以无限次调用每次只跑轻量的提示编码和掩码解码。下面这段是set_torch_image的核心问题它要解决的是图片进来后第一批重活在哪干# segment_anything/predictor.py self.original_size original_image_size self.input_size tuple(transformed_image.shape[-2:]) input_image self.model.preprocess(transformed_image) self.features self.model.image_encoder(input_image)对应架构的第一环图像特征self.features在这里被算出并缓存之后所有提示都只在这个固定特征上作业。记住这个64×64×256的特征图它是后文每一步的底片。进 1024 的方框归一化和补齐图片在进编码器前还要过一道预处理。Sam类segment_anything/modeling/sam.py的preprocess用 ImageNet 的均值方差归一化像素然后把图补成 1024×1024 的正方形# segment_anything/modeling/sam.py def preprocess(self, x: torch.Tensor) - torch.Tensor: x (x - self.pixel_mean) / self.pixel_std # ImageNet 均值/方差 h, w x.shape[-2:] padh self.image_encoder.img_size - h # 1024 - h padw self.image_encoder.img_size - w x F.pad(x, (0, padw, 0, padh)) return x为什么要垫边因为 ViT 的 patch 切分和位置嵌入都假设输入是正方形而真实图片长宽比各异。代价是多算了一些无意义的填充像素——不过这些像素在输出阶段会被原样裁掉前文存下的input_size和original_size就是为此准备的裁剪坐标。4096 个块图像如何变成 64×64 特征网格ImageEncoderViTsegment_anything/modeling/image_encoder.py接手后第一件事是 patch 化16×16 的卷积核一步把 1024×1024 的图切成 64×644096 个 token每个 768 维以默认的 ViT-B 计再加一份可学习的绝对位置嵌入。真正花心思的是注意力组织方式。全图 4096 个 token 做全局注意力的代价是 O(4096²)于是 SAM 的 12 层Block里绝大多数用窗口注意力window_size14每个 token 只和 14×14 窗口内的邻居算只留少数几层做全局注意力global_attn_indexes[2, 5, 8, 11]# segment_anything/modeling/image_encoder.py self.patch_embed PatchEmbed(kernel_size(16, 16), stride(16, 16), in_chans3, embed_dim768) for i in range(depth): block Block( dimembed_dim, num_headsnum_heads, # 指定层走全局注意力其余走窗口注意力 window_sizewindow_size if i not in global_attn_indexes else 0, input_size(img_size // patch_size, img_size // patch_size), )这是链路的第二环局部细节在窗口里算便宜全局关系靠那几层全局注意力兜底。ViT-L/ViT-H 只是把这个结构加深加宽24/32 层配方不变。Transformer 塔出来的特征是 768 通道但下游解码器只需要 256 维。编码器末尾的 neck 用1×1 卷积降维 3×3 卷积混邻域 LayerNorm两步完成self.neck nn.Sequential( nn.Conv2d(embed_dim, out_chans, kernel_size1, biasFalse), # 768 - 256 LayerNorm2d(out_chans), nn.Conv2d(out_chans, out_chans, kernel_size3, padding1, biasFalse), LayerNorm2d(out_chans), )到这里图片已经变成一张64×64×256的特征图——分辨率降到 1/16但语义密度足够支撑后面的掩码预测。把点击写成向量稀疏提示编码掩码预测需要回答提示到底指向哪儿这项工作由PromptEncodersegment_anything/modeling/prompt_encoder.py完成。坐标怎么变成向量SAM 没有用常见的正弦位置编码而是用一个固定的随机高斯矩阵PositionEmbeddingRandomcoords 2 * coords - 1 # [0,1] - [-1,1] coords coords self.positional_encoding_gaussian_matrix coords 2 * np.pi * coords return torch.cat([torch.sin(coords), torch.cos(coords)], dim-1)坐标先归一化到 [-1,1]乘以一个训练时就不动的高斯矩阵再套 sin/cos——相当于给每个坐标点一份随机频率指纹。好处是不用学习、不依赖分辨率点击落在哪个像素都是确定可算的。光有位置还不够模型还得知道点的身份。_embed_points给位置编码叠加了与标签绑定的可学习向量point_embedding self.pe_layer.forward_with_coords(points, self.input_image_size) point_embedding[labels -1] 0.0 point_embedding[labels -1] self.not_a_point_embed.weight point_embedding[labels 0] self.point_embeddings[0].weight # 负点 point_embedding[labels 1] self.point_embeddings[1].weight # 正点几个容易忽略的细节坐标会先加 0.5 对齐到像素中心一次预测没有任何点时会补一个label-1的非点占位 token保证解码器永远至少有一个稀疏提示框提示不单独建模而是拆成两个对角点各自叠加第 3、4 个可学习嵌入——框被降维成了带角色的点。PromptEncoder的输出分两路稀疏路点/框B×N×256进解码器的 token 序列稠密路上一次的掩码经卷积下采样成256×64×64直接加到图像特征上。没有掩码输入时稠密路用一个可学习的no_mask_embed铺满形状保持一致。256×256 的掩码和它的质量分MaskDecoder 内部链路的最后一环在MaskDecodersegment_anything/modeling/mask_decoder.py。它的输入很杂——图像特征、稠密提示、稀疏提示——统一的办法是全部 token 化1 个 IoU token、4 个掩码 token拼上用户的点/框 token一起塞进一个只有 2 层的TwoWayTransformer# segment_anything/modeling/mask_decoder.py output_tokens torch.cat([self.iou_token.weight, self.mask_tokens.weight], dim0) tokens torch.cat((output_tokens, sparse_prompt_embeddings), dim1) src torch.repeat_interleave(image_embeddings, tokens.shape[0], dim0) src src dense_prompt_embeddings pos_src torch.repeat_interleave(image_pe, tokens.shape[0], dim0) hs, src self.transformer(src, pos_src, tokens)注意repeat_interleave图像特征被复制 4 份每份配一个掩码 token 独立做双向注意力token 看图像、图像也看 token。这就是多候选掩码的算力来源——一次前向四个掩码并行。得到 token 输出后掩码不是直接卷积出来的而是走超网络hypernetwork路线每个掩码 token 过一个 MLP生成一组逐像素的线性权重与上采样到 256×256 的图像特征做点积upscaled_embedding self.output_upscaling(src) # 64x64 - 256x256 for i in range(self.num_mask_tokens): hyper_in_list.append(self.output_hypernetworks_mlpsi) hyper_in torch.stack(hyper_in_list, dim1) masks (hyper_in upscaled_embedding.view(b, c, h * w)).view(b, -1, h, w) iou_pred self.iou_prediction_head(iou_token_out)同一份上采样特征四套不同的像素权重得到四张不同的掩码——结构上就保证了多候选必然不同。IoU token 则单独走一个 MLP 头给四个候选各打一个质量分。最后按multimask_output切片False时只取第 0 号最优掩码True时取后 3 个去歧义候选让调用者或人自己挑。还有一处闭环值得留意predict的第三个返回值就是这里算出的 256×256 低分辨率 logitslow_res_masks它可以作为下一轮predict的mask_input喂回去。交互精修先框后点、越点越准的循环靠的就是这条回路。最后一英里裁掉填充回到原图尺寸解码器吐出的掩码只有 256×256离原图上的像素级掩码还差两步。Sam.postprocess_masks先把 logits 双线性插值到 1024×1024按之前存下的input_size裁掉右下的填充区再插值回用户原始分辨率最后在forward里以 0 为阈值二值化masks F.interpolate(masks, (img_size, img_size), modebilinear, align_cornersFalse) masks masks[..., : input_size[0], : input_size[1]] # 裁掉填充 masks F.interpolate(masks, original_size, modebilinear, align_cornersFalse) return masks self.mask_threshold # threshold 0到此一张完整的推理链路走完规范化 → 编码 → 提示编码 → 解码 → 回贴原图。整条链里最重的是第二步最轻的是后三步这正是SamPredictor敢让你无限次点击的底气。不点也能分自动掩码生成在复用同一条链路点哪儿分哪儿之外SAM 还能不给任何提示就把整张图扫一遍。SamAutomaticMaskGeneratorsegment_anything/automatic_mask_generator.py的做法相当朴素在图上铺 32×32 的均匀点网格默认points_per_side32共 1024 个点逐点调用前面讲的predict然后靠三道筛子过滤# SamAutomaticMaskGenerator 默认参数节选 points_per_batch64, pred_iou_thresh0.88, # 模型自报的 IoU 质量分下限 stability_score_thresh0.95, # 掩码对阈值扰动的稳定性下限 box_nms_thresh0.7, # 重复掩码的 NMS 去重它不是目标检测器而是网格扫描 质量过滤预测 IoU 太低的点通常落在纹理混乱的空白处被丢弃同一物体被相邻点重复命中时用 NMS 只留一个多尺度裁剪crop产生的重复结果则偏好来自更小块裁剪的掩码。每个掩码会附送predicted_iou和stability_score两个自评分方便下游再筛。代价与边界把链路走完之后几个工程上的取舍也就清楚了。重活全在编码器。默认 ViT-H 是 32 层、1280 维1024 分辨率写死在输入端长边缩到 1024。所以生产环境的标准打法是编码器只跑一次之后要么反复提示要么把解码器单独搬走——scripts/export_onnx_model.py就是把MaskDecoder导出为 ONNXbackbone 留在 PyTorch。仓库自带的 demo/ 是一个 React 单页应用浏览器里多线程跑导出的 ONNX 解码器体验即点即出掩码是低分辨率重采样来的。真正的分割决策发生在 256×256回原图靠双线性插值——所以头发丝级别的细边界会偏软需要靠多点交互迭代收边。提示的坐标有约定。点必须在模型输入坐标系里predict内部会用ResizeLongestSide帮你从原图坐标换算而框本质是两个带角色的点不是真正的矩形语义。下一步动手克隆仓库安装git clone https://gitcode.com/GitHub_Trending/se/segment-anything然后pip install -e .再按 README.md 下载对应 checkpoint先跑 notebooks/predictor_example.ipynb对同一张图连续点正点/负点观察masks与iou分数怎么变并把第一次的low_res_masks传回第二次调用看迭代效果再跑 notebooks/automatic_mask_generator_example.ipynb 看整图扫描配合scripts/amg.py试命令行批量出掩码想读源码入口是 segment_anything/modeling/sam.py是总装线image_encoder.py、prompt_encoder.py、mask_decoder.py分别对应链路的三段transformer.py里的TwoWayTransformer是解码器的心脏。【免费下载链接】segment-anythingThe repository provides code for running inference with the SegmentAnything Model (SAM), links for downloading the trained model checkpoints, and example notebooks that show how to use the model.项目地址: https://gitcode.com/GitHub_Trending/se/segment-anything创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考