Reference for ultralytics/models/sam/modules/encoders.py
Note
This file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/models/sam/modules/encoders.py. If you spot a problem please help fix it by contributing a Pull Request 🛠️. Thank you 🙏!
ultralytics.models.sam.modules.encoders.ImageEncoderViT
ImageEncoderViT(
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, ...] = (),
)
Bases: Module
An image encoder using Vision Transformer (ViT) architecture for encoding images into a compact latent space.
This class processes images by splitting them into patches, applying transformer blocks, and generating a final encoded representation through a neck module.
Attributes:
Name | Type | Description |
---|---|---|
img_size |
int
|
Dimension of input images, assumed to be square. |
patch_embed |
PatchEmbed
|
Module for patch embedding. |
pos_embed |
Parameter | None
|
Absolute positional embedding for patches. |
blocks |
ModuleList
|
List of transformer blocks for processing patch embeddings. |
neck |
Sequential
|
Neck module to further process the output. |
Methods:
Name | Description |
---|---|
forward |
Processes input through patch embedding, positional embedding, blocks, and neck. |
Examples:
>>> import torch
>>> encoder = ImageEncoderViT(img_size=224, patch_size=16, embed_dim=768, depth=12, num_heads=12)
>>> input_image = torch.randn(1, 3, 224, 224)
>>> output = encoder(input_image)
>>> print(output.shape)
Parameters:
Name | Type | Description | Default |
---|---|---|---|
img_size
|
int
|
Input image size, assumed to be square. |
1024
|
patch_size
|
int
|
Size of image patches. |
16
|
in_chans
|
int
|
Number of input image channels. |
3
|
embed_dim
|
int
|
Dimension of patch embeddings. |
768
|
depth
|
int
|
Number of transformer blocks. |
12
|
num_heads
|
int
|
Number of attention heads in each block. |
12
|
mlp_ratio
|
float
|
Ratio of MLP hidden dimension to embedding dimension. |
4.0
|
out_chans
|
int
|
Number of output channels from the neck module. |
256
|
qkv_bias
|
bool
|
If True, adds learnable bias to query, key, value projections. |
True
|
norm_layer
|
Type[Module]
|
Type of normalization layer to use. |
LayerNorm
|
act_layer
|
Type[Module]
|
Type of activation layer to use. |
GELU
|
use_abs_pos
|
bool
|
If True, uses absolute positional embeddings. |
True
|
use_rel_pos
|
bool
|
If True, adds relative positional embeddings to attention maps. |
False
|
rel_pos_zero_init
|
bool
|
If True, initializes relative positional parameters to zero. |
True
|
window_size
|
int
|
Size of attention window for windowed attention blocks. |
0
|
global_attn_indexes
|
Tuple[int, ...]
|
Indices of blocks that use global attention. |
()
|
Attributes:
Name | Type | Description |
---|---|---|
img_size |
int
|
Dimension of input images. |
patch_embed |
PatchEmbed
|
Module for patch embedding. |
pos_embed |
Parameter | None
|
Absolute positional embedding for patches. |
blocks |
ModuleList
|
List of transformer blocks. |
neck |
Sequential
|
Neck module for final processing. |
Examples:
>>> encoder = ImageEncoderViT(img_size=224, patch_size=16, embed_dim=768, depth=12, num_heads=12)
>>> input_image = torch.randn(1, 3, 224, 224)
>>> output = encoder(input_image)
>>> print(output.shape)
Source code in ultralytics/models/sam/modules/encoders.py
48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 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 |
|
forward
Processes input through patch embedding, positional embedding, transformer blocks, and neck module.
Source code in ultralytics/models/sam/modules/encoders.py
ultralytics.models.sam.modules.encoders.PromptEncoder
PromptEncoder(
embed_dim: int,
image_embedding_size: Tuple[int, int],
input_image_size: Tuple[int, int],
mask_in_chans: int,
activation: Type[nn.Module] = nn.GELU,
)
Bases: Module
Encodes different types of prompts for input to SAM's mask decoder, producing sparse and dense embeddings.
Attributes:
Name | Type | Description |
---|---|---|
embed_dim |
int
|
Dimension of the embeddings. |
input_image_size |
Tuple[int, int]
|
Size of the input image as (H, W). |
image_embedding_size |
Tuple[int, int]
|
Spatial size of the image embedding as (H, W). |
pe_layer |
PositionEmbeddingRandom
|
Module for random position embedding. |
num_point_embeddings |
int
|
Number of point embeddings for different types of points. |
point_embeddings |
ModuleList
|
List of point embeddings. |
not_a_point_embed |
Embedding
|
Embedding for points that are not part of any label. |
mask_input_size |
Tuple[int, int]
|
Size of the input mask. |
mask_downscaling |
Sequential
|
Neural network for downscaling the mask. |
no_mask_embed |
Embedding
|
Embedding for cases where no mask is provided. |
Methods:
Name | Description |
---|---|
get_dense_pe |
Returns the positional encoding used to encode point prompts. |
forward |
Embeds different types of prompts, returning both sparse and dense embeddings. |
Examples:
>>> prompt_encoder = PromptEncoder(256, (64, 64), (1024, 1024), 16)
>>> points = (torch.rand(1, 5, 2), torch.randint(0, 4, (1, 5)))
>>> boxes = torch.rand(1, 2, 2)
>>> masks = torch.rand(1, 1, 256, 256)
>>> sparse_embeddings, dense_embeddings = prompt_encoder(points, boxes, masks)
>>> print(sparse_embeddings.shape, dense_embeddings.shape)
torch.Size([1, 7, 256]) torch.Size([1, 256, 64, 64])
This module encodes different types of prompts (points, boxes, masks) for input to SAM's mask decoder, producing both sparse and dense embeddings.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
embed_dim
|
int
|
The dimension of the embeddings. |
required |
image_embedding_size
|
Tuple[int, int]
|
The spatial size of the image embedding as (H, W). |
required |
input_image_size
|
Tuple[int, int]
|
The padded size of the input image as (H, W). |
required |
mask_in_chans
|
int
|
The number of hidden channels used for encoding input masks. |
required |
activation
|
Type[Module]
|
The activation function to use when encoding input masks. |
GELU
|
Attributes:
Name | Type | Description |
---|---|---|
embed_dim |
int
|
Dimension of the embeddings. |
input_image_size |
Tuple[int, int]
|
Size of the input image as (H, W). |
image_embedding_size |
Tuple[int, int]
|
Spatial size of the image embedding as (H, W). |
pe_layer |
PositionEmbeddingRandom
|
Module for random position embedding. |
num_point_embeddings |
int
|
Number of point embeddings for different types of points. |
point_embeddings |
ModuleList
|
List of point embeddings. |
not_a_point_embed |
Embedding
|
Embedding for points that are not part of any label. |
mask_input_size |
Tuple[int, int]
|
Size of the input mask. |
mask_downscaling |
Sequential
|
Neural network for downscaling the mask. |
Examples:
>>> prompt_encoder = PromptEncoder(256, (64, 64), (1024, 1024), 16)
>>> points = (torch.rand(1, 5, 2), torch.randint(0, 4, (1, 5)))
>>> boxes = torch.rand(1, 2, 2)
>>> masks = torch.rand(1, 1, 256, 256)
>>> sparse_embeddings, dense_embeddings = prompt_encoder(points, boxes, masks)
>>> print(sparse_embeddings.shape, dense_embeddings.shape)
torch.Size([1, 7, 256]) torch.Size([1, 256, 64, 64])
Source code in ultralytics/models/sam/modules/encoders.py
forward
forward(
points: Optional[Tuple[torch.Tensor, torch.Tensor]],
boxes: Optional[torch.Tensor],
masks: Optional[torch.Tensor],
) -> Tuple[torch.Tensor, torch.Tensor]
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. The first tensor contains coordinates with shape (B, N, 2), and the second tensor contains labels with shape (B, N). |
required |
boxes
|
Tensor | None
|
Boxes to embed with shape (B, M, 2, 2), where M is the number of boxes. |
required |
masks
|
Tensor | None
|
Masks to embed with shape (B, 1, H, W). |
required |
Returns:
Type | Description |
---|---|
Tuple[Tensor, Tensor]
|
A tuple containing: - sparse_embeddings (torch.Tensor): Sparse embeddings for points and boxes with shape (B, N, embed_dim). - dense_embeddings (torch.Tensor): Dense embeddings for masks of shape (B, embed_dim, embed_H, embed_W). |
Examples:
>>> encoder = PromptEncoder(256, (64, 64), (1024, 1024), 16)
>>> points = (torch.rand(1, 5, 2), torch.randint(0, 4, (1, 5)))
>>> boxes = torch.rand(1, 2, 2, 2)
>>> masks = torch.rand(1, 1, 256, 256)
>>> sparse_emb, dense_emb = encoder(points, boxes, masks)
>>> print(sparse_emb.shape, dense_emb.shape)
torch.Size([1, 7, 256]) torch.Size([1, 256, 64, 64])
Source code in ultralytics/models/sam/modules/encoders.py
get_dense_pe
Returns the dense positional encoding used for encoding point prompts.
This method generates a positional encoding for a dense set of points matching the shape of the image encoding. The encoding is used to provide spatial information to the model when processing point prompts.
Returns:
Type | Description |
---|---|
Tensor
|
Positional encoding tensor with shape (1, embed_dim, H, W), where H and W are the height and width of the image embedding size, respectively. |
Examples:
>>> prompt_encoder = PromptEncoder(256, (64, 64), (1024, 1024), 16)
>>> dense_pe = prompt_encoder.get_dense_pe()
>>> print(dense_pe.shape)
torch.Size([1, 256, 64, 64])
Source code in ultralytics/models/sam/modules/encoders.py
ultralytics.models.sam.modules.encoders.MemoryEncoder
Bases: Module
Encodes pixel features and masks into a memory representation for efficient image segmentation.
This class processes pixel-level features and masks, fusing them to generate encoded memory representations suitable for downstream tasks in image segmentation models like SAM (Segment Anything Model).
Attributes:
Name | Type | Description |
---|---|---|
mask_downsampler |
MaskDownSampler
|
Module for downsampling input masks. |
pix_feat_proj |
Conv2d
|
Convolutional layer for projecting pixel features. |
fuser |
Fuser
|
Module for fusing pixel features and masks. |
position_encoding |
PositionEmbeddingSine
|
Module for adding positional encoding to features. |
out_proj |
Module
|
Output projection layer, either nn.Identity or nn.Conv2d. |
Methods:
Name | Description |
---|---|
forward |
Processes input pixel features and masks to generate encoded memory representations. |
Examples:
>>> import torch
>>> encoder = MemoryEncoder(out_dim=256, in_dim=256)
>>> pix_feat = torch.randn(1, 256, 64, 64)
>>> masks = torch.randn(1, 1, 64, 64)
>>> encoded_feat, pos = encoder(pix_feat, masks)
>>> print(encoded_feat.shape, pos.shape)
torch.Size([1, 256, 64, 64]) torch.Size([1, 128, 64, 64])
Source code in ultralytics/models/sam/modules/encoders.py
forward
forward(
pix_feat: torch.Tensor, masks: torch.Tensor, skip_mask_sigmoid: bool = False
) -> Tuple[torch.Tensor, torch.Tensor]
Processes pixel features and masks to generate encoded memory representations for segmentation.
Source code in ultralytics/models/sam/modules/encoders.py
ultralytics.models.sam.modules.encoders.ImageEncoder
Bases: Module
Encodes images using a trunk-neck architecture, producing multiscale features and positional encodings.
This class combines a trunk network for feature extraction with a neck network for feature refinement and positional encoding generation. It can optionally discard the lowest resolution features.
Attributes:
Name | Type | Description |
---|---|---|
trunk |
Module
|
The trunk network for initial feature extraction. |
neck |
Module
|
The neck network for feature refinement and positional encoding generation. |
scalp |
int
|
Number of lowest resolution feature levels to discard. |
Methods:
Name | Description |
---|---|
forward |
Processes the input image through the trunk and neck networks. |
Examples:
>>> trunk = SomeTrunkNetwork()
>>> neck = SomeNeckNetwork()
>>> encoder = ImageEncoder(trunk, neck, scalp=1)
>>> image = torch.randn(1, 3, 224, 224)
>>> output = encoder(image)
>>> print(output.keys())
dict_keys(['vision_features', 'vision_pos_enc', 'backbone_fpn'])
Source code in ultralytics/models/sam/modules/encoders.py
forward
Encodes input through patch embedding, positional embedding, transformer blocks, and neck module.
Source code in ultralytics/models/sam/modules/encoders.py
ultralytics.models.sam.modules.encoders.FpnNeck
FpnNeck(
d_model: int,
backbone_channel_list: List[int],
kernel_size: int = 1,
stride: int = 1,
padding: int = 0,
fpn_interp_model: str = "bilinear",
fuse_type: str = "sum",
fpn_top_down_levels: Optional[List[int]] = None,
)
Bases: Module
A Feature Pyramid Network (FPN) neck variant for multiscale feature fusion in object detection models.
This FPN variant removes the output convolution and uses bicubic interpolation for feature resizing, similar to ViT positional embedding interpolation.
Attributes:
Name | Type | Description |
---|---|---|
position_encoding |
PositionEmbeddingSine
|
Sinusoidal positional encoding module. |
convs |
ModuleList
|
List of convolutional layers for each backbone level. |
backbone_channel_list |
List[int]
|
List of channel dimensions from the backbone. |
fpn_interp_model |
str
|
Interpolation mode for FPN feature resizing. |
fuse_type |
str
|
Type of feature fusion, either 'sum' or 'avg'. |
fpn_top_down_levels |
List[int]
|
Levels to have top-down features in outputs. |
Methods:
Name | Description |
---|---|
forward |
Performs forward pass through the FPN neck. |
Examples:
>>> backbone_channels = [64, 128, 256, 512]
>>> fpn_neck = FpnNeck(256, backbone_channels)
>>> inputs = [torch.rand(1, c, 32, 32) for c in backbone_channels]
>>> outputs, positions = fpn_neck(inputs)
>>> print(len(outputs), len(positions))
4 4
This FPN variant removes the output convolution and uses bicubic interpolation for feature resizing, similar to ViT positional embedding interpolation.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
d_model
|
int
|
Dimension of the model. |
required |
backbone_channel_list
|
List[int]
|
List of channel dimensions from the backbone. |
required |
kernel_size
|
int
|
Kernel size for the convolutional layers. |
1
|
stride
|
int
|
Stride for the convolutional layers. |
1
|
padding
|
int
|
Padding for the convolutional layers. |
0
|
fpn_interp_model
|
str
|
Interpolation mode for FPN feature resizing. |
'bilinear'
|
fuse_type
|
str
|
Type of feature fusion, either 'sum' or 'avg'. |
'sum'
|
fpn_top_down_levels
|
Optional[List[int]]
|
Levels to have top-down features in outputs. |
None
|
Examples:
>>> backbone_channels = [64, 128, 256, 512]
>>> fpn_neck = FpnNeck(256, backbone_channels)
>>> print(fpn_neck)
Source code in ultralytics/models/sam/modules/encoders.py
forward
Performs forward pass through the Feature Pyramid Network (FPN) neck.
This method processes a list of input tensors from the backbone through the FPN, applying lateral connections and top-down feature fusion. It generates output feature maps and corresponding positional encodings.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
xs
|
List[Tensor]
|
List of input tensors from the backbone, each with shape (B, C, H, W). |
required |
Returns:
Type | Description |
---|---|
Tuple[List[Tensor], List[Tensor]]
|
A tuple containing: - out (List[torch.Tensor]): List of output feature maps after FPN processing, each with shape (B, d_model, H, W). - pos (List[torch.Tensor]): List of positional encodings corresponding to each output feature map. |
Examples:
>>> fpn_neck = FpnNeck(d_model=256, backbone_channel_list=[64, 128, 256, 512])
>>> inputs = [torch.rand(1, c, 32, 32) for c in [64, 128, 256, 512]]
>>> outputs, positions = fpn_neck(inputs)
>>> print(len(outputs), len(positions))
4 4
Source code in ultralytics/models/sam/modules/encoders.py
ultralytics.models.sam.modules.encoders.Hiera
Hiera(
embed_dim: int = 96,
num_heads: int = 1,
drop_path_rate: float = 0.0,
q_pool: int = 3,
q_stride: Tuple[int, int] = (2, 2),
stages: Tuple[int, ...] = (2, 3, 16, 3),
dim_mul: float = 2.0,
head_mul: float = 2.0,
window_pos_embed_bkg_spatial_size: Tuple[int, int] = (14, 14),
window_spec: Tuple[int, ...] = (8, 4, 14, 7),
global_att_blocks: Tuple[int, ...] = (12, 16, 20),
return_interm_layers=True,
)
Bases: Module
Hierarchical vision transformer for efficient multiscale feature extraction in image processing tasks.
This class implements a Hiera model, which is a hierarchical vision transformer architecture designed for efficient multiscale feature extraction. It uses a series of transformer blocks organized into stages, with optional pooling and global attention mechanisms.
Attributes:
Name | Type | Description |
---|---|---|
window_spec |
Tuple[int, ...]
|
Window sizes for each stage. |
q_stride |
Tuple[int, int]
|
Downsampling stride between stages. |
stage_ends |
List[int]
|
Indices of the last block in each stage. |
q_pool_blocks |
List[int]
|
Indices of blocks where pooling is applied. |
return_interm_layers |
bool
|
Whether to return intermediate layer outputs. |
patch_embed |
PatchEmbed
|
Module for patch embedding. |
global_att_blocks |
Tuple[int, ...]
|
Indices of blocks with global attention. |
window_pos_embed_bkg_spatial_size |
Tuple[int, int]
|
Spatial size for window positional embedding background. |
pos_embed |
Parameter
|
Positional embedding for the background. |
pos_embed_window |
Parameter
|
Positional embedding for the window. |
blocks |
ModuleList
|
List of MultiScaleBlock modules. |
channel_list |
List[int]
|
List of output channel dimensions for each stage. |
Methods:
Name | Description |
---|---|
_get_pos_embed |
Generates positional embeddings by interpolating and combining window and background embeddings. |
forward |
Performs the forward pass through the Hiera model. |
Examples:
>>> model = Hiera(embed_dim=96, num_heads=1, stages=(2, 3, 16, 3))
>>> input_tensor = torch.randn(1, 3, 224, 224)
>>> output_features = model(input_tensor)
>>> for feat in output_features:
... print(feat.shape)
Source code in ultralytics/models/sam/modules/encoders.py
679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 |
|
forward
Performs forward pass through Hiera model, extracting multiscale features from input images.