micro-sam 1.8.0__py3-none-any.whl

micro_sam/__init__.py

@@ -0,0 +1,19 @@
+"""
+.. include:: ../doc/start_page.md
+.. include:: ../doc/installation.md
+.. include:: ../doc/annotation_tools.md
+.. include:: ../doc/cli_tools.md
+.. include:: ../doc/python_library.md
+.. include:: ../doc/finetuned_models.md
+.. include:: ../doc/apg.md
+.. include:: ../doc/data_submission.md
+.. include:: ../doc/faq.md
+.. include:: ../doc/contributing.md
+.. include:: ../doc/band.md
+"""
+
+import os
+
+from .__version__ import __version__
+
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"

micro_sam/__version__.py

@@ -0,0 +1 @@
+__version__ = "1.8.0"

micro_sam/_model_settings.py

@@ -0,0 +1,142 @@
+# The settings for the instance segmentation widget with ais.
+AIS_SETTINGS = {
+    "vit_t_lm": {
+        "center_distance_thresh": 0.5,
+        "boundary_distance_thresh": 0.5,
+        "distance_smoothing": 2.0,
+        "min_size": 100,
+    },
+    "vit_b_lm": {
+        "center_distance_thresh": 0.4,
+        "boundary_distance_thresh": 0.5,
+        "distance_smoothing": 2.0,
+        "min_size": 100,
+    },
+    "vit_l_lm": {
+        "center_distance_thresh": 0.4,
+        "boundary_distance_thresh": 0.4,
+        "distance_smoothing": 1.6,
+        "min_size": 100,
+    },
+    "vit_h_lm": {
+        "center_distance_thresh": 0.5,
+        "boundary_distance_thresh": 0.5,
+        "distance_smoothing": 1.4,
+        "min_size": 100,
+    },
+
+    "vit_t_em_organelles": {
+        "center_distance_thresh": 0.4,
+        "boundary_distance_thresh": 0.5,
+        "distance_smoothing": 1.2,
+        "min_size": 100,
+    },
+    "vit_b_em_organelles": {
+        "center_distance_thresh": 0.3,
+        "boundary_distance_thresh": 0.4,
+        "distance_smoothing": 1.2,
+        "min_size": 100,
+    },
+    "vit_l_em_organelles": {
+        "center_distance_thresh": 0.3,
+        "boundary_distance_thresh": 0.4,
+        "distance_smoothing": 1.2,
+        "min_size": 100,
+    },
+    "vit_h_em_organelles": {
+        "center_distance_thresh": 0.3,
+        "boundary_distance_thresh": 0.4,
+        "distance_smoothing": 1.2,
+        "min_size": 100,
+    }
+}
+
+# The settings for the instance segmentation widget with amg.
+AMG_SETTINGS = {
+    "vit_t_lm": {
+        "pred_iou_thresh": 0.6,
+        "stability_score_thresh": 0.65,
+        "min_object_size": 100,
+    },
+    "vit_b_lm": {
+        "pred_iou_thresh": 0.65,
+        "stability_score_thresh": 0.7,
+        "min_object_size": 100,
+    },
+    "vit_l_lm": {
+        "pred_iou_thresh": 0.65,
+        "stability_score_thresh": 0.73,
+        "min_object_size": 100,
+    },
+    "vit_h_lm": {
+        "pred_iou_thresh": 0.65,
+        "stability_score_thresh": 0.7,
+        "min_object_size": 100,
+    },
+
+    "vit_t_em_organelles": {
+        "pred_iou_thresh": 0.75,
+        "stability_score_thresh": 0.75,
+        "min_object_size": 100,
+    },
+    "vit_b_em_organelles": {
+        "pred_iou_thresh": 0.75,
+        "stability_score_thresh": 0.75,
+        "min_object_size": 100,
+    },
+    "vit_l_em_organelles": {
+        "pred_iou_thresh": 0.8,
+        "stability_score_thresh": 0.8,
+        "min_object_size": 100,
+    },
+    "vit_h_em_organelles": {
+        "pred_iou_thresh": 0.8,
+        "stability_score_thresh": 0.8,
+        "min_object_size": 100,
+    },
+}
+
+# The settings for the nd segment widget.
+ND_SEGMENT_SETTINGS = {
+    "vit_t_lm": {
+        "projection_mode": "box",
+        "iou_threshold": 0.8,
+        "box_extension": 0.025,
+    },
+    "vit_b_lm": {
+        "projection_mode": "box",
+        "iou_threshold": 0.8,
+        "box_extension": 0.025,
+    },
+    "vit_l_lm": {
+        "projection_mode": "box",
+        "iou_threshold": 0.8,
+        "box_extension": 0.025,
+    },
+    "vit_h_lm": {
+        "projection_mode": "box",
+        "iou_threshold": 0.8,
+        "box_extension": 0.0025,
+    },
+
+    "vit_t_em_organelles": {
+        "projection_mode": "single_point",
+        "iou_threshold": 0.6,
+        "box_extension": 0.025,
+    },
+    "vit_b_em_organelles": {
+        "projection_mode": "single_point",
+        "iou_threshold": 0.6,
+        "box_extension": 0.025,
+    },
+    "vit_l_em_organelles": {
+        "projection_mode": "single_point",
+        "iou_threshold": 0.6,
+        "box_extension": 0.025,
+    },
+    "vit_h_em_organelles": {
+        "projection_mode": "single_point",
+        "iou_threshold": 0.6,
+        "box_extension": 0.025,
+    }
+}

