| 15 | |
| 16 | |
| 17 | class PromptEncoder(nn.Module): |
| 18 | def __init__( |
| 19 | self, |
| 20 | embed_dim: int, |
| 21 | image_embedding_size: Tuple[int, int], |
| 22 | input_image_size: Tuple[int, int], |
| 23 | mask_in_chans: int, |
| 24 | activation: Type[nn.Module] = nn.GELU, |
| 25 | ) -> None: |
| 26 | """ |
| 27 | Encodes prompts for input to SAM's mask decoder. |
| 28 | |
| 29 | Arguments: |
| 30 | embed_dim (int): The prompts' embedding dimension |
| 31 | image_embedding_size (tuple(int, int)): The spatial size of the |
| 32 | image embedding, as (H, W). |
| 33 | input_image_size (int): The padded size of the image as input |
| 34 | to the image encoder, as (H, W). |
| 35 | mask_in_chans (int): The number of hidden channels used for |
| 36 | encoding input masks. |
| 37 | activation (nn.Module): The activation to use when encoding |
| 38 | input masks. |
| 39 | """ |
| 40 | super().__init__() |
| 41 | self.embed_dim = embed_dim |
| 42 | self.input_image_size = input_image_size |
| 43 | self.image_embedding_size = image_embedding_size |
| 44 | self.pe_layer = PositionEmbeddingRandom(embed_dim // 2) |
| 45 | |
| 46 | self.num_point_embeddings: int = 4 # pos/neg point + 2 box corners |
| 47 | point_embeddings = [ |
| 48 | nn.Embedding(1, embed_dim) for i in range(self.num_point_embeddings) |
| 49 | ] |
| 50 | self.point_embeddings = nn.ModuleList(point_embeddings) |
| 51 | self.not_a_point_embed = nn.Embedding(1, embed_dim) |
| 52 | |
| 53 | self.mask_input_size = ( |
| 54 | 4 * image_embedding_size[0], |
| 55 | 4 * image_embedding_size[1], |
| 56 | ) |
| 57 | self.mask_downscaling = nn.Sequential( |
| 58 | nn.Conv2d(1, mask_in_chans // 4, kernel_size=2, stride=2), |
| 59 | LayerNorm2d(mask_in_chans // 4), |
| 60 | activation(), |
| 61 | nn.Conv2d(mask_in_chans // 4, mask_in_chans, kernel_size=2, stride=2), |
| 62 | LayerNorm2d(mask_in_chans), |
| 63 | activation(), |
| 64 | nn.Conv2d(mask_in_chans, embed_dim, kernel_size=1), |
| 65 | ) |
| 66 | self.no_mask_embed = nn.Embedding(1, embed_dim) |
| 67 | |
| 68 | def get_dense_pe(self) -> torch.Tensor: |
| 69 | """ |
| 70 | Returns the positional encoding used to encode point prompts, |
| 71 | applied to a dense set of points the shape of the image encoding. |
| 72 | |
| 73 | Returns: |
| 74 | torch.Tensor: Positional encoding with shape |