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,
)
| 69 | ) |
| 70 | |
| 71 | def forward( |
| 72 | self, |
| 73 | image_embeddings: torch.Tensor, |
| 74 | image_pe: torch.Tensor, |
| 75 | sparse_prompt_embeddings: torch.Tensor, |
| 76 | dense_prompt_embeddings: torch.Tensor, |
| 77 | multimask_output: bool, |
| 78 | ) -> Tuple[torch.Tensor, torch.Tensor]: |
| 79 | """ |
| 80 | Predict masks given image and prompt embeddings. |
| 81 | |
| 82 | Arguments: |
| 83 | image_embeddings (torch.Tensor): the embeddings from the image encoder |
| 84 | image_pe (torch.Tensor): positional encoding with the shape of image_embeddings |
| 85 | sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes |
| 86 | dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs |
| 87 | multimask_output (bool): Whether to return multiple masks or a single |
| 88 | mask. |
| 89 | |
| 90 | Returns: |
| 91 | torch.Tensor: batched predicted masks |
| 92 | torch.Tensor: batched predictions of mask quality |
| 93 | """ |
| 94 | masks, iou_pred = self.predict_masks( |
| 95 | image_embeddings=image_embeddings, |
| 96 | image_pe=image_pe, |
| 97 | sparse_prompt_embeddings=sparse_prompt_embeddings, |
| 98 | dense_prompt_embeddings=dense_prompt_embeddings, |
| 99 | ) |
| 100 | |
| 101 | # Select the correct mask or masks for outptu |
| 102 | if multimask_output: |
| 103 | mask_slice = slice(1, None) |
| 104 | else: |
| 105 | mask_slice = slice(0, 1) |
| 106 | masks = masks[:, mask_slice, :, :] |
| 107 | iou_pred = iou_pred[:, mask_slice] |
| 108 | |
| 109 | # Prepare output |
| 110 | return masks, iou_pred |
| 111 | |
| 112 | def predict_masks( |
| 113 | self, |
nothing calls this directly
no test coverage detected