【技术追踪】SAM（Segment Anything Model）代码解析与结构绘制之Image Encoder
论文：Segment Anything
代码：https://github.com/facebookresearch/segment-anything
1. 使用SAM
尽管官方demo玩的很花很溜，但只有能够本地运行起来，才能够查看中间过程不是，基于这篇文章，使用官方的狗狗图像，采用sam_vit_b_01ec64.pth模型，给定point，完成狗狗的分割。
（1）狗狗图像：
（2）运行代码：
import cv2
import matplotlib.pyplot as plt
import numpy as np
from segment_anything import sam_model_registry, SamPredictor
import torch
def show_mask(mask, ax, random_color=False):
if random_color:
color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
else:
color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6])
h, w = mask.shape[-2:]
mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
ax.imshow(mask_image)
return mask_image
def show_points(coords, labels, ax, marker_size=375):
pos_points = coords[labels == 1]
neg_points = coords[labels == 0]
ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white',
linewidth=1.25)
ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white',
linewidth=1.25)
sam_checkpoint = "./sam_vit_b_01ec64.pth"
device = "cuda" if torch.cuda.is_available() else "cpu"
model_type = "vit_b"
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)
predictor = SamPredictor(sam)
image = cv2.imread("./test image/image dog.jpg")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
predictor.set_image(image)
input_point = np.array([[1300, 800]])
input_label = np.array([1])
plt.figure(figsize=(10, 10))
plt.imshow(image)
show_points(input_point, input_label, plt.gca())
plt.axis('off')
plt.show()
masks, scores, logits = predictor.predict(
point_coords=input_point,
point_labels=input_label,
multimask_output=True,
)
print(scores)
index = np.argmax(scores)
plt.figure(figsize=(10, 10))
plt.imshow(image)
show_mask(masks[index], plt.gca())
show_points(input_point, input_label, plt.gca())
plt.title(f"Mask {index + 1}, Score: {scores[index]:.3f}", fontsize=18)
plt.axis('off')
plt.show()
（3）输出结果：
2. Image Encoder代码解析
（1）set_image函数
位置：【segment_anything/predictor.py --> SamPredictor类 --> set_image函数】
作用： 图像预处理：缩放、转换为Tensor，通道调整，调用set_torch_image函数
本例中狗狗图像，即输入image的 [ H , W , C ] {[H, W, C]} [H,W,C] 大小为 [ 1365 , 2048 , 3 ] {[1365, 2048, 3]} [1365,2048,3]
def set_image(
self,
image: np.ndarray,
image_format: str = "RGB",
) -> None:
assert image_format in [
"RGB",
"BGR",
], f"image_format must be in ['RGB', 'BGR'], is {image_format}."
if image_format != self.model.image_format:
image = image[..., ::-1]
# Transform the image to the form expected by the model
# 输入image: ndarray->(H, W, 3)=(1365, 2048, 3)
# input_image: ndarray->(H*1024/W, 1024, 3)=(683, 1024, 3)
input_image = self.transform.apply_image(image) # 等比缩放图像至长边为1024
# 转换为tensor形式：input_image_torch: tensor->[683, 1024, 3]
input_image_torch = torch.as_tensor(input_image, device=self.device)
# 通道调整：input_image_torch: tensor->[1, 3, 683, 1024]
input_image_torch = input_image_torch.permute(2, 0, 1).contiguous()[None, :, :, :]
# 调用set_torch_image函数，传入参数input_image_torch与原始图像大小(1365, 2048)
self.set_torch_image(input_image_torch, image.shape[:2])
（2）set_torch_image函数
位置：【segment_anything/predictor.py --> SamPredictor类 --> set_torch_image函数】
作用： 图像预处理，调用image_encoder，实现图像嵌入
def set_torch_image(
self,
transformed_image: torch.Tensor,
original_image_size: Tuple[int, ...],
) -> None:
assert (
len(transformed_image.shape) == 4
and transformed_image.shape[1] == 3
and max(*transformed_image.shape[2:]) == self.model.image_encoder.img_size
), f"set_torch_image input must be BCHW with long side {self.model.image_encoder.img_size}."
self.reset_image()
self.original_size = original_image_size # 原始图像大小(H, W)=(1365, 2048)
self.input_size = tuple(transformed_image.shape[-2:]) # 输入图像大小(683, 1024)
# transformed_image.size():[1, 3, H*1024/W, 1024]————>归一化且填充到正方形
input_image = self.model.preprocess(transformed_image) # input_image.size():[1, 3, 1024, 1024]
self.features = self.model.image_encoder(input_image) # feature.size():[1, 256, 64, 64]
self.is_image_set = True
（3）preprocess函数
位置：【segment_anything/modeling/sam.py --> sam类 --> preprocess函数】
作用： 归一化图像并将其填充为正方形
def preprocess(self, x: torch.Tensor) -> torch.Tensor:
"""Normalize pixel values and pad to a square input."""
# 归一化, 均值和标准差已经定义好了, 至于为什么是这个哩, 猜测可能是整个数据集的
# pixel_mean=[123.675, 116.28, 103.53], pixel_std=[58.395, 57.12, 57.375]
x = (x - self.pixel_mean) / self.pixel_std
# Pad
h, w = x.shape[-2:] # 输入图像大小 h=683, w=1024
# Image Encoder的图像输入大小为1024
padh = self.image_encoder.img_size - h # 1024-683=341
padw = self.image_encoder.img_size - w # 1024-1024=0
x = F.pad(x, (0, padw, 0, padh)) # 补零填充, x.size=[1, 3, 1024, 1024]
return x
（4）ImageEncoderViT类
位置：【segment_anything/modeling/image_encoder.py -->ImageEncoderViT类】
作用： 实现图像嵌入，主要包括patch_embed、block和neck三个部分
class ImageEncoderViT(nn.Module):
def __init__(
self,
img_size: int = 1024,
patch_size: int = 16,
in_chans: int = 3,
embed_dim: int = 768,
depth: int = 12,
num_heads: int = 12,
mlp_ratio: float = 4.0,
out_chans: int = 256,
qkv_bias: bool = True,
norm_layer: Type[nn.Module] = nn.LayerNorm,
act_layer: Type[nn.Module] = nn.GELU,
use_abs_pos: bool = True,
use_rel_pos: bool = False,
rel_pos_zero_init: bool = True,
window_size: int = 0,
global_attn_indexes: Tuple[int, ...] = (),
) -> None:
"""
Args:
img_size (int): Input image size.
patch_size (int): Patch size.
in_chans (int): Number of input image channels.
embed_dim (int): Patch embedding dimension.
depth (int): Depth of ViT.
num_heads (int): Number of attention heads in each ViT block.
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
qkv_bias (bool): If True, add a learnable bias to query, key, value.
norm_layer (nn.Module): Normalization layer.
act_layer (nn.Module): Activation layer.
use_abs_pos (bool): If True, use absolute positional embeddings.
use_rel_pos (bool): If True, add relative positional embeddings to the attention map.
rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
window_size (int): Window size for window attention blocks.
global_attn_indexes (list): Indexes for blocks using global attention.
"""
super().__init__()
self.img_size = img_size # 输入图像大小1024
# 将图像划分为Patch
self.patch_embed = PatchEmbed(
kernel_size=(patch_size, patch_size), # 卷积核大小(16, 16)
stride=(patch_size, patch_size), # 卷积核步长(16, 16)
in_chans=in_chans, # 输入图像通道=3
embed_dim=embed_dim, # patch嵌入维度=768
)
# 位置嵌入
self.pos_embed: Optional[nn.Parameter] = None
if use_abs_pos:
# Initialize absolute positional embedding with pretrain image size.
self.pos_embed = nn.Parameter(
torch.zeros(1, img_size // patch_size, img_size // patch_size, embed_dim)
) # 可学习参数[1, 64, 64, 768]
# Block模块
self.blocks = nn.ModuleList()
for i in range(depth):
block = Block(
dim=embed_dim, # 嵌入维度=768
num_heads=num_heads, # multi-head注意机制多头的数目=12
mlp_ratio=mlp_ratio, # MLP隐藏层的维度变换因子=4
qkv_bias=qkv_bias, # qkv全连接层的偏置=True
norm_layer=norm_layer, # 归一化层: nn.LayerNorm
act_layer=act_layer, # 激活函数层: nn.GELU
use_rel_pos=use_rel_pos, # 是否添加相对位置嵌入=False
rel_pos_zero_init=rel_pos_zero_init, # 零初始化相对位置参数=True
# sam_vit_b中global_attn_indexes=encoder_global_attn_indexes=[2, 5, 8, 11]
# 12个Block中的window_size[14,14,0,14,14,0,14,14,0,14,14,0]
window_size=window_size if i not in global_attn_indexes else 0,
input_size=(img_size // patch_size, img_size // patch_size), # 输入大小(64, 64)
)
self.blocks.append(block)
# 输出neck模块
self.neck = nn.Sequential(
nn.Conv2d(
embed_dim,
out_chans,
kernel_size=1,
bias=False,
),
LayerNorm2d(out_chans),
nn.Conv2d(
out_chans,
out_chans,
kernel_size=3,
padding=1,
bias=False,
),
LayerNorm2d(out_chans),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# 输入x.size():[1, 3, 1024, 1024]
x = self.patch_embed(x) # [1, 64, 64, 768]
# 添加位置嵌入
if self.pos_embed is not None:
x = x + self.pos_embed # [1, 64, 64, 768]
# attention模块
for blk in self.blocks:
x = blk(x) # [1, 64, 64, 768]
x = self.neck(x.permute(0, 3, 1, 2)) # 输出x.size():[1, 256, 64, 64]
return x
①patch_embed
class PatchEmbed(nn.Module):
def __init__(
self,
kernel_size: Tuple[int, int] = (16, 16),
stride: Tuple[int, int] = (16, 16),
padding: Tuple[int, int] = (0, 0),
in_chans: int = 3,
embed_dim: int = 768,
) -> None:
super().__init__()
self.proj = nn.Conv2d(
in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.proj(x) # [1, 3, 1024, 1024]——>[1, 768, 64, 64]
# B C H W -> B H W C
x = x.permute(0, 2, 3, 1) # [1, 64, 64, 768]
return x
②Block
class Block(nn.Module):
def __init__(
self,
dim: int,
num_heads: int,
mlp_ratio: float = 4.0,
qkv_bias: bool = True,
norm_layer: Type[nn.Module] = nn.LayerNorm,
act_layer: Type[nn.Module] = nn.GELU,
use_rel_pos: bool = False,
rel_pos_zero_init: bool = True,
window_size: int = 0,
input_size: Optional[Tuple[int, int]] = None,
) -> None:
super().__init__()
self.norm1 = norm_layer(dim) # 归一化层nn.LayerNorm
# attention模块
self.attn = Attention(
dim,
num_heads=num_heads,
qkv_bias=qkv_bias,
use_rel_pos=use_rel_pos,
rel_pos_zero_init=rel_pos_zero_init,
input_size=input_size if window_size == 0 else (window_size, window_size),
)
self.norm2 = norm_layer(dim) # 归一化层nn.LayerNorm
# MLP模块, mlp_ratio=4, act_layer=nn.GELU
self.mlp = MLPBlock(embedding_dim=dim, mlp_dim=int(dim * mlp_ratio), act=act_layer)
self.window_size = window_size # 窗口大小=14或0
def forward(self, x: torch.Tensor) -> torch.Tensor:
shortcut = x # [1, 64, 64, 768]
x = self.norm1(x)
# Window partition
if self.window_size > 0:
H, W = x.shape[1], x.shape[2] # H=64, W=64
x, pad_hw = window_partition(x, self.window_size) # x.size():[25, 14, 14, 768], Pad_hw.size():[70, 70]
x = self.attn(x) # [25, 14, 14, 768]
# Reverse window partition
if self.window_size > 0:
x = window_unpartition(x, self.window_size, pad_hw, (H, W)) # [1, 64, 64, 768]
x = shortcut + x # 残差连接
x = x + self.mlp(self.norm2(x)) # [1, 64, 64, 768]
return x
window_partition函数：不重叠窗口划分
def window_partition(x: torch.Tensor, window_size: int) -> Tuple[torch.Tensor, Tuple[int, int]]:
B, H, W, C = x.shape # [1, 64, 64, 768]
pad_h = (window_size - H % window_size) % window_size # 需要填充的高度=6
pad_w = (window_size - W % window_size) % window_size # 需要填充的宽度=6
if pad_h > 0 or pad_w > 0:
x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h)) # 填充为: [1, 70, 70, 768]
Hp, Wp = H + pad_h, W + pad_w # Hp=70, Wp=70
# 重塑为[1, 5, 14, 5, 14, 768]
x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C)
# [25, 14, 14, 768]
windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
return windows, (Hp, Wp)
Attention类：多头注意力机制
class Attention(nn.Module):
def __init__(
self,
dim: int,
num_heads: int = 8,
qkv_bias: bool = True,
use_rel_pos: bool = False,
rel_pos_zero_init: bool = True,
input_size: Optional[Tuple[int, int]] = None,
) -> None:
super().__init__()
self.num_heads = num_heads # head数目=12
head_dim = dim // num_heads # 768/12=64
self.scale = head_dim**-0.5 # 0.125
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) # (768, 768*3)
self.proj = nn.Linear(dim, dim)
self.use_rel_pos = use_rel_pos
if self.use_rel_pos:
assert (
input_size is not None
), "Input size must be provided if using relative positional encoding."
# initialize relative positional embeddings
self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim))
self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim))
def forward(self, x: torch.Tensor) -> torch.Tensor:
B, H, W, _ = x.shape # B=25, H=14, W=14
# qkv with shape (3, B, nHead, H * W, C)
# [25,14,14,768]->[25,14,14,2304]->[25,14*14,3,12,64]->[3,25,12,196,64]
qkv = self.qkv(x).reshape(B, H * W, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
# q, k, v with shape (B * nHead, H * W, C)=[25*12,14*14,64]=[300,196,64]
q, k, v = qkv.reshape(3, B * self.num_heads, H * W, -1).unbind(0)
attn = (q * self.scale) @ k.transpose(-2, -1) # [300,196,196]
# 使用相对位置编码
if self.use_rel_pos:
attn = add_decomposed_rel_pos(attn, q, self.rel_pos_h, self.rel_pos_w, (H, W), (H, W))
attn = attn.softmax(dim=-1) # [300,196,196]
# [300,196,196]->[300,196,64]->[25,12,14,14,64]->[25,14,14,12,64]->[25,14,14,768]
x = (attn @ v).view(B, self.num_heads, H, W, -1).permute(0, 2, 3, 1, 4).reshape(B, H, W, -1)
x = self.proj(x) # [25,14,14,768]
return x
获取相对位置编码：
def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor:
max_rel_dist = int(2 * max(q_size, k_size) - 1) # 27
# Interpolate rel pos if needed.
if rel_pos.shape[0] != max_rel_dist:
# Interpolate rel pos.
rel_pos_resized = F.interpolate(
rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1),
size=max_rel_dist,
mode="linear",
)
rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0)
else:
rel_pos_resized = rel_pos # [27,64]
# Scale the coords with short length if shapes for q and k are different.
# size[14,1]:[0,1,2,3,4,5,6,7,8,9,10,11,12,13]
q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0)
# size[1,14]:[0,1,2,3,4,5,6,7,8,9,10,11,12,13]
k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0)
# size[14,14]:相对位置编码,右上角为0,左下角为26,沿x=y对称
relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0)
return rel_pos_resized[relative_coords.long()] # [14,14,64]
relative_coords编码如下：
添加相对位置编码：
def add_decomposed_rel_pos(
attn: torch.Tensor,
q: torch.Tensor,
rel_pos_h: torch.Tensor,
rel_pos_w: torch.Tensor,
q_size: Tuple[int, int],
k_size: Tuple[int, int],
) -> torch.Tensor:
q_h, q_w = q_size # (14,14)
k_h, k_w = k_size # (14,14)
# rel_pos_h=rel_pos_w=[27,64]
Rh = get_rel_pos(q_h, k_h, rel_pos_h) # 获取相对位置编码(14,14,64)
Rw = get_rel_pos(q_w, k_w, rel_pos_w) # 获取相对位置编码(14,14,64)
B, _, dim = q.shape # B=300, dim=64
r_q = q.reshape(B, q_h, q_w, dim) # [300, 14, 14, 64]
rel_h = torch.einsum("bhwc,hkc->bhwk", r_q, Rh) # [300,14,14,14]
rel_w = torch.einsum("bhwc,wkc->bhwk", r_q, Rw) # [300,14,14,14]
# rel_h[:, :, :, :, None]=rel_w[:, :, :, None, :]=[300,14,14,14,1]
# attn=[300,196,196]->[300,14,14,14,14]->[300,196,196]
attn = (
attn.view(B, q_h, q_w, k_h, k_w) + rel_h[:, :, :, :, None] + rel_w[:, :, :, None, :]
).view(B, q_h * q_w, k_h * k_w)
return attn
window_unpartition函数：恢复原始中间特征尺寸
def window_unpartition(
windows: torch.Tensor, window_size: int, pad_hw: Tuple[int, int], hw: Tuple[int, int]
) -> torch.Tensor:
Hp, Wp = pad_hw # (70,70)
H, W = hw # (64,64)
B = windows.shape[0] // (Hp * Wp // window_size // window_size) # B=1
# [1,5,5,14,14,768]
x = windows.view(B, Hp // window_size, Wp // window_size, window_size, window_size, -1)
x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, Hp, Wp, -1) # [1,70,70,768]
if Hp > H or Wp > W:
x = x[:, :H, :W, :].contiguous() # 去掉填充元素[1,64,64,768]
return x
MLP模块：
class MLPBlock(nn.Module):
def __init__(
self,
embedding_dim: int,
mlp_dim: int,
act: Type[nn.Module] = nn.GELU,
) -> None:
super().__init__()
self.lin1 = nn.Linear(embedding_dim, mlp_dim)
self.lin2 = nn.Linear(mlp_dim, embedding_dim)
self.act = act()
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.lin2(self.act(self.lin1(x)))
3. ImageEncoderViT结构绘制
（1）结构打印
ImageEncoderViT(
(patch_embed): PatchEmbed(
(proj): Conv2d(3, 768, kernel_size=(16, 16), stride=(16, 16))
)
(blocks): ModuleList(
(0-11): 12 x Block(
(norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
(attn): Attention(
(qkv): Linear(in_features=768, out_features=2304, bias=True)
(proj): Linear(in_features=768, out_features=768, bias=True)
)
(norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
(mlp): MLPBlock(
(lin1): Linear(in_features=768, out_features=3072, bias=True)
(lin2): Linear(in_features=3072, out_features=768, bias=True)
(act): GELU(approximate='none')
)
)
)
(neck): Sequential(
(0): Conv2d(768, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(1): LayerNorm2d()
(2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(3): LayerNorm2d()
)
)
（2）结构绘制
