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):
(
self,
image_embeddings: torch.Tensor,
image_pe: torch.Tensor,
sparse_prompt_embeddings: torch.Tensor,
dense_prompt_embeddings: torch.Tensor,
multimask_output: bool,
repeat_image: bool,
high_res_features: Optional[List[torch.Tensor]] = None,
)
| 108 | self.dynamic_multimask_stability_thresh = dynamic_multimask_stability_thresh |
| 109 | |
| 110 | def forward( |
| 111 | self, |
| 112 | image_embeddings: torch.Tensor, |
| 113 | image_pe: torch.Tensor, |
| 114 | sparse_prompt_embeddings: torch.Tensor, |
| 115 | dense_prompt_embeddings: torch.Tensor, |
| 116 | multimask_output: bool, |
| 117 | repeat_image: bool, |
| 118 | high_res_features: Optional[List[torch.Tensor]] = None, |
| 119 | ) -> Tuple[torch.Tensor, torch.Tensor]: |
| 120 | """ |
| 121 | Predict masks given image and prompt embeddings. |
| 122 | |
| 123 | Arguments: |
| 124 | image_embeddings (torch.Tensor): the embeddings from the image encoder |
| 125 | image_pe (torch.Tensor): positional encoding with the shape of image_embeddings |
| 126 | sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes |
| 127 | dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs |
| 128 | multimask_output (bool): Whether to return multiple masks or a single |
| 129 | mask. |
| 130 | |
| 131 | Returns: |
| 132 | torch.Tensor: batched predicted masks |
| 133 | torch.Tensor: batched predictions of mask quality |
| 134 | torch.Tensor: batched SAM token for mask output |
| 135 | """ |
| 136 | masks, iou_pred, mask_tokens_out, object_score_logits = self.predict_masks( |
| 137 | image_embeddings=image_embeddings, |
| 138 | image_pe=image_pe, |
| 139 | sparse_prompt_embeddings=sparse_prompt_embeddings, |
| 140 | dense_prompt_embeddings=dense_prompt_embeddings, |
| 141 | repeat_image=repeat_image, |
| 142 | high_res_features=high_res_features, |
| 143 | ) |
| 144 | |
| 145 | # Select the correct mask or masks for output |
| 146 | if multimask_output: |
| 147 | masks = masks[:, 1:, :, :] |
| 148 | iou_pred = iou_pred[:, 1:] |
| 149 | elif self.dynamic_multimask_via_stability and not self.training: |
| 150 | masks, iou_pred = self._dynamic_multimask_via_stability(masks, iou_pred) |
| 151 | else: |
| 152 | masks = masks[:, 0:1, :, :] |
| 153 | iou_pred = iou_pred[:, 0:1] |
| 154 | |
| 155 | if multimask_output and self.use_multimask_token_for_obj_ptr: |
| 156 | sam_tokens_out = mask_tokens_out[:, 1:] # [b, 3, c] shape |
| 157 | else: |
| 158 | # Take the mask output token. Here we *always* use the token for single mask output. |
| 159 | # At test time, even if we track after 1-click (and using multimask_output=True), |
| 160 | # we still take the single mask token here. The rationale is that we always track |
| 161 | # after multiple clicks during training, so the past tokens seen during training |
| 162 | # are always the single mask token (and we'll let it be the object-memory token). |
| 163 | sam_tokens_out = mask_tokens_out[:, 0:1] # [b, 1, c] shape |
| 164 | |
| 165 | # Prepare output |
| 166 | return masks, iou_pred, sam_tokens_out, object_score_logits |
| 167 |
nothing calls this directly
no test coverage detected