micro_sam.instance_segmentation

   1"""
   2Automated instance segmentation functionality.
   3The classes implemented here extend the automatic instance segmentation from Segment Anything:
   4https://computational-cell-analytics.github.io/micro-sam/micro_sam.html
   5"""
   6
   7import os
   8from abc import ABC
   9from copy import deepcopy
  10from collections import OrderedDict
  11from typing import Any, Dict, List, Optional, Tuple, Union
  12
  13import vigra
  14import numpy as np
  15import torch
  16
  17from skimage.measure import label, regionprops
  18from skimage.segmentation import relabel_sequential
  19from torchvision.ops.boxes import batched_nms, box_area
  20
  21from torch_em.model import UNETR
  22from torch_em.util.segmentation import watershed_from_center_and_boundary_distances
  23
  24from nifty.tools import blocking
  25
  26import segment_anything.utils.amg as amg_utils
  27from segment_anything.predictor import SamPredictor
  28
  29from . import util
  30from ._vendored import batched_mask_to_box, mask_to_rle_pytorch
  31
  32#
  33# Utility Functionality
  34#
  35
  36
  37class _FakeInput:
  38    def __init__(self, shape):
  39        self.shape = shape
  40
  41    def __getitem__(self, index):
  42        block_shape = tuple(ind.stop - ind.start for ind in index)
  43        return np.zeros(block_shape, dtype="float32")
  44
  45
  46def mask_data_to_segmentation(
  47    masks: List[Dict[str, Any]],
  48    with_background: bool,
  49    min_object_size: int = 0,
  50    max_object_size: Optional[int] = None,
  51    label_masks: bool = True,
  52) -> np.ndarray:
  53    """Convert the output of the automatic mask generation to an instance segmentation.
  54
  55    Args:
  56        masks: The outputs generated by AutomaticMaskGenerator or EmbeddingMaskGenerator.
  57            Only supports output_mode=binary_mask.
  58        with_background: Whether the segmentation has background. If yes this function assures that the largest
  59            object in the output will be mapped to zero (the background value).
  60        min_object_size: The minimal size of an object in pixels.
  61        max_object_size: The maximal size of an object in pixels.
  62        label_masks: Whether to apply connected components to the result before removing small objects.
  63
  64    Returns:
  65        The instance segmentation.
  66    """
  67
  68    masks = sorted(masks, key=(lambda x: x["area"]), reverse=True)
  69    # we could also get the shape from the crop box
  70    shape = next(iter(masks))["segmentation"].shape
  71    segmentation = np.zeros(shape, dtype="uint32")
  72
  73    def require_numpy(mask):
  74        return mask.cpu().numpy() if torch.is_tensor(mask) else mask
  75
  76    seg_id = 1
  77    for mask in masks:
  78        if mask["area"] < min_object_size:
  79            continue
  80        if max_object_size is not None and mask["area"] > max_object_size:
  81            continue
  82
  83        this_seg_id = mask.get("seg_id", seg_id)
  84        segmentation[require_numpy(mask["segmentation"])] = this_seg_id
  85        seg_id = this_seg_id + 1
  86
  87    if label_masks:
  88        segmentation = label(segmentation).astype(segmentation.dtype)
  89
  90    seg_ids, sizes = np.unique(segmentation, return_counts=True)
  91
  92    # In some cases objects may be smaller than peviously calculated,
  93    # since they are covered by other objects. We ensure these also get
  94    # filtered out here.
  95    filter_ids = seg_ids[sizes < min_object_size]
  96
  97    # If we run segmentation with background we also map the largest segment
  98    # (the most likely background object) to zero. This is often zero already,
  99    # but it does not hurt to reset that to zero either.
 100    if with_background:
 101        bg_id = seg_ids[np.argmax(sizes)]
 102        filter_ids = np.concatenate([filter_ids, [bg_id]])
 103
 104    segmentation[np.isin(segmentation, filter_ids)] = 0
 105    segmentation = relabel_sequential(segmentation)[0]
 106
 107    return segmentation
 108
 109
 110#
 111# Classes for automatic instance segmentation
 112#
 113
 114
 115class AMGBase(ABC):
 116    """Base class for the automatic mask generators.
 117    """
 118    def __init__(self):
 119        # the state that has to be computed by the 'initialize' method of the child classes
 120        self._is_initialized = False
 121        self._crop_list = None
 122        self._crop_boxes = None
 123        self._original_size = None
 124
 125    @property
 126    def is_initialized(self):
 127        """Whether the mask generator has already been initialized.
 128        """
 129        return self._is_initialized
 130
 131    @property
 132    def crop_list(self):
 133        """The list of mask data after initialization.
 134        """
 135        return self._crop_list
 136
 137    @property
 138    def crop_boxes(self):
 139        """The list of crop boxes.
 140        """
 141        return self._crop_boxes
 142
 143    @property
 144    def original_size(self):
 145        """The original image size.
 146        """
 147        return self._original_size
 148
 149    def _postprocess_batch(
 150        self,
 151        data,
 152        crop_box,
 153        original_size,
 154        pred_iou_thresh,
 155        stability_score_thresh,
 156        box_nms_thresh,
 157    ):
 158        orig_h, orig_w = original_size
 159
 160        # filter by predicted IoU
 161        if pred_iou_thresh > 0.0:
 162            keep_mask = data["iou_preds"] > pred_iou_thresh
 163            data.filter(keep_mask)
 164
 165        # filter by stability score
 166        if stability_score_thresh > 0.0:
 167            keep_mask = data["stability_score"] >= stability_score_thresh
 168            data.filter(keep_mask)
 169
 170        # filter boxes that touch crop boundaries
 171        keep_mask = ~amg_utils.is_box_near_crop_edge(data["boxes"], crop_box, [0, 0, orig_w, orig_h])
 172        if not torch.all(keep_mask):
 173            data.filter(keep_mask)
 174
 175        # remove duplicates within this crop.
 176        keep_by_nms = batched_nms(
 177            data["boxes"].float(),
 178            data["iou_preds"],
 179            torch.zeros_like(data["boxes"][:, 0]),  # categories
 180            iou_threshold=box_nms_thresh,
 181        )
 182        data.filter(keep_by_nms)
 183
 184        # return to the original image frame
 185        data["boxes"] = amg_utils.uncrop_boxes_xyxy(data["boxes"], crop_box)
 186        data["crop_boxes"] = torch.tensor([crop_box for _ in range(len(data["rles"]))])
 187        # the data from embedding based segmentation doesn't have the points
 188        # so we skip if the corresponding key can't be found
 189        try:
 190            data["points"] = amg_utils.uncrop_points(data["points"], crop_box)
 191        except KeyError:
 192            pass
 193
 194        return data
 195
 196    def _postprocess_small_regions(self, mask_data, min_area, nms_thresh):
 197
 198        if len(mask_data["rles"]) == 0:
 199            return mask_data
 200
 201        # filter small disconnected regions and holes
 202        new_masks = []
 203        scores = []
 204        for rle in mask_data["rles"]:
 205            mask = amg_utils.rle_to_mask(rle)
 206
 207            mask, changed = amg_utils.remove_small_regions(mask, min_area, mode="holes")
 208            unchanged = not changed
 209            mask, changed = amg_utils.remove_small_regions(mask, min_area, mode="islands")
 210            unchanged = unchanged and not changed
 211
 212            new_masks.append(torch.as_tensor(mask, dtype=torch.int).unsqueeze(0))
 213            # give score=0 to changed masks and score=1 to unchanged masks
 214            # so NMS will prefer ones that didn't need postprocessing
 215            scores.append(float(unchanged))
 216
 217        # recalculate boxes and remove any new duplicates
 218        masks = torch.cat(new_masks, dim=0)
 219        boxes = batched_mask_to_box(masks.to(torch.bool))  # Casting this to boolean as we work with one-hot labels.
 220        keep_by_nms = batched_nms(
 221            boxes.float(),
 222            torch.as_tensor(scores, dtype=torch.float),
 223            torch.zeros_like(boxes[:, 0]),  # categories
 224            iou_threshold=nms_thresh,
 225        )
 226
 227        # only recalculate RLEs for masks that have changed
 228        for i_mask in keep_by_nms:
 229            if scores[i_mask] == 0.0:
 230                mask_torch = masks[i_mask].unsqueeze(0)
 231                # mask_data["rles"][i_mask] = amg_utils.mask_to_rle_pytorch(mask_torch)[0]
 232                mask_data["rles"][i_mask] = mask_to_rle_pytorch(mask_torch)[0]
 233                mask_data["boxes"][i_mask] = boxes[i_mask]  # update res directly
 234        mask_data.filter(keep_by_nms)
 235
 236        return mask_data
 237
 238    def _postprocess_masks(self, mask_data, min_mask_region_area, box_nms_thresh, crop_nms_thresh, output_mode):
 239        # filter small disconnected regions and holes in masks
 240        if min_mask_region_area > 0:
 241            mask_data = self._postprocess_small_regions(
 242                mask_data,
 243                min_mask_region_area,
 244                max(box_nms_thresh, crop_nms_thresh),
 245            )
 246
 247        # encode masks
 248        if output_mode == "coco_rle":
 249            mask_data["segmentations"] = [amg_utils.coco_encode_rle(rle) for rle in mask_data["rles"]]
 250        elif output_mode == "binary_mask":
 251            mask_data["segmentations"] = [amg_utils.rle_to_mask(rle) for rle in mask_data["rles"]]
 252        else:
 253            mask_data["segmentations"] = mask_data["rles"]
 254
 255        # write mask records
 256        curr_anns = []
 257        for idx in range(len(mask_data["segmentations"])):
 258            ann = {
 259                "segmentation": mask_data["segmentations"][idx],
 260                "area": amg_utils.area_from_rle(mask_data["rles"][idx]),
 261                "bbox": amg_utils.box_xyxy_to_xywh(mask_data["boxes"][idx]).tolist(),
 262                "predicted_iou": mask_data["iou_preds"][idx].item(),
 263                "stability_score": mask_data["stability_score"][idx].item(),
 264                "crop_box": amg_utils.box_xyxy_to_xywh(mask_data["crop_boxes"][idx]).tolist(),
 265            }
 266            # the data from embedding based segmentation doesn't have the points
 267            # so we skip if the corresponding key can't be found
 268            try:
 269                ann["point_coords"] = [mask_data["points"][idx].tolist()]
 270            except KeyError:
 271                pass
 272            curr_anns.append(ann)
 273
 274        return curr_anns
 275
 276    def _to_mask_data(self, masks, iou_preds, crop_box, original_size, points=None):
 277        orig_h, orig_w = original_size
 278
 279        # serialize predictions and store in MaskData
 280        data = amg_utils.MaskData(masks=masks.flatten(0, 1), iou_preds=iou_preds.flatten(0, 1))
 281        if points is not None:
 282            data["points"] = torch.as_tensor(points.repeat(masks.shape[1], axis=0), dtype=torch.float)
 283
 284        del masks
 285
 286        # calculate the stability scores
 287        data["stability_score"] = amg_utils.calculate_stability_score(
 288            data["masks"], self._predictor.model.mask_threshold, self._stability_score_offset
 289        )
 290
 291        # threshold masks and calculate boxes
 292        data["masks"] = data["masks"] > self._predictor.model.mask_threshold
 293        data["masks"] = data["masks"].type(torch.bool)
 294        data["boxes"] = batched_mask_to_box(data["masks"])
 295
 296        # compress to RLE
 297        data["masks"] = amg_utils.uncrop_masks(data["masks"], crop_box, orig_h, orig_w)
 298        # data["rles"] = amg_utils.mask_to_rle_pytorch(data["masks"])
 299        data["rles"] = mask_to_rle_pytorch(data["masks"])
 300        del data["masks"]
 301
 302        return data
 303
 304    def get_state(self) -> Dict[str, Any]:
 305        """Get the initialized state of the mask generator.
 306
 307        Returns:
 308            State of the mask generator.
 309        """
 310        if not self.is_initialized:
 311            raise RuntimeError("The state has not been computed yet. Call initialize first.")
 312
 313        return {"crop_list": self.crop_list, "crop_boxes": self.crop_boxes, "original_size": self.original_size}
 314
 315    def set_state(self, state: Dict[str, Any]) -> None:
 316        """Set the state of the mask generator.
 317
 318        Args:
 319            state: The state of the mask generator, e.g. from serialized state.
 320        """
 321        self._crop_list = state["crop_list"]
 322        self._crop_boxes = state["crop_boxes"]
 323        self._original_size = state["original_size"]
 324        self._is_initialized = True
 325
 326    def clear_state(self):
 327        """Clear the state of the mask generator.
 328        """
 329        self._crop_list = None
 330        self._crop_boxes = None
 331        self._original_size = None
 332        self._is_initialized = False
 333
 334
 335class AutomaticMaskGenerator(AMGBase):
 336    """Generates an instance segmentation without prompts, using a point grid.
 337
 338    This class implements the same logic as
 339    https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/automatic_mask_generator.py
 340    It decouples the computationally expensive steps of generating masks from the cheap post-processing operation
 341    to filter these masks to enable grid search and interactively changing the post-processing.
 342
 343    Use this class as follows:
 344    ```python
 345    amg = AutomaticMaskGenerator(predictor)
 346    amg.initialize(image)  # Initialize the masks, this takes care of all expensive computations.
 347    masks = amg.generate(pred_iou_thresh=0.8)  # Generate the masks. This is fast and enables testing parameters
 348    ```
 349
 350    Args:
 351        predictor: The segment anything predictor.
 352        points_per_side: The number of points to be sampled along one side of the image.
 353            If None, `point_grids` must provide explicit point sampling.
 354        points_per_batch: The number of points run simultaneously by the model.
 355            Higher numbers may be faster but use more GPU memory.
 356        crop_n_layers: If >0, the mask prediction will be run again on crops of the image.
 357        crop_overlap_ratio: Sets the degree to which crops overlap.
 358        crop_n_points_downscale_factor: How the number of points is downsampled when predicting with crops.
 359        point_grids: A lisst over explicit grids of points used for sampling masks.
 360            Normalized to [0, 1] with respect to the image coordinate system.
 361        stability_score_offset: The amount to shift the cutoff when calculating the stability score.
 362    """
 363    def __init__(
 364        self,
 365        predictor: SamPredictor,
 366        points_per_side: Optional[int] = 32,
 367        points_per_batch: Optional[int] = None,
 368        crop_n_layers: int = 0,
 369        crop_overlap_ratio: float = 512 / 1500,
 370        crop_n_points_downscale_factor: int = 1,
 371        point_grids: Optional[List[np.ndarray]] = None,
 372        stability_score_offset: float = 1.0,
 373    ):
 374        super().__init__()
 375
 376        if points_per_side is not None:
 377            self.point_grids = amg_utils.build_all_layer_point_grids(
 378                points_per_side,
 379                crop_n_layers,
 380                crop_n_points_downscale_factor,
 381            )
 382        elif point_grids is not None:
 383            self.point_grids = point_grids
 384        else:
 385            raise ValueError("Can't have both points_per_side and point_grid be None or not None.")
 386
 387        self._predictor = predictor
 388        self._points_per_side = points_per_side
 389
 390        # we set the points per batch to 16 for mps for performance reasons
 391        # and otherwise keep them at the default of 64
 392        if points_per_batch is None:
 393            points_per_batch = 16 if str(predictor.device) == "mps" else 64
 394        self._points_per_batch = points_per_batch
 395
 396        self._crop_n_layers = crop_n_layers
 397        self._crop_overlap_ratio = crop_overlap_ratio
 398        self._crop_n_points_downscale_factor = crop_n_points_downscale_factor
 399        self._stability_score_offset = stability_score_offset
 400
 401    def _process_batch(self, points, im_size, crop_box, original_size):
 402        # run model on this batch
 403        transformed_points = self._predictor.transform.apply_coords(points, im_size)
 404        in_points = torch.as_tensor(transformed_points, device=self._predictor.device, dtype=torch.float)
 405        in_labels = torch.ones(in_points.shape[0], dtype=torch.int, device=in_points.device)
 406        masks, iou_preds, _ = self._predictor.predict_torch(
 407            in_points[:, None, :],
 408            in_labels[:, None],
 409            multimask_output=True,
 410            return_logits=True,
 411        )
 412        data = self._to_mask_data(masks, iou_preds, crop_box, original_size, points=points)
 413        del masks
 414        return data
 415
 416    def _process_crop(self, image, crop_box, crop_layer_idx, precomputed_embeddings, pbar_init=None, pbar_update=None):
 417        # Crop the image and calculate embeddings.
 418        x0, y0, x1, y1 = crop_box
 419        cropped_im = image[y0:y1, x0:x1, :]
 420        cropped_im_size = cropped_im.shape[:2]
 421
 422        if not precomputed_embeddings:
 423            self._predictor.set_image(cropped_im)
 424
 425        # Get the points for this crop.
 426        points_scale = np.array(cropped_im_size)[None, ::-1]
 427        points_for_image = self.point_grids[crop_layer_idx] * points_scale
 428
 429        # Generate masks for this crop in batches.
 430        data = amg_utils.MaskData()
 431        n_batches = len(points_for_image) // self._points_per_batch +\
 432            int(len(points_for_image) % self._points_per_batch != 0)
 433        if pbar_init is not None:
 434            pbar_init(n_batches, "Predict masks for point grid prompts")
 435
 436        for (points,) in amg_utils.batch_iterator(self._points_per_batch, points_for_image):
 437            batch_data = self._process_batch(points, cropped_im_size, crop_box, self.original_size)
 438            data.cat(batch_data)
 439            del batch_data
 440            if pbar_update is not None:
 441                pbar_update(1)
 442
 443        if not precomputed_embeddings:
 444            self._predictor.reset_image()
 445
 446        return data
 447
 448    @torch.no_grad()
 449    def initialize(
 450        self,
 451        image: np.ndarray,
 452        image_embeddings: Optional[util.ImageEmbeddings] = None,
 453        i: Optional[int] = None,
 454        verbose: bool = False,
 455        pbar_init: Optional[callable] = None,
 456        pbar_update: Optional[callable] = None,
 457    ) -> None:
 458        """Initialize image embeddings and masks for an image.
 459
 460        Args:
 461            image: The input image, volume or timeseries.
 462            image_embeddings: Optional precomputed image embeddings.
 463                See `util.precompute_image_embeddings` for details.
 464            i: Index for the image data. Required if `image` has three spatial dimensions
 465                or a time dimension and two spatial dimensions.
 466            verbose: Whether to print computation progress.
 467            pbar_init: Callback to initialize an external progress bar. Must accept number of steps and description.
 468                Can be used together with pbar_update to handle napari progress bar in other thread.
 469                To enables using this function within a threadworker.
 470            pbar_update: Callback to update an external progress bar.
 471        """
 472        original_size = image.shape[:2]
 473        self._original_size = original_size
 474
 475        crop_boxes, layer_idxs = amg_utils.generate_crop_boxes(
 476            original_size, self._crop_n_layers, self._crop_overlap_ratio
 477        )
 478
 479        # We can set fixed image embeddings if we only have a single crop box (the default setting).
 480        # Otherwise we have to recompute the embeddings for each crop and can't precompute.
 481        if len(crop_boxes) == 1:
 482            if image_embeddings is None:
 483                image_embeddings = util.precompute_image_embeddings(self._predictor, image)
 484            util.set_precomputed(self._predictor, image_embeddings, i=i)
 485            precomputed_embeddings = True
 486        else:
 487            precomputed_embeddings = False
 488
 489        # we need to cast to the image representation that is compatible with SAM
 490        image = util._to_image(image)
 491
 492        _, pbar_init, pbar_update, pbar_close = util.handle_pbar(verbose, pbar_init, pbar_update)
 493
 494        crop_list = []
 495        for crop_box, layer_idx in zip(crop_boxes, layer_idxs):
 496            crop_data = self._process_crop(
 497                image, crop_box, layer_idx,
 498                precomputed_embeddings=precomputed_embeddings,
 499                pbar_init=pbar_init, pbar_update=pbar_update,
 500            )
 501            crop_list.append(crop_data)
 502        pbar_close()
 503
 504        self._is_initialized = True
 505        self._crop_list = crop_list
 506        self._crop_boxes = crop_boxes
 507
 508    @torch.no_grad()
 509    def generate(
 510        self,
 511        pred_iou_thresh: float = 0.88,
 512        stability_score_thresh: float = 0.95,
 513        box_nms_thresh: float = 0.7,
 514        crop_nms_thresh: float = 0.7,
 515        min_mask_region_area: int = 0,
 516        output_mode: str = "binary_mask",
 517    ) -> List[Dict[str, Any]]:
 518        """Generate instance segmentation for the currently initialized image.
 519
 520        Args:
 521            pred_iou_thresh: Filter threshold in [0, 1], using the mask quality predicted by the model.
 522            stability_score_thresh: Filter threshold in [0, 1], using the stability of the mask
 523                under changes to the cutoff used to binarize the model prediction.
 524            box_nms_thresh: The IoU threshold used by nonmax suppression to filter duplicate masks.
 525            crop_nms_thresh: The IoU threshold used by nonmax suppression to filter duplicate masks between crops.
 526            min_mask_region_area: Minimal size for the predicted masks.
 527            output_mode: The form masks are returned in.
 528
 529        Returns:
 530            The instance segmentation masks.
 531        """
 532        if not self.is_initialized:
 533            raise RuntimeError("AutomaticMaskGenerator has not been initialized. Call initialize first.")
 534
 535        data = amg_utils.MaskData()
 536        for data_, crop_box in zip(self.crop_list, self.crop_boxes):
 537            crop_data = self._postprocess_batch(
 538                data=deepcopy(data_),
 539                crop_box=crop_box, original_size=self.original_size,
 540                pred_iou_thresh=pred_iou_thresh,
 541                stability_score_thresh=stability_score_thresh,
 542                box_nms_thresh=box_nms_thresh
 543            )
 544            data.cat(crop_data)
 545
 546        if len(self.crop_boxes) > 1 and len(data["crop_boxes"]) > 0:
 547            # Prefer masks from smaller crops
 548            scores = 1 / box_area(data["crop_boxes"])
 549            scores = scores.to(data["boxes"].device)
 550            keep_by_nms = batched_nms(
 551                data["boxes"].float(),
 552                scores,
 553                torch.zeros_like(data["boxes"][:, 0]),  # categories
 554                iou_threshold=crop_nms_thresh,
 555            )
 556            data.filter(keep_by_nms)
 557
 558        data.to_numpy()
 559        masks = self._postprocess_masks(data, min_mask_region_area, box_nms_thresh, crop_nms_thresh, output_mode)
 560        return masks
 561
 562
 563# Helper function for tiled embedding computation and checking consistent state.
 564def _process_tiled_embeddings(predictor, image, image_embeddings, tile_shape, halo):
 565    if image_embeddings is None:
 566        if tile_shape is None or halo is None:
 567            raise ValueError("To compute tiled embeddings the parameters tile_shape and halo have to be passed.")
 568        image_embeddings = util.precompute_image_embeddings(predictor, image, tile_shape=tile_shape, halo=halo)
 569
 570    # Use tile shape and halo from the precomputed embeddings if not given.
 571    # Otherwise check that they are consistent.
 572    feats = image_embeddings["features"]
 573    tile_shape_, halo_ = feats.attrs["tile_shape"], feats.attrs["halo"]
 574    if tile_shape is None:
 575        tile_shape = tile_shape_
 576    elif tile_shape != tile_shape_:
 577        raise ValueError(
 578            f"Inconsistent tile_shape parameter {tile_shape} with precomputed embeedings: {tile_shape_}."
 579        )
 580    if halo is None:
 581        halo = halo_
 582    elif halo != halo_:
 583        raise ValueError(f"Inconsistent halo parameter {halo} with precomputed embeedings: {halo_}.")
 584
 585    return image_embeddings, tile_shape, halo
 586
 587
 588class TiledAutomaticMaskGenerator(AutomaticMaskGenerator):
 589    """Generates an instance segmentation without prompts, using a point grid.
 590
 591    Implements the same functionality as `AutomaticMaskGenerator` but for tiled embeddings.
 592
 593    Args:
 594        predictor: The segment anything predictor.
 595        points_per_side: The number of points to be sampled along one side of the image.
 596            If None, `point_grids` must provide explicit point sampling.
 597        points_per_batch: The number of points run simultaneously by the model.
 598            Higher numbers may be faster but use more GPU memory.
 599        point_grids: A lisst over explicit grids of points used for sampling masks.
 600            Normalized to [0, 1] with respect to the image coordinate system.
 601        stability_score_offset: The amount to shift the cutoff when calculating the stability score.
 602    """
 603
 604    # We only expose the arguments that make sense for the tiled mask generator.
 605    # Anything related to crops doesn't make sense, because we re-use that functionality
 606    # for tiling, so these parameters wouldn't have any effect.
 607    def __init__(
 608        self,
 609        predictor: SamPredictor,
 610        points_per_side: Optional[int] = 32,
 611        points_per_batch: int = 64,
 612        point_grids: Optional[List[np.ndarray]] = None,
 613        stability_score_offset: float = 1.0,
 614    ) -> None:
 615        super().__init__(
 616            predictor=predictor,
 617            points_per_side=points_per_side,
 618            points_per_batch=points_per_batch,
 619            point_grids=point_grids,
 620            stability_score_offset=stability_score_offset,
 621        )
 622
 623    @torch.no_grad()
 624    def initialize(
 625        self,
 626        image: np.ndarray,
 627        image_embeddings: Optional[util.ImageEmbeddings] = None,
 628        i: Optional[int] = None,
 629        tile_shape: Optional[Tuple[int, int]] = None,
 630        halo: Optional[Tuple[int, int]] = None,
 631        verbose: bool = False,
 632        pbar_init: Optional[callable] = None,
 633        pbar_update: Optional[callable] = None,
 634    ) -> None:
 635        """Initialize image embeddings and masks for an image.
 636
 637        Args:
 638            image: The input image, volume or timeseries.
 639            image_embeddings: Optional precomputed image embeddings.
 640                See `util.precompute_image_embeddings` for details.
 641            i: Index for the image data. Required if `image` has three spatial dimensions
 642                or a time dimension and two spatial dimensions.
 643            tile_shape: The tile shape for embedding prediction.
 644            halo: The overlap of between tiles.
 645            verbose: Whether to print computation progress.
 646            pbar_init: Callback to initialize an external progress bar. Must accept number of steps and description.
 647                Can be used together with pbar_update to handle napari progress bar in other thread.
 648                To enables using this function within a threadworker.
 649            pbar_update: Callback to update an external progress bar.
 650        """
 651        original_size = image.shape[:2]
 652        self._original_size = original_size
 653
 654        image_embeddings, tile_shape, halo = _process_tiled_embeddings(
 655            self._predictor, image, image_embeddings, tile_shape, halo
 656        )
 657
 658        tiling = blocking([0, 0], original_size, tile_shape)
 659        n_tiles = tiling.numberOfBlocks
 660
 661        # The crop box is always the full local tile.
 662        tiles = [tiling.getBlockWithHalo(tile_id, list(halo)).outerBlock for tile_id in range(n_tiles)]
 663        crop_boxes = [[tile.begin[1], tile.begin[0], tile.end[1], tile.end[0]] for tile in tiles]
 664
 665        _, pbar_init, pbar_update, pbar_close = util.handle_pbar(verbose, pbar_init, pbar_update)
 666        pbar_init(n_tiles, "Compute masks for tile")
 667
 668        # We need to cast to the image representation that is compatible with SAM.
 669        image = util._to_image(image)
 670
 671        mask_data = []
 672        for tile_id in range(n_tiles):
 673            # set the pre-computed embeddings for this tile
 674            features = image_embeddings["features"][tile_id]
 675            tile_embeddings = {
 676                "features": features,
 677                "input_size": features.attrs["input_size"],
 678                "original_size": features.attrs["original_size"],
 679            }
 680            util.set_precomputed(self._predictor, tile_embeddings, i)
 681
 682            # compute the mask data for this tile and append it
 683            this_mask_data = self._process_crop(
 684                image, crop_box=crop_boxes[tile_id], crop_layer_idx=0, precomputed_embeddings=True
 685            )
 686            mask_data.append(this_mask_data)
 687            pbar_update(1)
 688        pbar_close()
 689
 690        # set the initialized data
 691        self._is_initialized = True
 692        self._crop_list = mask_data
 693        self._crop_boxes = crop_boxes
 694
 695
 696#
 697# Instance segmentation functionality based on fine-tuned decoder
 698#
 699
 700
 701class DecoderAdapter(torch.nn.Module):
 702    """Adapter to contain the UNETR decoder in a single module.
 703
 704    To apply the decoder on top of pre-computed embeddings for
 705    the segmentation functionality.
 706    See also: https://github.com/constantinpape/torch-em/blob/main/torch_em/model/unetr.py
 707    """
 708    def __init__(self, unetr):
 709        super().__init__()
 710
 711        self.base = unetr.base
 712        self.out_conv = unetr.out_conv
 713        self.deconv_out = unetr.deconv_out
 714        self.decoder_head = unetr.decoder_head
 715        self.final_activation = unetr.final_activation
 716        self.postprocess_masks = unetr.postprocess_masks
 717
 718        self.decoder = unetr.decoder
 719        self.deconv1 = unetr.deconv1
 720        self.deconv2 = unetr.deconv2
 721        self.deconv3 = unetr.deconv3
 722        self.deconv4 = unetr.deconv4
 723
 724    def forward(self, input_, input_shape, original_shape):
 725        z12 = input_
 726
 727        z9 = self.deconv1(z12)
 728        z6 = self.deconv2(z9)
 729        z3 = self.deconv3(z6)
 730        z0 = self.deconv4(z3)
 731
 732        updated_from_encoder = [z9, z6, z3]
 733
 734        x = self.base(z12)
 735        x = self.decoder(x, encoder_inputs=updated_from_encoder)
 736        x = self.deconv_out(x)
 737
 738        x = torch.cat([x, z0], dim=1)
 739        x = self.decoder_head(x)
 740
 741        x = self.out_conv(x)
 742        if self.final_activation is not None:
 743            x = self.final_activation(x)
 744
 745        x = self.postprocess_masks(x, input_shape, original_shape)
 746        return x
 747
 748
 749def get_unetr(
 750    image_encoder: torch.nn.Module,
 751    decoder_state: Optional[OrderedDict[str, torch.Tensor]] = None,
 752    device: Optional[Union[str, torch.device]] = None,
 753) -> torch.nn.Module:
 754    """Get UNETR model for automatic instance segmentation.
 755
 756    Args:
 757        image_encoder: The image encoder of the SAM model.
 758            This is used as encoder by the UNETR too.
 759        decoder_state: Optional decoder state to initialize the weights
 760            of the UNETR decoder.
 761        device: The device.
 762    Returns:
 763        The UNETR model.
 764    """
 765    device = util.get_device(device)
 766
 767    unetr = UNETR(
 768        backbone="sam",
 769        encoder=image_encoder,
 770        out_channels=3,
 771        use_sam_stats=True,
 772        final_activation="Sigmoid",
 773        use_skip_connection=False,
 774        resize_input=True,
 775    )
 776    if decoder_state is not None:
 777        unetr_state_dict = unetr.state_dict()
 778        for k, v in unetr_state_dict.items():
 779            if not k.startswith("encoder"):
 780                unetr_state_dict[k] = decoder_state[k]
 781        unetr.load_state_dict(unetr_state_dict)
 782
 783    unetr.to(device)
 784    return unetr
 785
 786
 787def get_decoder(
 788    image_encoder: torch.nn.Module,
 789    decoder_state: OrderedDict[str, torch.Tensor],
 790    device: Optional[Union[str, torch.device]] = None,
 791) -> DecoderAdapter:
 792    """Get decoder to predict outputs for automatic instance segmentation
 793
 794    Args:
 795        image_encoder: The image encoder of the SAM model.
 796        decoder_state: State to initialize the weights of the UNETR decoder.
 797        device: The device.
 798    Returns:
 799        The decoder for instance segmentation.
 800    """
 801    unetr = get_unetr(image_encoder, decoder_state, device)
 802    return DecoderAdapter(unetr)
 803
 804
 805def get_predictor_and_decoder(
 806    model_type: str,
 807    checkpoint_path: Union[str, os.PathLike],
 808    device: Optional[Union[str, torch.device]] = None,
 809    peft_kwargs: Optional[Dict] = None,
 810) -> Tuple[SamPredictor, DecoderAdapter]:
 811    """Load the SAM model (predictor) and instance segmentation decoder.
 812
 813    This requires a checkpoint that contains the state for both predictor
 814    and decoder.
 815
 816    Args:
 817        model_type: The type of the image encoder used in the SAM model.
 818        checkpoint_path: Path to the checkpoint from which to load the data.
 819        device: The device.
 820        lora_rank: The rank for low rank adaptation of the attention layers.
 821
 822    Returns:
 823        The SAM predictor.
 824        The decoder for instance segmentation.
 825    """
 826    device = util.get_device(device)
 827    predictor, state = util.get_sam_model(
 828        model_type=model_type,
 829        checkpoint_path=checkpoint_path,
 830        device=device,
 831        return_state=True,
 832        peft_kwargs=peft_kwargs,
 833    )
 834    if "decoder_state" not in state:
 835        raise ValueError(
 836            f"The checkpoint at '{checkpoint_path}' or the chosen model '{model_type}' does not contain a decoder state"
 837        )
 838    decoder = get_decoder(predictor.model.image_encoder, state["decoder_state"], device)
 839    return predictor, decoder
 840
 841
 842class InstanceSegmentationWithDecoder:
 843    """Generates an instance segmentation without prompts, using a decoder.
 844
 845    Implements the same interface as `AutomaticMaskGenerator`.
 846
 847    Use this class as follows:
 848    ```python
 849    segmenter = InstanceSegmentationWithDecoder(predictor, decoder)
 850    segmenter.initialize(image)   # Predict the image embeddings and decoder outputs.
 851    masks = segmenter.generate(center_distance_threshold=0.75)  # Generate the instance segmentation.
 852    ```
 853
 854    Args:
 855        predictor: The segment anything predictor.
 856        decoder: The decoder to predict intermediate representations
 857            for instance segmentation.
 858    """
 859    def __init__(
 860        self,
 861        predictor: SamPredictor,
 862        decoder: torch.nn.Module,
 863    ) -> None:
 864        self._predictor = predictor
 865        self._decoder = decoder
 866
 867        # The decoder outputs.
 868        self._foreground = None
 869        self._center_distances = None
 870        self._boundary_distances = None
 871
 872        self._is_initialized = False
 873
 874    @property
 875    def is_initialized(self):
 876        """Whether the mask generator has already been initialized.
 877        """
 878        return self._is_initialized
 879
 880    @torch.no_grad()
 881    def initialize(
 882        self,
 883        image: np.ndarray,
 884        image_embeddings: Optional[util.ImageEmbeddings] = None,
 885        i: Optional[int] = None,
 886        verbose: bool = False,
 887        pbar_init: Optional[callable] = None,
 888        pbar_update: Optional[callable] = None,
 889    ) -> None:
 890        """Initialize image embeddings and decoder predictions for an image.
 891
 892        Args:
 893            image: The input image, volume or timeseries.
 894            image_embeddings: Optional precomputed image embeddings.
 895                See `util.precompute_image_embeddings` for details.
 896            i: Index for the image data. Required if `image` has three spatial dimensions
 897                or a time dimension and two spatial dimensions.
 898            verbose: Whether to be verbose.
 899            pbar_init: Callback to initialize an external progress bar. Must accept number of steps and description.
 900                Can be used together with pbar_update to handle napari progress bar in other thread.
 901                To enables using this function within a threadworker.
 902            pbar_update: Callback to update an external progress bar.
 903        """
 904        _, pbar_init, pbar_update, pbar_close = util.handle_pbar(verbose, pbar_init, pbar_update)
 905        pbar_init(1, "Initialize instance segmentation with decoder")
 906
 907        if image_embeddings is None:
 908            image_embeddings = util.precompute_image_embeddings(self._predictor, image)
 909
 910        # Get the image embeddings from the predictor.
 911        self._predictor = util.set_precomputed(self._predictor, image_embeddings, i=i)
 912        embeddings = self._predictor.features
 913        input_shape = tuple(self._predictor.input_size)
 914        original_shape = tuple(self._predictor.original_size)
 915
 916        # Run prediction with the UNETR decoder.
 917        output = self._decoder(embeddings, input_shape, original_shape).cpu().numpy().squeeze(0)
 918        assert output.shape[0] == 3, f"{output.shape}"
 919        pbar_update(1)
 920        pbar_close()
 921
 922        # Set the state.
 923        self._foreground = output[0]
 924        self._center_distances = output[1]
 925        self._boundary_distances = output[2]
 926        self._is_initialized = True
 927
 928    def _to_masks(self, segmentation, output_mode):
 929        if output_mode != "binary_mask":
 930            raise NotImplementedError
 931
 932        props = regionprops(segmentation)
 933        ndim = segmentation.ndim
 934        assert ndim in (2, 3)
 935
 936        shape = segmentation.shape
 937        if ndim == 2:
 938            crop_box = [0, shape[1], 0, shape[0]]
 939        else:
 940            crop_box = [0, shape[2], 0, shape[1], 0, shape[0]]
 941
 942        # go from skimage bbox in format [y0, x0, y1, x1] to SAM format [x0, w, y0, h]
 943        def to_bbox_2d(bbox):
 944            y0, x0 = bbox[0], bbox[1]
 945            w = bbox[3] - x0
 946            h = bbox[2] - y0
 947            return [x0, w, y0, h]
 948
 949        def to_bbox_3d(bbox):
 950            z0, y0, x0 = bbox[0], bbox[1], bbox[2]
 951            w = bbox[5] - x0
 952            h = bbox[4] - y0
 953            d = bbox[3] - y0
 954            return [x0, w, y0, h, z0, d]
 955
 956        to_bbox = to_bbox_2d if ndim == 2 else to_bbox_3d
 957        masks = [
 958            {
 959                "segmentation": segmentation == prop.label,
 960                "area": prop.area,
 961                "bbox": to_bbox(prop.bbox),
 962                "crop_box": crop_box,
 963                "seg_id": prop.label,
 964            } for prop in props
 965        ]
 966        return masks
 967
 968    def generate(
 969        self,
 970        center_distance_threshold: float = 0.5,
 971        boundary_distance_threshold: float = 0.5,
 972        foreground_threshold: float = 0.5,
 973        foreground_smoothing: float = 1.0,
 974        distance_smoothing: float = 1.6,
 975        min_size: int = 0,
 976        output_mode: Optional[str] = "binary_mask",
 977    ) -> List[Dict[str, Any]]:
 978        """Generate instance segmentation for the currently initialized image.
 979
 980        Args:
 981            center_distance_threshold: Center distance predictions below this value will be
 982                used to find seeds (intersected with thresholded boundary distance predictions).
 983            boundary_distance_threshold: Boundary distance predictions below this value will be
 984                used to find seeds (intersected with thresholded center distance predictions).
 985            foreground_smoothing: Sigma value for smoothing the foreground predictions, to avoid
 986                checkerboard artifacts in the prediction.
 987            foreground_threshold: Foreground predictions above this value will be used as foreground mask.
 988            distance_smoothing: Sigma value for smoothing the distance predictions.
 989            min_size: Minimal object size in the segmentation result.
 990            output_mode: The form masks are returned in. Pass None to directly return the instance segmentation.
 991
 992        Returns:
 993            The instance segmentation masks.
 994        """
 995        if not self.is_initialized:
 996            raise RuntimeError("InstanceSegmentationWithDecoder has not been initialized. Call initialize first.")
 997
 998        if foreground_smoothing > 0:
 999            foreground = vigra.filters.gaussianSmoothing(self._foreground, foreground_smoothing)