micro_sam/_test_util.py

@@ -0,0 +1,22 @@
+import numpy as np
+
+
+def check_layer_initialization(viewer, expected_shape):
+    """Utility function to check the initial layer setup is correct."""
+
+    assert len(viewer.layers) == 6
+    expected_layer_names = [
+        "image", "auto_segmentation", "committed_objects", "current_object", "point_prompts", "prompts"
+    ]
+
+    for layer_name in expected_layer_names:
+        assert layer_name in viewer.layers
+
+    # Check prompt layers
+    assert viewer.layers["prompts"].data == []  # shape data is list, not numpy array
+    np.testing.assert_equal(viewer.layers["point_prompts"].data, 0)
+
+    # Check segmentation layers.
+    for layer_name in ["auto_segmentation", "committed_objects", "current_object"]:
+        assert viewer.layers[layer_name].data.shape == expected_shape
+        np.testing.assert_equal(viewer.layers[layer_name].data, 0)

micro_sam/_vendored.py

@@ -0,0 +1,152 @@
+"""Functions from other third party libraries.
+
+We can remove these functions once the bugs affecting our code is fixed upstream.
+
+The license type of the thrid party software project must be compatible with
+the software license the micro-sam project is distributed under.
+"""
+
+from typing import Any, Dict, List, Literal
+
+import numpy as np
+
+import torch
+
+try:
+    from numba import njit
+    HAVE_NUMBA = True
+except (ImportError, SystemError):
+    HAVE_NUMBA = False
+
+    def njit(func):
+        return func
+
+from bioimage_cpp.utils import compute_rle as _compute_rle_cpp
+
+
+def _compute_rle_bioimage_cpp(mask):
+    # bioimage-cpp returns an int64 numpy array; convert to a list to match the
+    # numpy/numba implementations and the COCO/SAM RLE format.
+    return _compute_rle_cpp(mask).tolist()
+
+
+def batched_mask_to_box(masks: torch.Tensor) -> torch.Tensor:
+    """Calculates boxes in XYXY format around masks. Return [0, 0, 0, 0] for an empty mask.
+
+    For input shape C1xC2x...xHxW, the output shape is C1xC2x...x4.
+
+    This function is adapted from https://github.com/facebookresearch/segment-anything/segment_anything/util/amg.py
+    so that it also supports tensors that have MPS device.
+    It further ensures that inputs are boolean tensors, otherwise the function yields wrong results.
+    See https://github.com/facebookresearch/segment-anything/issues/552 for details.
+    """
+    assert masks.dtype == torch.bool, masks.dtype
+
+    # torch.max below raises an error on empty inputs, just skip in this case
+    if torch.numel(masks) == 0:
+        return torch.zeros(*masks.shape[:-2], 4, device=masks.device)
+
+    # Normalize shape to CxHxW
+    shape = masks.shape
+    h, w = shape[-2:]
+    if len(shape) > 2:
+        masks = masks.flatten(0, -3)
+    else:
+        masks = masks.unsqueeze(0)
+
+    # Get top and bottom edges
+    in_height, _ = torch.max(masks, dim=-1)
+    in_height_coords = in_height * torch.arange(h, dtype=torch.int, device=in_height.device)[None, :]
+    bottom_edges, _ = torch.max(in_height_coords, dim=-1)
+    in_height_coords = in_height_coords + h * (~in_height)
+    in_height_coords = in_height_coords.type(torch.int)
+    top_edges, _ = torch.min(in_height_coords, dim=-1)
+
+    # Get left and right edges
+    in_width, _ = torch.max(masks, dim=-2)
+    in_width_coords = in_width * torch.arange(w, dtype=torch.int, device=in_width.device)[None, :]
+    right_edges, _ = torch.max(in_width_coords, dim=-1)
+    in_width_coords = in_width_coords + w * (~in_width)
+    in_width_coords = in_width_coords.type(torch.int)
+    left_edges, _ = torch.min(in_width_coords, dim=-1)
+
+    # If the mask is empty the right edge will be to the left of the left edge.
+    # Replace these boxes with [0, 0, 0, 0]
+    empty_filter = (right_edges < left_edges) | (bottom_edges < top_edges)
+    out = torch.stack([left_edges, top_edges, right_edges, bottom_edges], dim=-1)
+    out = out * (~empty_filter).unsqueeze(-1)
+
+    # Return to original shape
+    if len(shape) > 2:
+        out = out.reshape(*shape[:-2], 4)
+    else:
+        out = out[0]
+
+    return out
+
+
+@njit
+def _compute_rle_numba(mask):
+    val = mask[0]
+    counts = [int(x) for x in range(0)] if val == 0 else [0]
+    count = 0
+    for m in mask:
+        if val == m:
+            count += 1
+        else:
+            val = m
+            counts.append(count)
+            count = 1
+    counts.append(count)
+    return counts
+
+
+def _compute_rle_numpy(mask):
+    diffs = mask[1:] != mask[:-1]  # pairwise unequal (string safe)
+    indices = np.append(np.where(diffs), len(mask) - 1)  # must include last element position
+    # count needs to start with 0 if the mask begins with 1
+    counts = [] if mask[0] == 0 else [0]
+    # compute the actual RLE
+    counts += np.diff(np.append(-1, indices)).tolist()
+    return counts
+
+
+def mask_to_rle_pytorch(
+    tensor: torch.Tensor, rle_implementation: Literal["default", "numpy", "numba", "bioimage_cpp"] = "default"
+) -> List[Dict[str, Any]]:
+    """Calculates the runlength encoding of binary input masks.
+
+    This replaces the function in
+    https://github.com/facebookresearch/segment-anything/segment_anything/util/amg.py
+    with a version that computes the RLE purely on the CPU.
+    This does not lead to any performance deficits when running on the GPU, but it speeds the computation
+    up significantly compared to running this on an MPS device.
+    The RLE implementation is based on
+    https://stackoverflow.com/questions/1066758/find-length-of-sequences-of-identical-values-in-a-numpy-array-run-length-encodi
+    """
+    # Put in fortran order and flatten h, w
+    b, h, w = tensor.shape
+    tensor = tensor.permute(0, 2, 1).flatten(1)
+    tensor = tensor.detach().cpu().numpy()
+
+    if rle_implementation == "default":
+        rle_implementation = "bioimage_cpp"
+
+    if rle_implementation == "numba":
+        assert HAVE_NUMBA
+        rle_impl = _compute_rle_numba
+    elif rle_implementation == "numpy":
+        rle_impl = _compute_rle_numpy
+    elif rle_implementation == "bioimage_cpp":
+        rle_impl = _compute_rle_bioimage_cpp
+    else:
+        raise ValueError(
+            f"RLE implementation {rle_implementation} is not available. "
+            "Has to be one of 'numpy', 'numba' or 'bioimage_cpp'."
+        )
+
+    out = []
+    for mask in tensor:
+        counts = rle_impl(mask)
+        out.append({"size": [h, w], "counts": counts})
+    return out