| import torch |
| import torch.nn as nn |
| from ..modeling import Sam |
| from .amg import calculate_stability_score |
|
|
|
|
| class SamCoreMLModel(nn.Module): |
| """ |
| This model should not be called directly, but is used in CoreML export. |
| """ |
|
|
| def __init__( |
| self, |
| model: Sam, |
| use_stability_score: bool = False |
| ) -> None: |
| super().__init__() |
| self.mask_decoder = model.mask_decoder |
| self.model = model |
| self.img_size = model.image_encoder.img_size |
| self.use_stability_score = use_stability_score |
| self.stability_score_offset = 1.0 |
|
|
| def _embed_points(self, point_coords: torch.Tensor, point_labels: torch.Tensor) -> torch.Tensor: |
| point_coords = point_coords + 0.5 |
| point_coords = point_coords / self.img_size |
| point_embedding = self.model.prompt_encoder.pe_layer._pe_encoding(point_coords) |
| point_labels = point_labels.unsqueeze(-1).expand_as(point_embedding) |
|
|
| point_embedding = point_embedding * (point_labels != -1) |
| point_embedding = point_embedding + self.model.prompt_encoder.not_a_point_embed.weight * ( |
| point_labels == -1 |
| ) |
|
|
| for i in range(self.model.prompt_encoder.num_point_embeddings): |
| point_embedding = point_embedding + self.model.prompt_encoder.point_embeddings[ |
| i |
| ].weight * (point_labels == i) |
|
|
| return point_embedding |
|
|
| @torch.no_grad() |
| def forward( |
| self, |
| image_embeddings: torch.Tensor, |
| point_coords: torch.Tensor, |
| point_labels: torch.Tensor, |
| ): |
| sparse_embedding = self._embed_points(point_coords, point_labels) |
| dense_embedding = self.model.prompt_encoder.no_mask_embed.weight.reshape(1, -1, 1, 1) |
|
|
| masks, scores = self.model.mask_decoder.predict_masks( |
| image_embeddings=image_embeddings, |
| image_pe=self.model.prompt_encoder.get_dense_pe(), |
| sparse_prompt_embeddings=sparse_embedding, |
| dense_prompt_embeddings=dense_embedding, |
| ) |
|
|
| if self.use_stability_score: |
| scores = calculate_stability_score( |
| masks, self.model.mask_threshold, self.stability_score_offset |
| ) |
|
|
| return scores, masks |
|
|
