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:
    """Initialize 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.

    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: nn.Parameter | None = 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:
    """Initialize the PromptEncoder module for encoding various types of prompts.

    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.

    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
):
    """Initialize the MemoryEncoder for encoding pixel features and masks into memory representations.

    This encoder 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).

    Args:
        out_dim (int): Output dimension of the encoded features.
        in_dim (int): Input dimension of the pixel features.

    Examples:
        >>> 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])
    """
    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,
):
    """Initialize the ImageEncoder with trunk and neck networks for feature extraction and refinement.

    This encoder 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.

    Args:
        trunk (nn.Module): The trunk network for initial feature extraction.
        neck (nn.Module): The neck network for feature refinement and positional encoding generation.
        scalp (int): Number of lowest resolution feature levels to discard.

    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'])
    """
    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: list[int] | None = None,
):
    """Initialize 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
):
    """Initialize a Hiera model, a hierarchical vision transformer for efficient multiscale feature extraction.

    Hiera is a hierarchical vision transformer architecture designed for efficient multiscale feature extraction in
    image processing tasks. It uses a series of transformer blocks organized into stages, with optional pooling and
    global attention mechanisms.

    Args:
        embed_dim (int): Initial embedding dimension for the model.
        num_heads (int): Initial number of attention heads.
        drop_path_rate (float): Stochastic depth rate.
        q_pool (int): Number of query pooling stages.
        q_stride (tuple[int, int]): Downsampling stride between stages.
        stages (tuple[int, ...]): Number of blocks per stage.
        dim_mul (float): Dimension multiplier factor at stage transitions.
        head_mul (float): Head multiplier factor at stage transitions.
        window_pos_embed_bkg_spatial_size (tuple[int, int]): Spatial size for window positional embedding
            background.
        window_spec (tuple[int, ...]): Window sizes for each stage when not using global attention.
        global_att_blocks (tuple[int, ...]): Indices of blocks that use global attention.
        return_interm_layers (bool): Whether to return intermediate layer outputs.

    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)
    """
    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 attention?
    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]
    )