在 Devin 中试用 DeepWiki 处理私有仓库
DeepWiki

facebookresearch/segment-anything

菜单

图像处理管道

相关源文件

本页详细介绍了 Segment Anything Model (SAM) 中的图像处理管道，说明了输入图像如何被转换、处理并流经模型架构以生成掩模。该管道是确保各种尺寸和格式的图像能够被模型高效处理的关键组成部分。有关整体模型架构的信息，请参阅 SAM 模型架构。

管道概述

SAM 图像处理管道包含几个关键阶段，用于准备图像进行处理、提取特征以及将结果转换回原始图像空间。理解此管道对于有效使用模型和扩展其功能至关重要。

来源: segment_anything/predictor.py34-90 segment_anything/utils/transforms.py16-102

图像缩放与变换

管道的第一步是将输入图像转换为适合模型的格式。这由 ResizeLongestSide 类处理，该类在保持纵横比的同时将图像缩放到其最长维度的目标长度。

ResizeLongestSide 变换

来源: segment_anything/utils/transforms.py16-102

ResizeLongestSide 类提供了以下方法：

  1. 在保持纵横比的同时缩放图像
  2. 将坐标从原始图像空间转换为缩放后的空间
  3. 转换边界框以匹配缩放后的图像

此变换至关重要，因为：

  • SAM 的图像编码器期望特定的输入尺寸（通常为 1024×1024 像素）
  • 保持纵横比可防止图像特征失真
  • 一致的缩放确保基于坐标的提示能正确工作

图像变换过程

来源: segment_anything/predictor.py34-62 segment_anything/utils/transforms.py26-32

SamPredictor.set_image() 中的变换过程包括：

  1. 格式验证：确保图像是预期的颜色格式（RGB 或 BGR）
  2. 缩放：使用 ResizeLongestSide.apply_image() 缩放图像
  3. 张量转换：将 numpy 数组转换为 PyTorch 张量
  4. 维度置换：从 HWC（高度、宽度、通道）更改为 CHW（通道、高度、宽度）格式
  5. 添加批次维度：添加批次维度以创建 BCHW 格式

特征提取

图像经过正确变换后，会通过模型的预处理函数并进入图像编码器以提取图像嵌入。

来源: segment_anything/predictor.py62-90

特征提取过程

  1. 变换后的图像由模型的 preprocess 方法处理
  2. 预处理后的图像被传递给图像编码器
  3. 生成的特征（嵌入）存储在 SamPredictor 实例中
  4. 原始图像大小和输入大小（变换后）都会被存储以备后用

此步骤每张图像只执行一次，允许高效地处理多个提示，而无需重新编码图像。

提示变换

当 SAM 与交互式提示（点、框或掩模）一起使用时，这些提示必须转换为与缩放后的图像空间匹配。

来源: segment_anything/predictor.py136-152 segment_anything/utils/transforms.py33-53

提示变换

  1. 点坐标使用 ResizeLongestSide.apply_coords() 进行变换
  2. 边界框使用 ResizeLongestSide.apply_boxes() 进行变换
  3. 变换应用了与图像相同的缩放因子
  4. 变换后的提示被转换为 PyTorch 张量

这确保了在原始图像坐标空间中提供的提示能够与从缩放图像中提取的特征正确对齐。

掩模预测和后处理

在图像编码和提示变换后，模型会预测掩模，然后将这些掩模放大到原始图像尺寸。

来源: segment_anything/predictor.py168-243

掩模预测过程

  1. 提示通过模型的提示编码器编码为稀疏和密集嵌入
  2. 掩模解码器使用这些嵌入以及先前计算的图像嵌入
  3. 解码器生成低分辨率掩模（通常为 256×256）
  4. 这些掩模使用模型的 postprocess_masks 方法放大到原始图像大小
  5. 应用阈值将掩模从 logits 转换为二进制掩模（除非设置了 return_logits=True）

放大过程考虑了模型内部分辨率、变换后的输入尺寸和原始图像尺寸之间的缩放因子。

完整数据流

下图说明了通过图像处理管道的完整数据流

来源: segment_anything/predictor.py34-243

关键实现细节

图像大小处理

管道会仔细跟踪三种不同的图像大小：

  1. 原始大小：输入图像的尺寸（高，宽）
  2. 输入大小：经过 ResizeLongestSide 变换后的尺寸（高'，宽'）
  3. 模型大小：图像编码器使用的内部尺寸（通常为 1024×1024）

这些尺寸用于确保输入和输出的正确缩放。

维度格式化

SAM 期望图像采用 BCHW 格式（批次、通道、高度、宽度），这与常用的图像 HWC 格式不同。管道处理了此转换

Input: HWC format (numpy array)
↓ transform.apply_image()
Resized: HWC format (numpy array)
↓ torch.as_tensor()
Tensor: HWC format (torch tensor)
↓ permute(2, 0, 1).contiguous()[None, :, :, :]
Output: BCHW format (torch tensor)

来源: segment_anything/predictor.py55-60

缓存机制

管道中的一项关键优化是缓存图像嵌入。这使得：

  1. 每张图像计算一次图像嵌入
  2. 无需重新计算嵌入即可高效处理多个提示
  3. 实现低延迟的交互式分割

这是通过 set_image() 和 set_torch_image() 方法实现的，这些方法将嵌入存储在 self.features 属性中。

来源: segment_anything/predictor.py34-90 segment_anything/predictor.py245-256

图像处理 API

图像处理管道的主要接口是 SamPredictor 类，它提供了以下关键方法：

方法目的
set_image(image, image_format)处理图像并计算嵌入
set_torch_image(transformed_image, original_image_size)设置已转换的图像
predict(point_coords, point_labels, box, mask_input, ...)从提示预测掩模
get_image_embedding()返回存储的图像嵌入
reset_image()清除当前设置的图像

来源: segment_anything/predictor.py17-268

这些方法封装了图像处理管道的各个阶段，为在应用程序中使用 SAM 提供了清晰的接口。