Merge pull request #852 from computational-cell-analytics/dev

Prepare release 1.2.1

Author: constantinpape

doc/bioimageio/histopathology_v1.md

@@ -0,0 +1,16 @@
+# Segment Anything for Histopathology
+
+This is a [Segment Anything]https://segment-anything.com/) model that was specialized for histopathology with [micro_sam](https://github.com/computational-cell-analytics/micro-sam).
+This model uses a %s vision transformer as image encoder.
+
+Segment Anything is a model for interactive and automatic instance segmentation.
+We improve it for histopathology by finetuning on a large and diverse microscopy dataset.
+It should perform well for nucleus segmentation in histopathology datasets.
+
+See [the dataset overview](https://github.com/computational-cell-analytics/micro-sam/blob/master/doc/datasets/histopathology_v%i.md) for further informations on the training data and the [micro_sam documentation](https://computational-cell-analytics.github.io/micro-sam/micro_sam.html) for details on how to use the model for interactive and automatic segmentation.
+
+## Validation
+
+The easiest way to validate the model is to visually check the segmentation quality for your data.
+If you have annotations you can use for validation you can also quantitative validation, see [here for details](https://computational-cell-analytics.github.io/micro-sam/micro_sam.html#9-how-can-i-evaluate-a-model-i-have-finetuned).
+Please note that the required quality for segmentation always depends on the analysis task you want to solve.

finetuning/livecell_finetuning.py

@@ -56,6 +56,19 @@ def finetune_livecell(args):
     train_loader, val_loader = get_dataloaders(patch_shape=patch_shape, data_path=args.input_path)
     scheduler_kwargs = {"mode": "min", "factor": 0.9, "patience": 10, "verbose": True}
 
+    # NOTE: memory req. for all vit_b models (compared on A100 80GB)
+    # vit_b
+    # freeze_encoder: ~ 33.89 GB
+    # QLoRA: ~48.54 GB
+    # LoRA: ~48.62 GB
+    # FFT: ~49.56 GB
+
+    # vit_h
+    # freeze_encoder: ~36.05 GB
+    # QLoRA: ~ 65.68 GB
+    # LoRA: ~ 67.14 GB
+    # FFT: ~72.34 GB
+
     # Run training.
     sam_training.train_sam(
         name=checkpoint_name,
@@ -72,7 +85,7 @@ def finetune_livecell(args):
         save_root=args.save_root,
         scheduler_kwargs=scheduler_kwargs,
         save_every_kth_epoch=args.save_every_kth_epoch,
-        peft_kwargs={"rank": args.lora_rank} if args.lora_rank is not None else None,
+        peft_kwargs={"rank": args.lora_rank, "quantize": True} if args.lora_rank is not None else None,
     )
 
     if args.export_path is not None:
@@ -87,7 +100,7 @@ def finetune_livecell(args):
 def main():
     parser = argparse.ArgumentParser(description="Finetune Segment Anything for the LIVECell dataset.")
     parser.add_argument(
-        "--input_path", "-i", default="/scratch/projects/nim00007/sam/data/livecell/",
+        "--input_path", "-i", default="/mnt/vast-nhr/projects/cidas/cca/data/livecell/",
         help="The filepath to the LIVECell data. If the data does not exist yet it will be downloaded."
     )
     parser.add_argument(

micro_sam/__init__.py

@@ -9,6 +9,7 @@
 .. include:: ../doc/contributing.md
 .. include:: ../doc/band.md
 """
+
 import os
 
 from .__version__ import __version__

micro_sam/__version__.py

@@ -1 +1 @@
-__version__ = "1.2.0"
+__version__ = "1.2.1"

micro_sam/_vendored.py

@@ -1,13 +1,12 @@
-"""
-Functions from other third party libraries.
+"""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
+from typing import Any, Dict, List, Literal
 
 import numpy as np
 
@@ -109,7 +108,9 @@ def _compute_rle_numpy(mask):
     return counts
 
 
-def mask_to_rle_pytorch(tensor: torch.Tensor, rle_implementation: str = "default") -> List[Dict[str, Any]]:
+def mask_to_rle_pytorch(
+    tensor: torch.Tensor, rle_implementation: Literal["default", "numpy", "numba", "nifty"] = "default"
+) -> List[Dict[str, Any]]:
     """Calculates the runlength encoding of binary input masks.
 
     This replaces the function in

micro_sam/automatic_segmentation.py

@@ -120,7 +120,11 @@ def automatic_instance_segmentation(
             verbose=verbose,
         )
 
-        segmenter.initialize(image=image_data, image_embeddings=image_embeddings)
+        # If we run AIS with tiling then we use the same tile shape for the watershed postprocessing.
+        if isinstance(segmenter, InstanceSegmentationWithDecoder) and tile_shape is not None:
+            generate_kwargs.update({"tile_shape": tile_shape, "halo": halo})
+
+        segmenter.initialize(image=image_data, image_embeddings=image_embeddings, verbose=verbose)
         masks = segmenter.generate(**generate_kwargs)
 
         if len(masks) == 0:  # instance segmentation can have no masks, hence we just save empty labels

micro_sam/bioimageio/bioengine_export.py

@@ -3,6 +3,7 @@
 from typing import Optional, Union
 
 import torch
+
 from segment_anything.utils.onnx import SamOnnxModel
 
 try:
@@ -67,7 +68,7 @@ def export_image_encoder(
     model_type: str,
     output_root: Union[str, os.PathLike],
     export_name: Optional[str] = None,
-    checkpoint_path: Optional[str] = None,
+    checkpoint_path: Optional[Union[str, os.PathLike]] = None,
 ) -> None:
     """Export SAM image encoder to torchscript.
 
@@ -103,15 +104,16 @@ def export_image_encoder(
 
 
 def export_onnx_model(
-    model_type,
-    output_root,
-    opset: int,
+    model_type: str,
+    output_root: Union[str, os.PathLike],
+    opset: int = 17,
     export_name: Optional[str] = None,
     checkpoint_path: Optional[Union[str, os.PathLike]] = None,
     return_single_mask: bool = True,
     gelu_approximate: bool = False,
     use_stability_score: bool = False,
     return_extra_metrics: bool = False,
+    quantize_model: bool = False,
 ) -> None:
     """Export SAM prompt encoder and mask decoder to onnx.
 
@@ -122,14 +124,16 @@ def export_onnx_model(
     Args:
         model_type: The SAM model type.
         output_root: The output root directory where the exported model is saved.
-        opset: The ONNX opset version.
+        opset: The ONNX opset version. The recommended opset version is 17.
         export_name: The name of the exported model.
         checkpoint_path: Optional checkpoint for loading the SAM model.
         return_single_mask: Whether the mask decoder returns a single or multiple masks.
         gelu_approximate: Whether to use a GeLU approximation, in case the ONNX backend
             does not have an efficient GeLU implementation.
         use_stability_score: Whether to use the stability score instead of the predicted score.
         return_extra_metrics: Whether to return a larger set of metrics.
+        quantize_model: Whether to also export a quantized version of the model.
+            This only works for onnxruntime < 1.17.
     """
     if export_name is None:
         export_name = model_type
@@ -154,10 +158,7 @@ def export_onnx_model(
             if isinstance(m, torch.nn.GELU):
                 m.approximate = "tanh"
 
-    dynamic_axes = {
-        "point_coords": {1: "num_points"},
-        "point_labels": {1: "num_points"},
-    }
+    dynamic_axes = {"point_coords": {1: "num_points"}, "point_labels": {1: "num_points"}}
 
     embed_dim = sam.prompt_encoder.embed_dim
     embed_size = sam.prompt_encoder.image_embedding_size
@@ -202,14 +203,31 @@ def export_onnx_model(
         _ = ort_session.run(None, ort_inputs)
         print("Model has successfully been run with ONNXRuntime.")
 
+    # This requires onnxruntime < 1.17.
+    # See https://github.com/facebookresearch/segment-anything/issues/699#issuecomment-1984670808
+    if quantize_model:
+        assert onnxruntime_exists
+        from onnxruntime.quantization import QuantType
+        from onnxruntime.quantization.quantize import quantize_dynamic
+
+        quantized_path = os.path.join(weight_output_folder, "model_quantized.onnx")
+        quantize_dynamic(
+            model_input=weight_path,
+            model_output=quantized_path,
+            # optimize_model=True,
+            per_channel=False,
+            reduce_range=False,
+            weight_type=QuantType.QUInt8,
+        )
+
     config_output_path = os.path.join(output_folder, "config.pbtxt")
     with open(config_output_path, "w") as f:
         f.write(DECODER_CONFIG % name)
 
 
 def export_bioengine_model(
-    model_type,
-    output_root,
+    model_type: str,
+    output_root: Union[str, os.PathLike],
     opset: int,
     export_name: Optional[str] = None,
     checkpoint_path: Optional[Union[str, os.PathLike]] = None,

micro_sam/bioimageio/model_export.py

@@ -33,14 +33,13 @@
     "tags": ["segment-anything", "instance-segmentation"],
 }
 
+# Reference: https://github.com/bioimage-io/spec-bioimage-io/commit/39d343681d427ec93cf69eef7597d9eb9678deb1#diff-0bbdaa8196fa31f945afabcf04a4295ff098f1f24400ef9e59b0f684d411905eL269  # noqa
+# We had this parameter in bioimageio.spec. This has been removed. We just make a copy of the same parameter.
+ARBITRARY_SIZE = spec.ParameterizedSize(min=1, step=1)
+
+
+def _create_test_inputs_and_outputs(image, labels, model_type, checkpoint_path, tmp_dir):
 
-def _create_test_inputs_and_outputs(
-    image,
-    labels,
-    model_type,
-    checkpoint_path,
-    tmp_dir,
-):
     # For now we just generate a single box prompt here, but we could also generate more input prompts.
     generator = PointAndBoxPromptGenerator(
         n_positive_points=1,
@@ -59,10 +58,7 @@ def _create_test_inputs_and_outputs(
 
     # Generate logits from the two
     mask_prompts = np.stack(
-        [
-            _compute_logits_from_mask(labels == 1),
-            _compute_logits_from_mask(labels == 2),
-        ]
+        [_compute_logits_from_mask(labels == 1), _compute_logits_from_mask(labels == 2)]
     )[None]
 
     predictor = PredictorAdaptor(model_type=model_type)
@@ -104,11 +100,7 @@ def _create_test_inputs_and_outputs(
         "point_labels": point_label_path,
         "mask_prompts": mask_prompt_path,
     }
-    outputs = {
-        "mask": mask_path,
-        "score": score_path,
-        "embeddings": embed_path
-    }
+    outputs = {"mask": mask_path, "score": score_path, "embeddings": embed_path}
     return inputs, outputs
 
 
@@ -161,6 +153,7 @@ def _get_checkpoint(model_type, checkpoint_path, tmp_dir):
         return checkpoint_path, None
 
 
+# TODO: Update this with our latest yaml file updates.
 def _write_dependencies(dependency_file, require_mobile_sam):
     content = """name: sam
 channels:
@@ -215,7 +208,7 @@ def _check_model(model_description, input_paths, result_paths):
     image = xarray.DataArray(np.load(input_paths["image"]), dims=tuple("bcyx"))
     embeddings = xarray.DataArray(np.load(result_paths["embeddings"]), dims=tuple("bcyx"))
     box_prompts = xarray.DataArray(np.load(input_paths["box_prompts"]), dims=tuple("bic"))
-    point_prompts = xarray.DataArray(np.load(input_paths["point_prompts"]), dims=tuple("biic"))
+    point_prompts = xarray.DataArray(np.load(input_paths["point_prompts"]), dims=tuple("bhwc"))
     point_labels = xarray.DataArray(np.load(input_paths["point_labels"]), dims=tuple("bic"))
     mask_prompts = xarray.DataArray(np.load(input_paths["mask_prompts"]), dims=tuple("bicyx"))
 
@@ -303,8 +296,8 @@ def export_sam_model(
                     # NOTE: to support 1 and 3 channels we can add another preprocessing.
                     # Best solution: Have a pre-processing for this! (1C -> RGB)
                     spec.ChannelAxis(channel_names=[spec.Identifier(cname) for cname in "RGB"]),
-                    spec.SpaceInputAxis(id=spec.AxisId("y"), size=spec.ARBITRARY_SIZE),
-                    spec.SpaceInputAxis(id=spec.AxisId("x"), size=spec.ARBITRARY_SIZE),
+                    spec.SpaceInputAxis(id=spec.AxisId("y"), size=ARBITRARY_SIZE),
+                    spec.SpaceInputAxis(id=spec.AxisId("x"), size=ARBITRARY_SIZE),
                 ],
                 test_tensor=spec.FileDescr(source=input_paths["image"]),
                 data=spec.IntervalOrRatioDataDescr(type="uint8")
@@ -318,7 +311,7 @@ def export_sam_model(
                     spec.BatchAxis(size=1),
                     spec.IndexInputAxis(
                         id=spec.AxisId("object"),
-                        size=spec.ARBITRARY_SIZE
+                        size=ARBITRARY_SIZE
                     ),
                     spec.ChannelAxis(channel_names=[spec.Identifier(bname) for bname in "hwxy"]),
                 ],
@@ -334,11 +327,11 @@ def export_sam_model(
                     spec.BatchAxis(size=1),
                     spec.IndexInputAxis(
                         id=spec.AxisId("object"),
-                        size=spec.ARBITRARY_SIZE
+                        size=ARBITRARY_SIZE
                     ),
                     spec.IndexInputAxis(
                         id=spec.AxisId("point"),
-                        size=spec.ARBITRARY_SIZE
+                        size=ARBITRARY_SIZE
                     ),
                     spec.ChannelAxis(channel_names=[spec.Identifier(bname) for bname in "xy"]),
                 ],
@@ -354,11 +347,11 @@ def export_sam_model(
                     spec.BatchAxis(size=1),
                     spec.IndexInputAxis(
                         id=spec.AxisId("object"),
-                        size=spec.ARBITRARY_SIZE
+                        size=ARBITRARY_SIZE
                     ),
                     spec.IndexInputAxis(
                         id=spec.AxisId("point"),
-                        size=spec.ARBITRARY_SIZE
+                        size=ARBITRARY_SIZE
                     ),
                 ],
                 test_tensor=spec.FileDescr(source=input_paths["point_labels"]),
@@ -373,7 +366,7 @@ def export_sam_model(
                     spec.BatchAxis(size=1),
                     spec.IndexInputAxis(
                         id=spec.AxisId("object"),
-                        size=spec.ARBITRARY_SIZE
+                        size=ARBITRARY_SIZE
                     ),
                     spec.ChannelAxis(channel_names=["channel"]),
                     spec.SpaceInputAxis(id=spec.AxisId("y"), size=256),

micro_sam/bioimageio/predictor_adaptor.py

@@ -55,6 +55,9 @@ def forward(
             embeddings: precomputed image embeddings B x 256 x 64 x 64
 
         Returns:
+            The segmentation masks.
+            The scores for prediction quality.
+            The computed image embeddings.
         """
         batch_size = image.shape[0]
         if batch_size != 1:

micro_sam/evaluation/experiments.py

@@ -11,9 +11,7 @@
 
 
 def full_experiment_settings(
-    use_boxes: bool = False,
-    positive_range: Optional[List[int]] = None,
-    negative_range: Optional[List[int]] = None,
+    use_boxes: bool = False, positive_range: Optional[List[int]] = None, negative_range: Optional[List[int]] = None,
 ) -> ExperimentSettings:
     """The full experiment settings.