Reference for ultralytics/models/sam/modules/encoders.py

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:
    """
    Initializes an ImageEncoderViT instance for encoding images using Vision Transformer architecture.

    Args:
        img_size (int): Input image size, assumed to be square.
        patch_size (int): Size of image patches.
        in_chans (int): Number of input image channels.
        embed_dim (int): Dimension of patch embeddings.
        depth (int): Number of transformer blocks.
        num_heads (int): Number of attention heads in each block.
        mlp_ratio (float): Ratio of MLP hidden dimension to embedding dimension.
        out_chans (int): Number of output channels from the neck module.
        qkv_bias (bool): If True, adds learnable bias to query, key, value projections.
        norm_layer (Type[nn.Module]): Type of normalization layer to use.
        act_layer (Type[nn.Module]): Type of activation layer to use.
        use_abs_pos (bool): If True, uses absolute positional embeddings.
        use_rel_pos (bool): If True, adds relative positional embeddings to attention maps.
        rel_pos_zero_init (bool): If True, initializes relative positional parameters to zero.
        window_size (int): Size of attention window for windowed attention blocks.
        global_attn_indexes (Tuple[int, ...]): Indices of blocks that use global attention.

    Attributes:
        img_size (int): Dimension of input images.
        patch_embed (PatchEmbed): Module for patch embedding.
        pos_embed (nn.Parameter | None): Absolute positional embedding for patches.
        blocks (nn.ModuleList): List of transformer blocks.
        neck (nn.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)
    """
    super().__init__()
    self.img_size = img_size

    self.patch_embed = PatchEmbed(
        kernel_size=(patch_size, patch_size),
        stride=(patch_size, patch_size),
        in_chans=in_chans,
        embed_dim=embed_dim,
    )

    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))

    self.blocks = nn.ModuleList()
    for i in range(depth):
        block = Block(
            dim=embed_dim,
            num_heads=num_heads,
            mlp_ratio=mlp_ratio,
            qkv_bias=qkv_bias,
            norm_layer=norm_layer,
            act_layer=act_layer,
            use_rel_pos=use_rel_pos,
            rel_pos_zero_init=rel_pos_zero_init,
            window_size=window_size if i not in global_attn_indexes else 0,
            input_size=(img_size // patch_size, img_size // patch_size),
        )
        self.blocks.append(block)

    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 __init__(
    self,
    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,
) -> None:
    """
    Initializes the PromptEncoder module for encoding various types of prompts.

    This module encodes different types of prompts (points, boxes, masks) for input to SAM's mask decoder,
    producing both sparse and dense embeddings.

    Args:
        embed_dim (int): The dimension of the embeddings.
        image_embedding_size (Tuple[int, int]): The spatial size of the image embedding as (H, W).
        input_image_size (Tuple[int, int]): The padded size of the input image as (H, W).
        mask_in_chans (int): The number of hidden channels used for encoding input masks.
        activation (Type[nn.Module]): The activation function to use when encoding input masks.

    Attributes:
        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 (nn.ModuleList): List of point embeddings.
        not_a_point_embed (nn.Embedding): Embedding for points that are not part of any label.
        mask_input_size (Tuple[int, int]): Size of the input mask.
        mask_downscaling (nn.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])
    """
    super().__init__()
    self.embed_dim = embed_dim
    self.input_image_size = input_image_size
    self.image_embedding_size = image_embedding_size
    self.pe_layer = PositionEmbeddingRandom(embed_dim // 2)

    self.num_point_embeddings: int = 4  # pos/neg point + 2 box corners
    point_embeddings = [nn.Embedding(1, embed_dim) for _ in range(self.num_point_embeddings)]
    self.point_embeddings = nn.ModuleList(point_embeddings)
    self.not_a_point_embed = nn.Embedding(1, embed_dim)

    self.mask_input_size = (4 * image_embedding_size[0], 4 * image_embedding_size[1])
    self.mask_downscaling = nn.Sequential(
        nn.Conv2d(1, mask_in_chans // 4, kernel_size=2, stride=2),
        LayerNorm2d(mask_in_chans // 4),
        activation(),
        nn.Conv2d(mask_in_chans // 4, mask_in_chans, kernel_size=2, stride=2),
        LayerNorm2d(mask_in_chans),
        activation(),
        nn.Conv2d(mask_in_chans, embed_dim, kernel_size=1),
    )
    self.no_mask_embed = nn.Embedding(1, embed_dim)

def __init__(
    self,
    out_dim,
    in_dim=256,  # in_dim of pix_feats
):
    """Initializes the MemoryEncoder for encoding pixel features and masks into memory representations."""
    super().__init__()

    self.mask_downsampler = MaskDownSampler(kernel_size=3, stride=2, padding=1)

    self.pix_feat_proj = nn.Conv2d(in_dim, in_dim, kernel_size=1)
    self.fuser = Fuser(CXBlock(dim=256), num_layers=2)
    self.position_encoding = PositionEmbeddingSine(num_pos_feats=64)
    self.out_proj = nn.Identity()
    if out_dim != in_dim:
        self.out_proj = nn.Conv2d(in_dim, out_dim, kernel_size=1)

def __init__(
    self,
    trunk: nn.Module,
    neck: nn.Module,
    scalp: int = 0,
):
    """Initializes the ImageEncoder with trunk and neck networks for feature extraction and refinement."""
    super().__init__()
    self.trunk = trunk
    self.neck = neck
    self.scalp = scalp
    assert (
        self.trunk.channel_list == self.neck.backbone_channel_list
    ), f"Channel dims of trunk {self.trunk.channel_list} and neck {self.neck.backbone_channel_list} do not match."

def __init__(
    self,
    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,
):
    """
    Initializes a modified Feature Pyramid Network (FPN) neck.

    This FPN variant removes the output convolution and uses bicubic interpolation for feature resizing,
    similar to ViT positional embedding interpolation.

    Args:
        d_model (int): Dimension of the model.
        backbone_channel_list (List[int]): List of channel dimensions from the backbone.
        kernel_size (int): Kernel size for the convolutional layers.
        stride (int): Stride for the convolutional layers.
        padding (int): Padding for the convolutional layers.
        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 (Optional[List[int]]): Levels to have top-down features in outputs.

    Examples:
        >>> backbone_channels = [64, 128, 256, 512]
        >>> fpn_neck = FpnNeck(256, backbone_channels)
        >>> print(fpn_neck)
    """
    super().__init__()
    self.position_encoding = PositionEmbeddingSine(num_pos_feats=256)
    self.convs = nn.ModuleList()
    self.backbone_channel_list = backbone_channel_list
    for dim in backbone_channel_list:
        current = nn.Sequential()
        current.add_module(
            "conv",
            nn.Conv2d(
                in_channels=dim,
                out_channels=d_model,
                kernel_size=kernel_size,
                stride=stride,
                padding=padding,
            ),
        )

        self.convs.append(current)
    self.fpn_interp_model = fpn_interp_model
    assert fuse_type in {"sum", "avg"}
    self.fuse_type = fuse_type

    # levels to have top-down features in its outputs
    # e.g. if fpn_top_down_levels is [2, 3], then only outputs of level 2 and 3
    # have top-down propagation, while outputs of level 0 and level 1 have only
    # lateral features from the same backbone level.
    if fpn_top_down_levels is None:
        # default is to have top-down features on all levels
        fpn_top_down_levels = range(len(self.convs))
    self.fpn_top_down_levels = list(fpn_top_down_levels)

def __init__(
    self,
    embed_dim: int = 96,  # initial embed dim
    num_heads: int = 1,  # initial number of heads
    drop_path_rate: float = 0.0,  # stochastic depth
    q_pool: int = 3,  # number of q_pool stages
    q_stride: Tuple[int, int] = (2, 2),  # downsample stride bet. stages
    stages: Tuple[int, ...] = (2, 3, 16, 3),  # blocks per stage
    dim_mul: float = 2.0,  # dim_mul factor at stage shift
    head_mul: float = 2.0,  # head_mul factor at stage shift
    window_pos_embed_bkg_spatial_size: Tuple[int, int] = (14, 14),
    # window size per stage, when not using global att.
    window_spec: Tuple[int, ...] = (
        8,
        4,
        14,
        7,
    ),
    # global attn in these blocks
    global_att_blocks: Tuple[int, ...] = (
        12,
        16,
        20,
    ),
    return_interm_layers=True,  # return feats from every stage
):
    """Initializes the Hiera model, configuring its hierarchical vision transformer architecture."""
    super().__init__()

    assert len(stages) == len(window_spec)
    self.window_spec = window_spec

    depth = sum(stages)
    self.q_stride = q_stride
    self.stage_ends = [sum(stages[:i]) - 1 for i in range(1, len(stages) + 1)]
    assert 0 <= q_pool <= len(self.stage_ends[:-1])
    self.q_pool_blocks = [x + 1 for x in self.stage_ends[:-1]][:q_pool]
    self.return_interm_layers = return_interm_layers

    self.patch_embed = PatchEmbed(
        embed_dim=embed_dim,
        kernel_size=(7, 7),
        stride=(4, 4),
        padding=(3, 3),
    )
    # Which blocks have global att?
    self.global_att_blocks = global_att_blocks

    # Windowed positional embedding (https://arxiv.org/abs/2311.05613)
    self.window_pos_embed_bkg_spatial_size = window_pos_embed_bkg_spatial_size
    self.pos_embed = nn.Parameter(torch.zeros(1, embed_dim, *self.window_pos_embed_bkg_spatial_size))
    self.pos_embed_window = nn.Parameter(torch.zeros(1, embed_dim, self.window_spec[0], self.window_spec[0]))

    dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule

    cur_stage = 1
    self.blocks = nn.ModuleList()

    for i in range(depth):
        dim_out = embed_dim
        # lags by a block, so first block of
        # next stage uses an initial window size
        # of previous stage and final window size of current stage
        window_size = self.window_spec[cur_stage - 1]

        if self.global_att_blocks is not None:
            window_size = 0 if i in self.global_att_blocks else window_size

        if i - 1 in self.stage_ends:
            dim_out = int(embed_dim * dim_mul)
            num_heads = int(num_heads * head_mul)
            cur_stage += 1

        block = MultiScaleBlock(
            dim=embed_dim,
            dim_out=dim_out,
            num_heads=num_heads,
            drop_path=dpr[i],
            q_stride=self.q_stride if i in self.q_pool_blocks else None,
            window_size=window_size,
        )

        embed_dim = dim_out
        self.blocks.append(block)

    self.channel_list = (
        [self.blocks[i].dim_out for i in self.stage_ends[::-1]]
        if return_interm_layers
        else [self.blocks[-1].dim_out]
    )