Reference for ultralytics/models/sam/sam3/geometry_encoders.py#

class Prompt:
    """Utility class to manipulate geometric box prompts.

    Sequences follow PyTorch convention (sequence first, batch second), while masks are batch-first.

    Attributes:
        box_embeddings (torch.Tensor | None): Boxes with shape (N_boxes, B, 4) in normalized CxCyWH format.
        box_mask (torch.Tensor | None): Padding mask with shape (B, N_boxes), True for padded entries. Defaults to all
            False if not provided.
        box_labels (torch.Tensor | None): Long tensor with shape (N_boxes, B) of positive (1) or negative (0) labels.
            Defaults to all positive if not provided.
    """

    def __init__(self, box_embeddings=None, box_mask=None, box_labels=None):
        """Initialize the Prompt object."""
        # Check for null prompt
        if box_embeddings is None:
            self.box_embeddings = None
            self.box_labels = None
            self.box_mask = None
            return

        # Get sequence length, batch size, and device
        box_seq_len = box_embeddings.shape[0]
        bs = box_embeddings.shape[1]
        device = box_embeddings.device

        # Initialize labels and attention mask if not provided
        if box_labels is None:
            box_labels = torch.ones(box_seq_len, bs, device=device, dtype=torch.long)
        if box_mask is None:
            box_mask = torch.zeros(bs, box_seq_len, device=device, dtype=torch.bool)

        # Dimension checks
        assert list(box_embeddings.shape[:2]) == [box_seq_len, bs], (
            f"Wrong dimension for box embeddings. Expected [{box_seq_len}, {bs}, *] got {box_embeddings.shape}"
        )
        assert box_embeddings.shape[-1] == 4, (
            f"Expected box embeddings to have 4 coordinates, got {box_embeddings.shape[-1]}"
        )
        assert list(box_mask.shape) == [bs, box_seq_len], (
            f"Wrong dimension for box mask. Expected [{bs}, {box_seq_len}] got {box_mask.shape}"
        )
        assert list(box_labels.shape) == [box_seq_len, bs], (
            f"Wrong dimension for box labels. Expected [{box_seq_len}, {bs}] got {box_labels.shape}"
        )

        # Device checks
        assert box_embeddings.device == device, (
            f"Expected box embeddings to be on device {device}, got {box_embeddings.device}"
        )
        assert box_mask.device == device, f"Expected box mask to be on device {device}, got {box_mask.device}"
        assert box_labels.device == device, f"Expected box labels to be on device {device}, got {box_labels.device}"

        self.box_embeddings = box_embeddings
        self.box_mask = box_mask
        self.box_labels = box_labels

def append_boxes(self, boxes, labels=None, mask=None):
    """Append box prompts to existing prompts.

    Args:
        boxes (torch.Tensor): Tensor of shape (N_new_boxes, B, 4) with normalized box coordinates.
        labels (torch.Tensor | None): Optional tensor of shape (N_new_boxes, B) with positive/negative labels.
        mask (torch.Tensor | None): Optional tensor of shape (B, N_new_boxes) for attention mask.
    """
    if self.box_embeddings is None:
        # First boxes - initialize
        self.box_embeddings = boxes
        bs = boxes.shape[1]
        box_seq_len = boxes.shape[0]

        if labels is None:
            labels = torch.ones(box_seq_len, bs, device=boxes.device, dtype=torch.long)
        if mask is None:
            mask = torch.zeros(bs, box_seq_len, device=boxes.device, dtype=torch.bool)

        self.box_labels = labels
        self.box_mask = mask
        return

    # Append to existing boxes
    bs = self.box_embeddings.shape[1]
    assert boxes.shape[1] == bs, f"Batch size mismatch: expected {bs}, got {boxes.shape[1]}"

    if labels is None:
        labels = torch.ones(boxes.shape[0], bs, device=boxes.device, dtype=torch.long)
    if mask is None:
        mask = torch.zeros(bs, boxes.shape[0], dtype=torch.bool, device=boxes.device)

    assert list(boxes.shape[:2]) == list(labels.shape[:2]), (
        f"Shape mismatch between boxes {boxes.shape} and labels {labels.shape}"
    )

    # Concatenate using the helper function
    self.box_labels, _ = concat_padded_sequences(
        self.box_labels.unsqueeze(-1), self.box_mask, labels.unsqueeze(-1), mask
    )
    self.box_labels = self.box_labels.squeeze(-1)
    self.box_embeddings, self.box_mask = concat_padded_sequences(self.box_embeddings, self.box_mask, boxes, mask)

