docstr: Predicts masks end-to-end from provided images and prompts.
  If prompts are not known in advance, using SamPredictor is
  recommended over calling the model directly.
  
  Arguments:
    batched_input (list(dict)): A list over input images, each a
      dictionary with the following keys. A prompt key can be
      excluded if it is not present.
        'image': The image as a torch tensor in 3xHxW format,
          already transformed for input to the model.
        'original_size': (tuple(int, int)) The original size of
          the image before transformation, as (H, W).
        'point_coords': (torch.Tensor) Batched point prompts for
          this image, with shape BxNx2. Already transformed to the
          input frame of the model.
        'point_labels': (torch.Tensor) Batched labels for point prompts,
          with shape BxN.
        'boxes': (torch.Tensor) Batched box inputs, with shape Bx4.
          Already transformed to the input frame of the model.
        'mask_inputs': (torch.Tensor) Batched mask inputs to the model,
          in the form Bx1xHxW.
    multimask_output (bool): Whether the model should predict multiple
      disambiguating masks, or return a single mask.
  
  Returns:
    (list(dict)): A list over input images, where each element is
      as dictionary with the following keys.
        'masks': (torch.Tensor) Batched binary mask predictions,
          with shape BxCxHxW, where B is the number of input prompts,
          C is determined by multimask_output, and (H, W) is the
          original size of the image.
        'iou_predictions': (torch.Tensor) The model's predictions
          of mask quality, in shape BxC.
        'low_res_logits': (torch.Tensor) Low resolution logits with
          shape BxCxHxW, where H=W=256. Can be passed as mask input
          to subsequent iterations of prediction.
docstr_tok: ["Predicts","masks","end-to-end","from","provided","images","and","prompts",".","If","prompts","are","not","known","in","advance",",","using","SamPredictor","is","recommended","over","calling","the","model","directly",".","Arguments",":","batched_input","(","list","(","dict",")",")",":","A","list","over","input","images",",","each","a","dictionary","with","the","following","keys",".","A","prompt","key","can","be","excluded","if","it","is","not","present",".","'image","'",":","The","image","as","a","torch","tensor","in","3xHxW","format",",","already","transformed","for","input","to","the","model",".","'original_size","'",":","(","tuple","(","int",",","int",")",")","The","original","size","of","the","image","before","transformation",",","as","(","H",",","W",")",".","'point_coords","'",":","(","torch.Tensor",")","Batched","point","prompts","for","this","image",",","with","shape","BxNx2",".","Already","transformed","to","the","input","frame","of","the","model",".","'point_labels","'",":","(","torch.Tensor",")","Batched","labels","for","point","prompts",",","with","shape","BxN",".","'boxes","'",":","(","torch.Tensor",")","Batched","box","inputs",",","with","shape","Bx4",".","Already","transformed","to","the","input","frame","of","the","model",".","'mask_inputs","'",":","(","torch.Tensor",")","Batched","mask","inputs","to","the","model",",","in","the","form","Bx1xHxW",".","multimask_output","(","bool",")",":","Whether","the","model","should","predict","multiple","disambiguating","masks",",","or","return","a","single","mask",".","Returns",":","(","list","(","dict",")",")",":","A","list","over","input","images",",","where","each","element","is","as","dictionary","with","the","following","keys",".","'masks","'",":","(","torch.Tensor",")","Batched","binary","mask","predictions",",","with","shape","BxCxHxW",",","where","B","is","the","number","of","input","prompts",",","C","is","determined","by","multimask_output",",","and","(","H",",","W",")","is","the","original","size","of","the","image",".","'iou_predictions","'",":","(","torch.Tensor",")","The","model","'s","predictions","of","mask","quality",",","in","shape","BxC",".","'low_res_logits","'",":","(","torch.Tensor",")","Low","resolution","logits","with","shape","BxCxHxW",",","where","H=W=256",".","Can","be","passed","as","mask","input","to","subsequent","iterations","of","prediction","."]
code: def forward(
      self,
      batched_input: List[Dict[str, Any]],
      multimask_output: bool,
  ) -> List[Dict[str, torch.Tensor]]:
      """
      Predicts masks end-to-end from provided images and prompts.
      If prompts are not known in advance, using SamPredictor is
      recommended over calling the model directly.
  
      Arguments:
        batched_input (list(dict)): A list over input images, each a
          dictionary with the following keys. A prompt key can be
          excluded if it is not present.
            'image': The image as a torch tensor in 3xHxW format,
              already transformed for input to the model.
            'original_size': (tuple(int, int)) The original size of
              the image before transformation, as (H, W).
            'point_coords': (torch.Tensor) Batched point prompts for
              this image, with shape BxNx2. Already transformed to the
              input frame of the model.
            'point_labels': (torch.Tensor) Batched labels for point prompts,
              with shape BxN.
            'boxes': (torch.Tensor) Batched box inputs, with shape Bx4.
              Already transformed to the input frame of the model.
            'mask_inputs': (torch.Tensor) Batched mask inputs to the model,
              in the form Bx1xHxW.
        multimask_output (bool): Whether the model should predict multiple
          disambiguating masks, or return a single mask.
  
      Returns:
        (list(dict)): A list over input images, where each element is
          as dictionary with the following keys.
            'masks': (torch.Tensor) Batched binary mask predictions,
              with shape BxCxHxW, where B is the number of input prompts,
              C is determined by multimask_output, and (H, W) is the
              original size of the image.
            'iou_predictions': (torch.Tensor) The model's predictions
              of mask quality, in shape BxC.
            'low_res_logits': (torch.Tensor) Low resolution logits with
              shape BxCxHxW, where H=W=256. Can be passed as mask input
              to subsequent iterations of prediction.
      """
      input_images = torch.stack(
          [self.preprocess(x["image"]) for x in batched_input], dim=0
      )
      image_embeddings = self.image_encoder(input_images)
  
      outputs = []
      for image_record, curr_embedding in zip(
          batched_input, image_embeddings
      ):
          if "point_coords" in image_record:
              points = (
                  image_record["point_coords"],
                  image_record["point_labels"],
              )
          else:
              points = None
          sparse_embeddings, dense_embeddings = self.prompt_encoder(
              points=points,
              boxes=image_record.get("boxes", None),
              masks=image_record.get("mask_inputs", None),
          )
          low_res_masks, iou_predictions = self.mask_decoder(
              image_embeddings=curr_embedding.unsqueeze(0),
              image_pe=self.prompt_encoder.get_dense_pe(),
              sparse_prompt_embeddings=sparse_embeddings,
              dense_prompt_embeddings=dense_embeddings,
              multimask_output=multimask_output,
          )
          masks = self.postprocess_masks(
              low_res_masks,
              input_size=image_record["image"].shape[-2:],
              original_size=image_record["original_size"],
          )
          masks = masks > self.mask_threshold
          outputs.append(
              {
                  "masks": masks,
                  "iou_predictions": iou_predictions,
                  "low_res_logits": low_res_masks,
              }
          )
      return outputs
  
file_path: segment-anything/segment_anything/modeling/sam.py
imports: import torch
  from torch import nn
  from torch.nn import functional
  from typing import Any, Dict, List, Tuple
  from .image_encoder import ImageEncoderViT
  from .mask_decoder import MaskDecoder
  from .prompt_encoder import PromptEncoder
label_desc: Use image node_id 3 for calling the Sam obj's underlying member method code with example usage: obj.forward(batched_input, multimask_output) and returns: outputs