Image Segmentation

def forward_image(self, samples: torch.Tensor):
    """Extract vision features from images."""
    # Forward through ViT backbone
    sam3_features, sam3_pos, sam2_features, sam2_pos = \
        self.vision_backbone.forward(samples)
    
    # Return feature pyramid and positional encodings
    return {
        "vision_features": sam3_features[-1],  # Highest level
        "vision_pos_enc": sam3_pos,
        "backbone_fpn": sam3_features,  # All pyramid levels
        "sam2_backbone_out": sam2_output,  # For tracker
    }

def _encode_prompt(
    self,
    backbone_out,
    find_input,
    geometric_prompt,
    encode_text=True,
):
    # Get text features if using text prompts
    txt_feats = backbone_out["language_features"][:, txt_ids]
    txt_masks = backbone_out["language_mask"][txt_ids]
    
    # Get image features for geometric encoding
    img_feats, img_pos_embeds, vis_feat_sizes = \
        self._get_img_feats(backbone_out, img_ids)
    
    # Encode geometric prompts (boxes, points, masks)
    geo_feats, geo_masks = self.geometry_encoder(
        geo_prompt=geometric_prompt,
        img_feats=img_feats,
        img_sizes=vis_feat_sizes,
        img_pos_embeds=img_pos_embeds,
    )
    
    # Concatenate text and geometric features
    if encode_text:
        prompt = torch.cat([txt_feats, geo_feats], dim=0)
        prompt_mask = torch.cat([txt_masks, geo_masks], dim=1)
    else:
        prompt = geo_feats
        prompt_mask = geo_masks
    
    return prompt, prompt_mask, backbone_out

def forward(
    self,
    src: List[Tensor],  # Multi-level image features
    prompt: Tensor,     # Text + geometric prompts
    src_pos: List[Tensor],  # Positional encodings
    prompt_key_padding_mask: Tensor,
    ...
):
    # Add pooled text to image features for fusion
    if self.add_pooled_text_to_img_feat:
        pooled_text = pool_text_feat(prompt, prompt_key_padding_mask)
        pooled_text = self.text_pooling_proj(pooled_text)[..., None, None]
        src = [x.add_(pooled_text) for x in src]
    
    # Flatten multi-level features
    src_flatten, lvl_pos_embed_flatten = \
        self._prepare_multilevel_features(src, src_pos)
    
    # Process through encoder layers
    output = src_flatten
    for layer in self.layers:
        output = layer(
            tgt=output,
            memory=prompt.transpose(0, 1),  # Cross-attend to prompts
            query_pos=lvl_pos_embed_flatten,
            memory_key_padding_mask=prompt_key_padding_mask,
        )
    
    return {"memory": output, ...}

def forward(
    self,
    tgt,  # Object queries
    memory,  # Encoded image features
    memory_text: Tensor,  # Text features
    text_attention_mask: Tensor,
    ...
):
    bs = memory.shape[1]
    
    # Initialize queries
    query_embed = self.query_embed.weight
    tgt = query_embed.unsqueeze(1).repeat(1, bs, 1)
    
    # Initialize reference boxes
    reference_boxes = self.reference_points.weight.unsqueeze(1)
    reference_boxes = reference_boxes.repeat(1, bs, 1).sigmoid()
    
    intermediate = []
    intermediate_ref_boxes = [reference_boxes]
    
    # Process through decoder layers
    for layer_idx, layer in enumerate(self.layers):
        # Self-attention among queries
        # Cross-attention to text
        # Cross-attention to image
        tgt, presence_out = layer(
            tgt=tgt,
            memory=memory,
            memory_text=memory_text,
            text_attention_mask=text_attention_mask,
            ...
        )
        
        # Iteratively refine boxes
        delta_boxes = self.bbox_embed(tgt)
        reference_boxes = (inverse_sigmoid(reference_boxes) + delta_boxes).sigmoid()
        
        intermediate.append(self.norm(tgt))
        intermediate_ref_boxes.append(reference_boxes)
    
    return intermediate, intermediate_ref_boxes, ...

def _get_rpb_matrix(self, reference_boxes, feat_size):
    """Compute relative positional bias from boxes."""
    H, W = feat_size
    boxes_xyxy = box_cxcywh_to_xyxy(reference_boxes)
    
    # Compute deltas between boxes and spatial positions
    deltas_y = coords_h.view(1, -1, 1) - boxes_xyxy[..., 1:4:2]
    deltas_x = coords_w.view(1, -1, 1) - boxes_xyxy[..., 0:3:2]
    
    # Apply log-space encoding
    deltas_x_log = torch.sign(deltas_x) * torch.log2(torch.abs(deltas_x) + 1.0)
    deltas_y_log = torch.sign(deltas_y) * torch.log2(torch.abs(deltas_y) + 1.0)
    
    # Project to attention bias
    B_x = self.boxRPB_embed_x(deltas_x)  # bs, num_queries, W, n_heads
    B_y = self.boxRPB_embed_y(deltas_y)  # bs, num_queries, H, n_heads
    
    # Combine into 2D bias
    B = B_y.unsqueeze(3) + B_x.unsqueeze(2)  # bs, num_queries, H, W, n_heads
    return B.flatten(2, 3).permute(0, 3, 1, 2)  # bs, n_heads, num_queries, HW

def _run_segmentation_heads(
    self,
    out,
    backbone_out,
    img_ids,
    vis_feat_sizes,
    encoder_hidden_states,
    prompt,
    prompt_mask,
    hs,  # Decoder output queries
):
    if self.segmentation_head is not None:
        # Run mask prediction head
        seg_head_outputs = self.segmentation_head(
            backbone_feats=backbone_out["backbone_fpn"],
            obj_queries=hs,  # Use decoder queries
            image_ids=img_ids,
            encoder_hidden_states=encoder_hidden_states,
            prompt=prompt,
            prompt_mask=prompt_mask,
        )
        
        # Store mask predictions
        for k, v in seg_head_outputs.items():
            if k == "pred_masks":
                out[k] = v