Spaces:
Running
Running
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # All rights reserved. | |
| # This source code is licensed under the license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| import torch | |
| import torch.nn as nn | |
| from torch.nn import functional as F | |
| from typing import Tuple | |
| from ..modeling import Sam | |
| from .amg import calculate_stability_score | |
| class SamOnnxModel(nn.Module): | |
| """ | |
| This model should not be called directly, but is used in ONNX export. | |
| It combines the prompt encoder, mask decoder, and mask postprocessing of Sam, | |
| with some functions modified to enable model tracing. Also supports extra | |
| options controlling what information. See the ONNX export script for details. | |
| """ | |
| def __init__( | |
| self, | |
| model: Sam, | |
| return_single_mask: bool, | |
| use_stability_score: bool = False, | |
| return_extra_metrics: bool = False, | |
| ) -> None: | |
| super().__init__() | |
| self.mask_decoder = model.mask_decoder | |
| self.model = model | |
| self.img_size = model.image_encoder.img_size | |
| self.return_single_mask = return_single_mask | |
| self.use_stability_score = use_stability_score | |
| self.stability_score_offset = 1.0 | |
| self.return_extra_metrics = return_extra_metrics | |
| def resize_longest_image_size( | |
| input_image_size: torch.Tensor, longest_side: int | |
| ) -> torch.Tensor: | |
| input_image_size = input_image_size.to(torch.float32) | |
| scale = longest_side / torch.max(input_image_size) | |
| transformed_size = scale * input_image_size | |
| transformed_size = torch.floor(transformed_size + 0.5).to(torch.int64) | |
| return transformed_size | |
| 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 | |
| def _embed_masks(self, input_mask: torch.Tensor, has_mask_input: torch.Tensor) -> torch.Tensor: | |
| mask_embedding = has_mask_input * self.model.prompt_encoder.mask_downscaling(input_mask) | |
| mask_embedding = mask_embedding + ( | |
| 1 - has_mask_input | |
| ) * self.model.prompt_encoder.no_mask_embed.weight.reshape(1, -1, 1, 1) | |
| return mask_embedding | |
| def mask_postprocessing(self, masks: torch.Tensor, orig_im_size: torch.Tensor) -> torch.Tensor: | |
| masks = F.interpolate( | |
| masks, | |
| size=(self.img_size, self.img_size), | |
| mode="bilinear", | |
| align_corners=False, | |
| ) | |
| prepadded_size = self.resize_longest_image_size(orig_im_size, self.img_size).to(torch.int64) | |
| masks = masks[..., : prepadded_size[0], : prepadded_size[1]] # type: ignore | |
| orig_im_size = orig_im_size.to(torch.int64) | |
| h, w = orig_im_size[0], orig_im_size[1] | |
| masks = F.interpolate(masks, size=(h, w), mode="bilinear", align_corners=False) | |
| return masks | |
| def select_masks( | |
| self, masks: torch.Tensor, iou_preds: torch.Tensor, num_points: int | |
| ) -> Tuple[torch.Tensor, torch.Tensor]: | |
| # Determine if we should return the multiclick mask or not from the number of points. | |
| # The reweighting is used to avoid control flow. | |
| score_reweight = torch.tensor( | |
| [[1000] + [0] * (self.model.mask_decoder.num_mask_tokens - 1)] | |
| ).to(iou_preds.device) | |
| score = iou_preds + (num_points - 2.5) * score_reweight | |
| best_idx = torch.argmax(score, dim=1) | |
| masks = masks[torch.arange(masks.shape[0]), best_idx, :, :].unsqueeze(1) | |
| iou_preds = iou_preds[torch.arange(masks.shape[0]), best_idx].unsqueeze(1) | |
| return masks, iou_preds | |
| def forward( | |
| self, | |
| image_embeddings: torch.Tensor, | |
| point_coords: torch.Tensor, | |
| point_labels: torch.Tensor, | |
| mask_input: torch.Tensor, | |
| has_mask_input: torch.Tensor, | |
| orig_im_size: torch.Tensor, | |
| ): | |
| sparse_embedding = self._embed_points(point_coords, point_labels) | |
| dense_embedding = self._embed_masks(mask_input, has_mask_input) | |
| 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 | |
| ) | |
| if self.return_single_mask: | |
| masks, scores = self.select_masks(masks, scores, point_coords.shape[1]) | |
| upscaled_masks = self.mask_postprocessing(masks, orig_im_size) | |
| if self.return_extra_metrics: | |
| stability_scores = calculate_stability_score( | |
| upscaled_masks, self.model.mask_threshold, self.stability_score_offset | |
| ) | |
| areas = (upscaled_masks > self.model.mask_threshold).sum(-1).sum(-1) | |
| return upscaled_masks, scores, stability_scores, areas, masks | |
| return upscaled_masks, scores, masks | |