class SequenceGeometryEncoder(nn.Module):
    """Encoder for geometric box prompts. Assumes boxes are passed in the "normalized CxCyWH" format.

    Boxes can be encoded with any of the three possibilities:
    - direct projection: linear projection from coordinate space to d_model
    - pooling: RoI align features from the backbone
    - pos encoder: sinusoidal position encoding of the box center plus its width and height

    These three options are mutually compatible and will be summed if multiple are selected.

    The encoded sequence can be further processed with a transformer.
    """

    def __init__(
        self,
        boxes_direct_project: bool,
        boxes_pool: bool,
        boxes_pos_enc: bool,
        d_model: int,
        pos_enc,
        num_layers: int,
        layer: nn.Module,
        roi_size: int = 7,
        add_cls: bool = True,
        add_post_encode_proj: bool = True,
        use_act_ckpt: bool = False,
    ):
        """Initialize the SequenceGeometryEncoder."""
        super().__init__()

        self.d_model = d_model
        self.pos_enc = pos_enc
        self.roi_size = roi_size

        self.label_embed = torch.nn.Embedding(2, self.d_model)  # positive/negative box labels

        # CLS token for pooling
        self.cls_embed = None
        if add_cls:
            self.cls_embed = torch.nn.Embedding(1, self.d_model)

        # Box encoding modules
        assert boxes_direct_project or boxes_pos_enc or boxes_pool, "Error: need at least one way to encode boxes"
        self.boxes_direct_project = nn.Linear(4, self.d_model) if boxes_direct_project else None
        self.boxes_pool_project = nn.Conv2d(self.d_model, self.d_model, self.roi_size) if boxes_pool else None
        self.boxes_pos_enc_project = nn.Linear(self.d_model + 2, self.d_model) if boxes_pos_enc else None

        self.final_proj = None
        if add_post_encode_proj:
            self.final_proj = nn.Linear(self.d_model, self.d_model)
            self.norm = nn.LayerNorm(self.d_model)

        self.img_pre_norm = nn.Identity()
        if self.boxes_pool_project is not None:
            self.img_pre_norm = nn.LayerNorm(self.d_model)

        self.encode = None
        if num_layers > 0:
            assert add_cls, "It's currently highly recommended to add a CLS when using a transformer"
            self.encode = _get_clones(layer, num_layers)
            self.encode_norm = nn.LayerNorm(self.d_model)

        self.use_act_ckpt = use_act_ckpt

def _encode_boxes(self, boxes, boxes_mask, boxes_labels, img_feats: torch.Tensor):
    """Encode boxes using configured encoding methods."""
    boxes_embed = None
    n_boxes, bs = boxes.shape[:2]

    if self.boxes_direct_project is not None:
        proj = self.boxes_direct_project(boxes.to(img_feats.dtype))
        boxes_embed = proj

    if self.boxes_pool_project is not None:
        H, W = img_feats.shape[-2:]

        # Convert boxes to xyxy format and denormalize
        boxes_xyxy = xywh2xyxy(boxes.to(img_feats.dtype))
        scale = torch.tensor([W, H, W, H], dtype=boxes_xyxy.dtype)
        scale = scale.to(device=boxes_xyxy.device, non_blocking=True)
        scale = scale.view(1, 1, 4)
        boxes_xyxy = boxes_xyxy * scale

        # RoI align
        # Scoped for import ultralytics speed: ROI align requires optional torchvision ops.
        from torchvision.ops import roi_align

        sampled = roi_align(img_feats, boxes_xyxy.transpose(0, 1).unbind(0), self.roi_size)
        assert list(sampled.shape) == [
            bs * n_boxes,
            self.d_model,
            self.roi_size,
            self.roi_size,
        ]
        proj = self.boxes_pool_project(sampled)
        proj = proj.view(bs, n_boxes, self.d_model).transpose(0, 1)

        if boxes_embed is None:
            boxes_embed = proj
        else:
            boxes_embed = boxes_embed + proj

    if self.boxes_pos_enc_project is not None:
        cx, cy, w, h = boxes.unbind(-1)
        enc = self.pos_enc.encode_boxes(cx.flatten(), cy.flatten(), w.flatten(), h.flatten())
        enc = enc.view(boxes.shape[0], boxes.shape[1], enc.shape[-1])

        proj = self.boxes_pos_enc_project(enc.to(img_feats.dtype))
        if boxes_embed is None:
            boxes_embed = proj
        else:
            boxes_embed = boxes_embed + proj

    # Add label embeddings
    type_embed = self.label_embed(boxes_labels.long())
    return type_embed + boxes_embed, boxes_mask

