docstr: Predict masks given image and prompt embeddings.
  
  Arguments:
    image_embeddings (torch.Tensor): the embeddings from the image encoder
    image_pe (torch.Tensor): positional encoding with the shape of image_embeddings
    sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes
    dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs
    multimask_output (bool): Whether to return multiple masks or a single
      mask.
  
  Returns:
    torch.Tensor: batched predicted masks
    torch.Tensor: batched predictions of mask quality
docstr_tok: ["Predict","masks","given","image","and","prompt","embeddings",".","Arguments",":","image_embeddings","(","torch.Tensor",")",":","the","embeddings","from","the","image","encoder","image_pe","(","torch.Tensor",")",":","positional","encoding","with","the","shape","of","image_embeddings","sparse_prompt_embeddings","(","torch.Tensor",")",":","the","embeddings","of","the","points","and","boxes","dense_prompt_embeddings","(","torch.Tensor",")",":","the","embeddings","of","the","mask","inputs","multimask_output","(","bool",")",":","Whether","to","return","multiple","masks","or","a","single","mask",".","Returns",":","torch.Tensor",":","batched","predicted","masks","torch.Tensor",":","batched","predictions","of","mask","quality"]
code: def forward(
      self,
      image_embeddings: torch.Tensor,
      image_pe: torch.Tensor,
      sparse_prompt_embeddings: torch.Tensor,
      dense_prompt_embeddings: torch.Tensor,
      multimask_output: bool,
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """
      Predict masks given image and prompt embeddings.
  
      Arguments:
        image_embeddings (torch.Tensor): the embeddings from the image encoder
        image_pe (torch.Tensor): positional encoding with the shape of image_embeddings
        sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes
        dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs
        multimask_output (bool): Whether to return multiple masks or a single
          mask.
  
      Returns:
        torch.Tensor: batched predicted masks
        torch.Tensor: batched predictions of mask quality
      """
      masks, iou_pred = self.predict_masks(
          image_embeddings=image_embeddings,
          image_pe=image_pe,
          sparse_prompt_embeddings=sparse_prompt_embeddings,
          dense_prompt_embeddings=dense_prompt_embeddings,
      )
  
      # Select the correct mask or masks for output
      if multimask_output:
          mask_slice = slice(1, None)
      else:
          mask_slice = slice(0, 1)
      masks = masks[:, mask_slice, :, :]
      iou_pred = iou_pred[:, mask_slice]
  
      # Prepare output
      return masks, iou_pred
  
file_path: segment-anything/segment_anything/modeling/mask_decoder.py
imports: import torch
  from torch import nn
  from torch.nn import functional
  from typing import List, Tuple, Type
  from .common import LayerNorm2d
label_desc: Use image node_id 2 for calling the MaskDecoder obj's underlying member method code with example usage: obj.forward(image_embeddings, image_pe, sparse_prompt_embeddings, dense_prompt_embeddings, multimask_output) and returns: masks,  iou_pred