micro_sam.instance_segmentation

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