def forward(self, geo_prompt: Prompt, img_feats, img_sizes, img_pos_embeds=None):
    """Encode geometric box prompts.

    Args:
        geo_prompt (Prompt): Prompt object containing box embeddings, masks, and labels.
        img_feats (list[torch.Tensor]): List of sequence-first image features from the backbone, each with shape
            (H*W, B, C).
        img_sizes (list[tuple[int, int]]): List of (H, W) tuples for each feature level.
        img_pos_embeds (list[torch.Tensor] | None): Optional position embeddings for image features.

    Returns:
        final_embeds (torch.Tensor): Encoded prompt embeddings with shape (N, B, d_model).
        final_mask (torch.Tensor): Padding mask with shape (B, N), True for padded entries.
    """
    boxes = geo_prompt.box_embeddings
    boxes_mask = geo_prompt.box_mask
    boxes_labels = geo_prompt.box_labels

    seq_first_img_feats = img_feats[-1]  # [H*W, B, C]
    seq_first_img_pos_embeds = (
        img_pos_embeds[-1] if img_pos_embeds is not None else torch.zeros_like(seq_first_img_feats)
    )

    # Prepare image features for pooling if needed
    if self.boxes_pool_project is not None:
        assert len(img_feats) == len(img_sizes)
        cur_img_feat = img_feats[-1]
        cur_img_feat = self.img_pre_norm(cur_img_feat)
        H, W = img_sizes[-1]
        assert cur_img_feat.shape[0] == H * W
        N, C = cur_img_feat.shape[-2:]
        # Reshape to NxCxHxW
        cur_img_feat = cur_img_feat.permute(1, 2, 0)
        cur_img_feat = cur_img_feat.view(N, C, H, W)
        img_feats = cur_img_feat

    final_embeds, final_mask = self._encode_boxes(
        boxes=boxes,
        boxes_mask=boxes_mask,
        boxes_labels=boxes_labels,
        img_feats=img_feats,
    )

    bs = final_embeds.shape[1]
    assert final_mask.shape[0] == bs

    # Add CLS token if configured
    if self.cls_embed is not None:
        cls = self.cls_embed.weight.view(1, 1, self.d_model).repeat(1, bs, 1)
        cls_mask = torch.zeros(bs, 1, dtype=final_mask.dtype, device=final_mask.device)
        final_embeds, final_mask = concat_padded_sequences(final_embeds, final_mask, cls, cls_mask)

    # Final projection
    if self.final_proj is not None:
        final_embeds = self.norm(self.final_proj(final_embeds))

    # Transformer encoding layers
    if self.encode is not None:
        for lay in self.encode:
            final_embeds = lay(
                tgt=final_embeds,
                memory=seq_first_img_feats,
                tgt_key_padding_mask=final_mask,
                pos=seq_first_img_pos_embeds,
            )
        final_embeds = self.encode_norm(final_embeds)

    return final_embeds, final_mask

def is_right_padded(mask: torch.Tensor):
    """Return whether a padding mask (PyTorch convention, 1s for padded values) is right-padded."""
    return (mask.long() == torch.sort(mask.long(), dim=-1)[0]).all()

def concat_padded_sequences(seq1, mask1, seq2, mask2, return_index: bool = False):
    """Concatenate two right-padded sequences into a contiguous, right-padded sequence.

    Following PyTorch convention, tensors are sequence-first and masks are batch-first, with 1s for padded values.

    Args:
        seq1 (torch.Tensor): A tensor of shape (seq1_length, batch_size, hidden_size).
        mask1 (torch.Tensor): A tensor of shape (batch_size, seq1_length).
        seq2 (torch.Tensor): A tensor of shape (seq2_length, batch_size, hidden_size).
        mask2 (torch.Tensor): A tensor of shape (batch_size, seq2_length).
        return_index (bool): If True, also return the index of the ids of the elements of seq2 in the concatenated
            sequence, which can be used to retrieve the elements of seq2.

    Returns:
        (tuple): (concatenated_sequence, concatenated_mask) if return_index is False, otherwise (concatenated_sequence,
            concatenated_mask, index).
    """
    seq1_length, batch_size, hidden_size = seq1.shape
    seq2_length, batch_size, hidden_size = seq2.shape

    assert batch_size == seq1.size(1) == seq2.size(1) == mask1.size(0) == mask2.size(0)
    assert hidden_size == seq1.size(2) == seq2.size(2)
    assert seq1_length == mask1.size(1)
    assert seq2_length == mask2.size(1)

    torch._assert(is_right_padded(mask1), "Mask is not right padded")
    torch._assert(is_right_padded(mask2), "Mask is not right padded")

    actual_seq1_lengths = (~mask1).sum(dim=-1)
    actual_seq2_lengths = (~mask2).sum(dim=-1)

    final_lengths = actual_seq1_lengths + actual_seq2_lengths
    max_length = seq1_length + seq2_length
    concatenated_mask = (
        torch.arange(max_length, device=seq2.device)[None].repeat(batch_size, 1) >= final_lengths[:, None]
    )

    # (max_len, batch_size, hidden_size)
    concatenated_sequence = torch.zeros((max_length, batch_size, hidden_size), device=seq2.device, dtype=seq2.dtype)
    concatenated_sequence[:seq1_length, :, :] = seq1

    # At this point, the element of seq1 are in the right place
    # We just need to shift the elements of seq2

    index = torch.arange(seq2_length, device=seq2.device)[:, None].repeat(1, batch_size)
    index = index + actual_seq1_lengths[None]

    concatenated_sequence = concatenated_sequence.scatter(0, index[:, :, None].expand(-1, -1, hidden_size), seq2)

    if return_index:
        return concatenated_sequence, concatenated_mask, index

    return concatenated_sequence, concatenated_mask