Reference for ultralytics/models/sam/modules/encoders.py
Note
Full source code for this file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/models/sam/modules/encoders.py. Help us fix any issues you see by submitting a Pull Request 🛠️. Thank you 🙏!
ultralytics.models.sam.modules.encoders.ImageEncoderViT
Bases: Module
Source code in ultralytics/models/sam/modules/encoders.py
__init__(img_size=1024, patch_size=16, in_chans=3, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4.0, out_chans=256, qkv_bias=True, norm_layer=nn.LayerNorm, act_layer=nn.GELU, use_abs_pos=True, use_rel_pos=False, rel_pos_zero_init=True, window_size=0, global_attn_indexes=())
Parameters:
Name | Type | Description | Default |
---|---|---|---|
img_size |
int
|
Input image size. |
1024
|
patch_size |
int
|
Patch size. |
16
|
in_chans |
int
|
Number of input image channels. |
3
|
embed_dim |
int
|
Patch embedding dimension. |
768
|
depth |
int
|
Depth of ViT. |
12
|
num_heads |
int
|
Number of attention heads in each ViT block. |
12
|
mlp_ratio |
float
|
Ratio of mlp hidden dim to embedding dim. |
4.0
|
qkv_bias |
bool
|
If True, add a learnable bias to query, key, value. |
True
|
norm_layer |
Module
|
Normalization layer. |
LayerNorm
|
act_layer |
Module
|
Activation layer. |
GELU
|
use_abs_pos |
bool
|
If True, use absolute positional embeddings. |
True
|
use_rel_pos |
bool
|
If True, add relative positional embeddings to the attention map. |
False
|
rel_pos_zero_init |
bool
|
If True, zero initialize relative positional parameters. |
True
|
window_size |
int
|
Window size for window attention blocks. |
0
|
global_attn_indexes |
list
|
Indexes for blocks using global attention. |
()
|
Source code in ultralytics/models/sam/modules/encoders.py
ultralytics.models.sam.modules.encoders.PromptEncoder
Bases: Module
Source code in ultralytics/models/sam/modules/encoders.py
111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 |
|
__init__(embed_dim, image_embedding_size, input_image_size, mask_in_chans, activation=nn.GELU)
Encodes prompts for input to SAM's mask decoder.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
embed_dim |
int
|
The prompts' embedding dimension |
required |
image_embedding_size |
tuple(int, int
|
The spatial size of the image embedding, as (H, W). |
required |
input_image_size |
int
|
The padded size of the image as input to the image encoder, as (H, W). |
required |
mask_in_chans |
int
|
The number of hidden channels used for encoding input masks. |
required |
activation |
Module
|
The activation to use when encoding input masks. |
GELU
|
Source code in ultralytics/models/sam/modules/encoders.py
forward(points, boxes, masks)
Embeds different types of prompts, returning both sparse and dense embeddings.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
points |
(tuple(Tensor, Tensor), None)
|
point coordinates and labels to embed. |
required |
boxes |
(Tensor, None)
|
boxes to embed |
required |
masks |
(Tensor, None)
|
masks to embed |
required |
Returns:
Type | Description |
---|---|
Tensor
|
torch.Tensor: sparse embeddings for the points and boxes, with shape BxNx(embed_dim), where N is determined by the number of input points and boxes. |
Tensor
|
torch.Tensor: dense embeddings for the masks, in the shape Bx(embed_dim)x(embed_H)x(embed_W) |
Source code in ultralytics/models/sam/modules/encoders.py
get_dense_pe()
Returns the positional encoding used to encode point prompts, applied to a dense set of points the shape of the image encoding.
Returns:
Type | Description |
---|---|
Tensor
|
torch.Tensor: Positional encoding with shape 1x(embed_dim)x(embedding_h)x(embedding_w) |
Source code in ultralytics/models/sam/modules/encoders.py
ultralytics.models.sam.modules.encoders.PositionEmbeddingRandom
Bases: Module
Positional encoding using random spatial frequencies.
Source code in ultralytics/models/sam/modules/encoders.py
forward(size)
Generate positional encoding for a grid of the specified size.
Source code in ultralytics/models/sam/modules/encoders.py
forward_with_coords(coords_input, image_size)
Positionally encode points that are not normalized to [0,1].
Source code in ultralytics/models/sam/modules/encoders.py
ultralytics.models.sam.modules.encoders.Block
Bases: Module
Transformer blocks with support of window attention and residual propagation blocks
Source code in ultralytics/models/sam/modules/encoders.py
__init__(dim, num_heads, mlp_ratio=4.0, qkv_bias=True, norm_layer=nn.LayerNorm, act_layer=nn.GELU, use_rel_pos=False, rel_pos_zero_init=True, window_size=0, input_size=None)
Parameters:
Name | Type | Description | Default |
---|---|---|---|
dim |
int
|
Number of input channels. |
required |
num_heads |
int
|
Number of attention heads in each ViT block. |
required |
mlp_ratio |
float
|
Ratio of mlp hidden dim to embedding dim. |
4.0
|
qkv_bias |
bool
|
If True, add a learnable bias to query, key, value. |
True
|
norm_layer |
Module
|
Normalization layer. |
LayerNorm
|
act_layer |
Module
|
Activation layer. |
GELU
|
use_rel_pos |
bool
|
If True, add relative positional embeddings to the attention map. |
False
|
rel_pos_zero_init |
bool
|
If True, zero initialize relative positional parameters. |
True
|
window_size |
int
|
Window size for window attention blocks. If it equals 0, then use global attention. |
0
|
input_size |
(tuple(int, int), None)
|
Input resolution for calculating the relative positional parameter size. |
None
|
Source code in ultralytics/models/sam/modules/encoders.py
ultralytics.models.sam.modules.encoders.Attention
Bases: Module
Multi-head Attention block with relative position embeddings.
Source code in ultralytics/models/sam/modules/encoders.py
__init__(dim, num_heads=8, qkv_bias=True, use_rel_pos=False, rel_pos_zero_init=True, input_size=None)
Parameters:
Name | Type | Description | Default |
---|---|---|---|
dim |
int
|
Number of input channels. |
required |
num_heads |
int
|
Number of attention heads. |
8
|
qkv_bias |
bool
|
If True, add a learnable bias to query, key, value. |
True
|
rel_pos_zero_init |
bool
|
If True, zero initialize relative positional parameters. |
True
|
input_size |
(tuple(int, int), None)
|
Input resolution for calculating the relative positional parameter size. |
None
|
Source code in ultralytics/models/sam/modules/encoders.py
ultralytics.models.sam.modules.encoders.PatchEmbed
Bases: Module
Image to Patch Embedding.
Source code in ultralytics/models/sam/modules/encoders.py
__init__(kernel_size=(16, 16), stride=(16, 16), padding=(0, 0), in_chans=3, embed_dim=768)
Parameters:
Name | Type | Description | Default |
---|---|---|---|
kernel_size |
Tuple
|
kernel size of the projection layer. |
(16, 16)
|
stride |
Tuple
|
stride of the projection layer. |
(16, 16)
|
padding |
Tuple
|
padding size of the projection layer. |
(0, 0)
|
in_chans |
int
|
Number of input image channels. |
3
|
embed_dim |
int
|
Patch embedding dimension. |
768
|
Source code in ultralytics/models/sam/modules/encoders.py
ultralytics.models.sam.modules.encoders.window_partition(x, window_size)
Partition into non-overlapping windows with padding if needed. Args: x (tensor): input tokens with [B, H, W, C]. window_size (int): window size.
Returns:
Name | Type | Description |
---|---|---|
windows |
Tensor
|
windows after partition with [B * num_windows, window_size, window_size, C]. |
(Hp, Wp)
|
padded height and width before partition |
Source code in ultralytics/models/sam/modules/encoders.py
ultralytics.models.sam.modules.encoders.window_unpartition(windows, window_size, pad_hw, hw)
Window unpartition into original sequences and removing padding. Args: windows (tensor): input tokens with [B * num_windows, window_size, window_size, C]. window_size (int): window size. pad_hw (Tuple): padded height and width (Hp, Wp). hw (Tuple): original height and width (H, W) before padding.
Returns:
Name | Type | Description |
---|---|---|
x |
Tensor
|
unpartitioned sequences with [B, H, W, C]. |
Source code in ultralytics/models/sam/modules/encoders.py
ultralytics.models.sam.modules.encoders.get_rel_pos(q_size, k_size, rel_pos)
Get relative positional embeddings according to the relative positions of query and key sizes. Args: q_size (int): size of query q. k_size (int): size of key k. rel_pos (Tensor): relative position embeddings (L, C).
Returns:
Type | Description |
---|---|
Tensor
|
Extracted positional embeddings according to relative positions. |
Source code in ultralytics/models/sam/modules/encoders.py
ultralytics.models.sam.modules.encoders.add_decomposed_rel_pos(attn, q, rel_pos_h, rel_pos_w, q_size, k_size)
Calculate decomposed Relative Positional Embeddings from :paper:mvitv2
.
https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py # noqa B950
Args:
attn (Tensor): attention map.
q (Tensor): query q in the attention layer with shape (B, q_h * q_w, C).
rel_pos_h (Tensor): relative position embeddings (Lh, C) for height axis.
rel_pos_w (Tensor): relative position embeddings (Lw, C) for width axis.
q_size (Tuple): spatial sequence size of query q with (q_h, q_w).
k_size (Tuple): spatial sequence size of key k with (k_h, k_w).
Returns:
Name | Type | Description |
---|---|---|
attn |
Tensor
|
attention map with added relative positional embeddings. |