micro_sam.bioimageio.predictor_adaptor

  1import warnings
  2from typing import Optional, Tuple
  3
  4import numpy as np
  5
  6import torch
  7from torch import nn
  8
  9from segment_anything.predictor import SamPredictor
 10
 11try:
 12    # Avoid import warnings from mobile_sam
 13    with warnings.catch_warnings():
 14        warnings.simplefilter("ignore")
 15        from mobile_sam import sam_model_registry
 16except ImportError:
 17    from segment_anything import sam_model_registry
 18
 19
 20class PredictorAdaptor(nn.Module):
 21    """Wrapper around the SamPredictor.
 22
 23    This model supports the same functionality as SamPredictor and can provide mask segmentations
 24    from box, point or mask input prompts.
 25
 26    If it was loaded from a checkpoint that also contains the state of an instance segmentation
 27    decoder, then calling it without any prompts will run automatic instance segmentation (AIS):
 28    the UNETR decoder predicts foreground and distance maps from the image embeddings and the
 29    instances are computed from them via a seeded watershed.
 30    Running AIS requires the `micro_sam` package.
 31
 32    Args:
 33        model_type: The type of the model for the image encoder.
 34            Can be one of 'vit_b', 'vit_l', 'vit_h' or 'vit_t'.
 35            For 'vit_t' support the 'mobile_sam' package has to be installed.
 36    """
 37    def __init__(self, model_type: str) -> None:
 38        super().__init__()
 39        self.sam_model = sam_model_registry[model_type]()
 40        self.sam = SamPredictor(self.sam_model)
 41        self.decoder = None
 42        # Cache the bounded SAM input and original size for invalidation.
 43        self._cached_input = None
 44        self._cached_original_size = None
 45
 46    def load_state_dict(self, state, **kwargs):
 47        # Finetuning checkpoints store SAM and decoder weights separately.
 48        if "model_state" in state:
 49            load_result = self.sam.model.load_state_dict(state["model_state"], **kwargs)
 50            decoder_state = state.get("decoder_state")
 51            if decoder_state is not None:
 52                from micro_sam.instance_segmentation import get_decoder
 53                device = next(self.sam.model.parameters()).device
 54                self.decoder = get_decoder(self.sam.model.image_encoder, decoder_state, device=device)
 55            return load_result
 56
 57        return self.sam.model.load_state_dict(state, **kwargs)
 58
 59    def _automatic_instance_segmentation(self, image: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
 60        """Run automatic instance segmentation with the decoder for the image embeddings
 61        that were set in the forward method.
 62
 63        Returns the instances as a stack of binary masks, so that the output signature
 64        matches the output of the prompt-based segmentation.
 65        """
 66        if self.decoder is None:
 67            raise ValueError(
 68                "This model was exported without an instance segmentation decoder, "
 69                "so it does not support automatic instance segmentation. "
 70                "At least one prompt input (box, point or mask) is required."
 71            )
 72        from micro_sam.instance_segmentation import InstanceSegmentationWithDecoder
 73
 74        segmenter = InstanceSegmentationWithDecoder(self.sam, self.decoder)
 75        image_embeddings = {
 76            "features": self.sam.features,
 77            "input_size": tuple(self.sam.input_size),
 78            "original_size": tuple(self.sam.original_size),
 79        }
 80        # The image is unused because the embeddings are precomputed.
 81        segmenter.initialize(image=image[0].permute(1, 2, 0).cpu().numpy(), image_embeddings=image_embeddings)
 82        segmentation = segmenter.generate(output_mode="instance_segmentation")
 83        seg_ids = np.unique(segmentation)
 84        seg_ids = seg_ids[seg_ids != 0]
 85        instance_masks = [segmentation == seg_id for seg_id in seg_ids]
 86
 87        height, width = self.sam.original_size
 88        if len(instance_masks) == 0:
 89            masks = torch.zeros((1, 0, 1, height, width), dtype=torch.uint8)
 90        else:
 91            masks = torch.from_numpy(np.stack(instance_masks)[None, :, None].astype("uint8"))
 92        # AIS does not predict mask quality.
 93        scores = torch.ones((1, masks.shape[1], 1), dtype=torch.float32)
 94        return masks, scores
 95
 96    @torch.no_grad()
 97    def forward(
 98        self,
 99        image: torch.Tensor,
100        box_prompts: Optional[torch.Tensor] = None,
101        point_prompts: Optional[torch.Tensor] = None,
102        point_labels: Optional[torch.Tensor] = None,
103        mask_prompts: Optional[torch.Tensor] = None,
104        embeddings: Optional[torch.Tensor] = None,
105    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
106        """
107
108        Args:
109            image: torch inputs of dimensions B x C x H x W
110            box_prompts: box coordinates of dimensions B x OBJECTS x 4
111            point_prompts: point coordinates of dimension B x OBJECTS x POINTS x 2
112            point_labels: point labels of dimension B x OBJECTS x POINTS
113            mask_prompts: mask prompts of dimension B x OBJECTS x 256 x 256
114            embeddings: precomputed image embeddings B x 256 x 64 x 64
115
116        Returns:
117            The segmentation masks.
118            The scores for prediction quality.
119            The computed image embeddings.
120        """
121        batch_size = image.shape[0]
122        if batch_size != 1:
123            raise ValueError
124
125        # Cast to float for MPS compatibility: F.interpolate with antialias=True
126        # only supports floating-point dtypes on MPS (Apple Silicon).
127        image_float = image.float() if not image.is_floating_point() else image
128        input_ = self.sam.transform.apply_image_torch(image_float)
129        original_image_size = tuple(image.shape[2:])
130
131        # Reuse embeddings only when their input and output geometry match.
132        if (
133            self.sam.is_image_set
134            and embeddings is None
135            and self._cached_input is not None
136            and self._cached_original_size == original_image_size
137            and input_.shape == self._cached_input.shape
138            and torch.equal(input_, self._cached_input)
139        ):
140            pass  # do nothing
141
142        # The embeddings are passed, so we set them.
143        elif embeddings is not None:
144            self.sam.features = embeddings
145            self.sam.orig_h, self.sam.orig_w = original_image_size
146            self.sam.input_h, self.sam.input_w = input_.shape[2:]
147            self.sam.is_image_set = True
148            self._cached_input = input_.detach().clone()
149            self._cached_original_size = original_image_size
150
151        # No embeddings were passed and we don't have embeddings for this image,
152        # so we compute them.
153        else:
154            self.sam.set_torch_image(input_, original_image_size=original_image_size)
155            self.sam.orig_h, self.sam.orig_w = self.sam.original_size
156            self.sam.input_h, self.sam.input_w = self.sam.input_size
157            self._cached_input = input_.detach().clone()
158            self._cached_original_size = original_image_size
159
160        assert self.sam.is_image_set, "The predictor has not yet been initialized."
161
162        # Ensure input size and original size are set.
163        self.sam.input_size = (self.sam.input_h, self.sam.input_w)
164        self.sam.original_size = (self.sam.orig_h, self.sam.orig_w)
165
166        # Preserve prompt-free SamPredictor inference without a decoder.
167        prompts = (box_prompts, point_prompts, mask_prompts)
168        if self.decoder is not None and all(prompt is None for prompt in prompts):
169            masks, scores = self._automatic_instance_segmentation(image)
170            embeddings = self.sam.get_image_embedding()
171            return masks, scores, embeddings
172
173        if box_prompts is None:
174            boxes = None
175        else:
176            boxes = self.sam.transform.apply_boxes_torch(box_prompts, original_size=self.sam.original_size)
177
178        if point_prompts is None:
179            point_coords = None
180        else:
181            assert point_labels is not None
182            point_coords = self.sam.transform.apply_coords_torch(point_prompts, original_size=self.sam.original_size)[0]
183            point_labels = point_labels[0]
184
185        if mask_prompts is None:
186            mask_input = None
187        else:
188            mask_input = mask_prompts[0]
189
190        masks, scores, _ = self.sam.predict_torch(
191            point_coords=point_coords,
192            point_labels=point_labels,
193            boxes=boxes,
194            mask_input=mask_input,
195            multimask_output=False
196        )
197
198        assert masks.shape[2:] == image.shape[2:], \
199            f"{masks.shape[2:]} is not as expected ({image.shape[2:]})"
200
201        # Ensure batch axis.
202        if masks.ndim == 4:
203            masks = masks[None]
204            assert scores.ndim == 2
205            scores = scores[None]
206
207        embeddings = self.sam.get_image_embedding()
208        return masks.to(dtype=torch.uint8), scores, embeddings