1000        else:
1001            foreground = self._foreground
1002        # Further optimization: parallel implementation using elf.parallel functionality.
1003        # (Make sure to expose n_threads to avoid over-subscription in case of outer parallelization)
1004        segmentation = watershed_from_center_and_boundary_distances(
1005            self._center_distances, self._boundary_distances, foreground,
1006            center_distance_threshold=center_distance_threshold,
1007            boundary_distance_threshold=boundary_distance_threshold,
1008            foreground_threshold=foreground_threshold,
1009            distance_smoothing=distance_smoothing,
1010            min_size=min_size,
1011        )
1012        if output_mode is not None:
1013            segmentation = self._to_masks(segmentation, output_mode)
1014        return segmentation
1015
1016    def get_state(self) -> Dict[str, Any]:
1017        """Get the initialized state of the instance segmenter.
1018
1019        Returns:
1020            Instance segmentation state.
1021        """
1022        if not self.is_initialized:
1023            raise RuntimeError("The state has not been computed yet. Call initialize first.")
1024
1025        return {
1026            "foreground": self._foreground,
1027            "center_distances": self._center_distances,
1028            "boundary_distances": self._boundary_distances,
1029        }
1030
1031    def set_state(self, state: Dict[str, Any]) -> None:
1032        """Set the state of the instance segmenter.
1033
1034        Args:
1035            state: The instance segmentation state
1036        """
1037        self._foreground = state["foreground"]
1038        self._center_distances = state["center_distances"]
1039        self._boundary_distances = state["boundary_distances"]
1040        self._is_initialized = True
1041
1042    def clear_state(self):
1043        """Clear the state of the instance segmenter.
1044        """
1045        self._foreground = None
1046        self._center_distances = None
1047        self._boundary_distances = None
1048        self._is_initialized = False
1049
1050
1051class TiledInstanceSegmentationWithDecoder(InstanceSegmentationWithDecoder):
1052    """Same as `InstanceSegmentationWithDecoder` but for tiled image embeddings.
1053    """
1054
1055    @torch.no_grad()
1056    def initialize(
1057        self,
1058        image: np.ndarray,
1059        image_embeddings: Optional[util.ImageEmbeddings] = None,
1060        i: Optional[int] = None,
1061        tile_shape: Optional[Tuple[int, int]] = None,
1062        halo: Optional[Tuple[int, int]] = None,
1063        verbose: bool = False,
1064        pbar_init: Optional[callable] = None,
1065        pbar_update: Optional[callable] = None,
1066    ) -> None:
1067        """Initialize image embeddings and decoder predictions for an image.
1068
1069        Args:
1070            image: The input image, volume or timeseries.
1071            image_embeddings: Optional precomputed image embeddings.
1072                See `util.precompute_image_embeddings` for details.
1073            i: Index for the image data. Required if `image` has three spatial dimensions
1074                or a time dimension and two spatial dimensions.
1075            verbose: Dummy input to be compatible with other function signatures.
1076            pbar_init: Callback to initialize an external progress bar. Must accept number of steps and description.
1077                Can be used together with pbar_update to handle napari progress bar in other thread.
1078                To enables using this function within a threadworker.
1079            pbar_update: Callback to update an external progress bar.
1080        """
1081        original_size = image.shape[:2]
1082        image_embeddings, tile_shape, halo = _process_tiled_embeddings(
1083            self._predictor, image, image_embeddings, tile_shape, halo
1084        )
1085        tiling = blocking([0, 0], original_size, tile_shape)
1086
1087        _, pbar_init, pbar_update, pbar_close = util.handle_pbar(verbose, pbar_init, pbar_update)
1088        pbar_init(tiling.numberOfBlocks, "Initialize tiled instance segmentation with decoder")
1089
1090        foreground = np.zeros(original_size, dtype="float32")
1091        center_distances = np.zeros(original_size, dtype="float32")
1092        boundary_distances = np.zeros(original_size, dtype="float32")
1093
1094        for tile_id in range(tiling.numberOfBlocks):
1095
1096            # Get the image embeddings from the predictor for this tile.
1097            self._predictor = util.set_precomputed(self._predictor, image_embeddings, i=i, tile_id=tile_id)
1098            embeddings = self._predictor.features
1099            input_shape = tuple(self._predictor.input_size)
1100            original_shape = tuple(self._predictor.original_size)
1101
1102            # Predict with the UNETR decoder for this tile.
1103            output = self._decoder(embeddings, input_shape, original_shape).cpu().numpy().squeeze(0)
1104            assert output.shape[0] == 3, f"{output.shape}"
1105
1106            # Set the predictions in the output for this tile.
1107            block = tiling.getBlockWithHalo(tile_id, halo=list(halo))
1108            local_bb = tuple(
1109                slice(beg, end) for beg, end in zip(block.innerBlockLocal.begin, block.innerBlockLocal.end)
1110            )
1111            inner_bb = tuple(slice(beg, end) for beg, end in zip(block.innerBlock.begin, block.innerBlock.end))
1112
1113            foreground[inner_bb] = output[0][local_bb]
1114            center_distances[inner_bb] = output[1][local_bb]
1115            boundary_distances[inner_bb] = output[2][local_bb]
1116        pbar_close()
1117
1118        # Set the state.
1119        self._foreground = foreground
1120        self._center_distances = center_distances
1121        self._boundary_distances = boundary_distances
1122        self._is_initialized = True
1123
1124
1125def get_amg(
1126    predictor: SamPredictor, is_tiled: bool, decoder: Optional[torch.nn.Module] = None, **kwargs,
1127) -> Union[AMGBase, InstanceSegmentationWithDecoder]:
1128    """Get the automatic mask generator class.
1129
1130    Args:
1131        predictor: The segment anything predictor.
1132        is_tiled: Whether tiled embeddings are used.
1133        decoder: Decoder to predict instacne segmmentation.
1134        kwargs: The keyword arguments for the amg class.
1135
1136    Returns:
1137        The automatic mask generator.
1138    """
1139    if decoder is None:
1140        segmenter_class = TiledAutomaticMaskGenerator if is_tiled else AutomaticMaskGenerator
1141        segmenter = segmenter_class(predictor, **kwargs)
1142    else:
1143        segmenter_class = TiledInstanceSegmentationWithDecoder if is_tiled else InstanceSegmentationWithDecoder
1144        segmenter = segmenter_class(predictor, decoder, **kwargs)
1145
1146    return segmenter