micro_sam.bioimageio.predictor_adaptor

  1import warnings
  2from typing import Optional, Tuple
  3
  4import torch
  5from torch import nn
  6
  7from segment_anything.predictor import SamPredictor
  8
  9try:
 10    # Avoid import warnings from mobile_sam
 11    with warnings.catch_warnings():
 12        warnings.simplefilter("ignore")
 13        from mobile_sam import sam_model_registry
 14except ImportError:
 15    from segment_anything import sam_model_registry
 16
 17
 18class PredictorAdaptor(nn.Module):
 19    """Wrapper around the SamPredictor.
 20
 21    This model supports the same functionality as SamPredictor and can provide mask segmentations
 22    from box, point or mask input prompts.
 23
 24    Args:
 25        model_type: The type of the model for the image encoder.
 26            Can be one of 'vit_b', 'vit_l', 'vit_h' or 'vit_t'.
 27            For 'vit_t' support the 'mobile_sam' package has to be installed.
 28    """
 29    def __init__(self, model_type: str) -> None:
 30        super().__init__()
 31        sam_model = sam_model_registry[model_type]()
 32        self.sam = SamPredictor(sam_model)
 33
 34    def load_state_dict(self, state):
 35        self.sam.model.load_state_dict(state)
 36
 37    @torch.no_grad()
 38    def forward(
 39        self,
 40        image: torch.Tensor,
 41        box_prompts: Optional[torch.Tensor] = None,
 42        point_prompts: Optional[torch.Tensor] = None,
 43        point_labels: Optional[torch.Tensor] = None,
 44        mask_prompts: Optional[torch.Tensor] = None,
 45        embeddings: Optional[torch.Tensor] = None,
 46    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
 47        """
 48
 49        Args:
 50            image: torch inputs of dimensions B x C x H x W
 51            box_prompts: box coordinates of dimensions B x OBJECTS x 4
 52            point_prompts: point coordinates of dimension B x OBJECTS x POINTS x 2
 53            point_labels: point labels of dimension B x OBJECTS x POINTS
 54            mask_prompts: mask prompts of dimension B x OBJECTS x 256 x 256
 55            embeddings: precomputed image embeddings B x 256 x 64 x 64
 56
 57        Returns:
 58        """
 59        batch_size = image.shape[0]
 60        if batch_size != 1:
 61            raise ValueError
 62
 63        # We have image embeddings set and image embeddings were not passed.
 64        if self.sam.is_image_set and embeddings is None:
 65            pass   # do nothing
 66
 67        # The embeddings are passed, so we set them.
 68        elif embeddings is not None:
 69            self.sam.features = embeddings
 70            self.sam.orig_h, self.sam.orig_w = image.shape[2:]
 71            self.sam.input_h, self.sam.input_w = self.sam.transform.apply_image_torch(image).shape[2:]
 72            self.sam.is_image_set = True
 73
 74        # We don't have image embeddings set and they were not passed.
 75        elif not self.sam.is_image_set:
 76            input_ = self.sam.transform.apply_image_torch(image)
 77            self.sam.set_torch_image(input_, original_image_size=image.shape[2:])
 78            self.sam.orig_h, self.sam.orig_w = self.sam.original_size
 79            self.sam.input_h, self.sam.input_w = self.sam.input_size
 80
 81        assert self.sam.is_image_set, "The predictor has not yet been initialized."
 82
 83        # Ensure input size and original size are set.
 84        self.sam.input_size = (self.sam.input_h, self.sam.input_w)
 85        self.sam.original_size = (self.sam.orig_h, self.sam.orig_w)
 86
 87        if box_prompts is None:
 88            boxes = None
 89        else:
 90            boxes = self.sam.transform.apply_boxes_torch(box_prompts, original_size=self.sam.original_size)
 91
 92        if point_prompts is None:
 93            point_coords = None
 94        else:
 95            assert point_labels is not None
 96            point_coords = self.sam.transform.apply_coords_torch(point_prompts, original_size=self.sam.original_size)[0]
 97            point_labels = point_labels[0]
 98
 99        if mask_prompts is None:
100            mask_input = None
101        else:
102            mask_input = mask_prompts[0]
103
104        masks, scores, _ = self.sam.predict_torch(
105            point_coords=point_coords,
106            point_labels=point_labels,
107            boxes=boxes,
108            mask_input=mask_input,
109            multimask_output=False
110        )
111
112        assert masks.shape[2:] == image.shape[2:], \
113            f"{masks.shape[2:]} is not as expected ({image.shape[2:]})"
114
115        # Ensure batch axis.
116        if masks.ndim == 4:
117            masks = masks[None]
118            assert scores.ndim == 2
119            scores = scores[None]
120
121        embeddings = self.sam.get_image_embedding()
122        return masks.to(dtype=torch.uint8), scores, embeddings