| 14 | |
| 15 | |
| 16 | class MaskDecoder(nn.Module): |
| 17 | def __init__( |
| 18 | self, |
| 19 | *, |
| 20 | transformer_dim: int, |
| 21 | transformer: nn.Module, |
| 22 | num_multimask_outputs: int = 3, |
| 23 | activation: Type[nn.Module] = nn.GELU, |
| 24 | iou_head_depth: int = 3, |
| 25 | iou_head_hidden_dim: int = 256, |
| 26 | ) -> None: |
| 27 | """ |
| 28 | Predicts masks given an image and prompt embeddings, using a |
| 29 | transformer architecture. |
| 30 | |
| 31 | Arguments: |
| 32 | transformer_dim (int): the channel dimension of the transformer |
| 33 | transformer (nn.Module): the transformer used to predict masks |
| 34 | num_multimask_outputs (int): the number of masks to predict |
| 35 | when disambiguating masks |
| 36 | activation (nn.Module): the type of activation to use when |
| 37 | upscaling masks |
| 38 | iou_head_depth (int): the depth of the MLP used to predict |
| 39 | mask quality |
| 40 | iou_head_hidden_dim (int): the hidden dimension of the MLP |
| 41 | used to predict mask quality |
| 42 | """ |
| 43 | super().__init__() |
| 44 | self.transformer_dim = transformer_dim |
| 45 | self.transformer = transformer |
| 46 | |
| 47 | self.num_multimask_outputs = num_multimask_outputs |
| 48 | |
| 49 | self.iou_token = nn.Embedding(1, transformer_dim) |
| 50 | self.num_mask_tokens = num_multimask_outputs + 1 |
| 51 | self.mask_tokens = nn.Embedding(self.num_mask_tokens, transformer_dim) |
| 52 | |
| 53 | self.output_upscaling = nn.Sequential( |
| 54 | nn.ConvTranspose2d(transformer_dim, transformer_dim // 4, kernel_size=2, stride=2), |
| 55 | LayerNorm2d(transformer_dim // 4), |
| 56 | activation(), |
| 57 | nn.ConvTranspose2d(transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2), |
| 58 | activation(), |
| 59 | ) |
| 60 | self.output_hypernetworks_mlps = nn.ModuleList( |
| 61 | [ |
| 62 | MLP(transformer_dim, transformer_dim, transformer_dim // 8, 3) |
| 63 | for i in range(self.num_mask_tokens) |
| 64 | ] |
| 65 | ) |
| 66 | |
| 67 | self.iou_prediction_head = MLP( |
| 68 | transformer_dim, iou_head_hidden_dim, self.num_mask_tokens, iou_head_depth |
| 69 | ) |
| 70 | |
| 71 | def forward( |
| 72 | self, |
| 73 | image_embeddings: torch.Tensor, |