| 13 | |
| 14 | |
| 15 | class MaskDecoder(nn.Module): |
| 16 | |
| 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 | use_high_res_features: bool = False, |
| 27 | iou_prediction_use_sigmoid=False, |
| 28 | dynamic_multimask_via_stability=False, |
| 29 | dynamic_multimask_stability_delta=0.05, |
| 30 | dynamic_multimask_stability_thresh=0.98, |
| 31 | pred_obj_scores: bool = False, |
| 32 | pred_obj_scores_mlp: bool = False, |
| 33 | use_multimask_token_for_obj_ptr: bool = False, |
| 34 | ) -> None: |
| 35 | """ |
| 36 | Predicts masks given an image and prompt embeddings, using a |
| 37 | transformer architecture. |
| 38 | |
| 39 | Arguments: |
| 40 | transformer_dim (int): the channel dimension of the transformer |
| 41 | transformer (nn.Module): the transformer used to predict masks |
| 42 | num_multimask_outputs (int): the number of masks to predict |
| 43 | when disambiguating masks |
| 44 | activation (nn.Module): the type of activation to use when |
| 45 | upscaling masks |
| 46 | iou_head_depth (int): the depth of the MLP used to predict |
| 47 | mask quality |
| 48 | iou_head_hidden_dim (int): the hidden dimension of the MLP |
| 49 | used to predict mask quality |
| 50 | """ |
| 51 | super().__init__() |
| 52 | self.transformer_dim = transformer_dim |
| 53 | self.transformer = transformer |
| 54 | |
| 55 | self.num_multimask_outputs = num_multimask_outputs |
| 56 | |
| 57 | self.iou_token = nn.Embedding(1, transformer_dim) |
| 58 | self.num_mask_tokens = num_multimask_outputs + 1 |
| 59 | self.mask_tokens = nn.Embedding(self.num_mask_tokens, transformer_dim) |
| 60 | |
| 61 | self.pred_obj_scores = pred_obj_scores |
| 62 | if self.pred_obj_scores: |
| 63 | self.obj_score_token = nn.Embedding(1, transformer_dim) |
| 64 | self.use_multimask_token_for_obj_ptr = use_multimask_token_for_obj_ptr |
| 65 | |
| 66 | self.output_upscaling = nn.Sequential( |
| 67 | nn.ConvTranspose2d(transformer_dim, transformer_dim // 4, kernel_size=2, stride=2), |
| 68 | LayerNorm2d(transformer_dim // 4), |
| 69 | activation(), |
| 70 | nn.ConvTranspose2d(transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2), |
| 71 | activation(), |
| 72 | ) |