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

New release

Author: constantinpape

.github/workflows/test.yaml

@@ -20,9 +20,7 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-latest, windows-latest, macos-latest]
-        # 3.12 currently not supported due to issues with nifty.
-        # python-version: ["3.11", "3.12"]
-        python-version: ["3.11"]
+        python-version: ["3.11", "3.12"]
 
     steps:
       - name: Checkout

.gitignore

@@ -195,6 +195,8 @@ iterative_prompting_results/
 *.tif
 *.zip
 *MACOSX
+hela_ctc
+clf-test-data
 
 # Related to i2k workshop folders.
 data/

environment.yaml

@@ -5,7 +5,7 @@ dependencies:
     - nifty >=1.2.3
     - imagecodecs
     - magicgui
-    - napari >=0.5.0,<0.6.0
+    - napari
     - natsort
     - pip
     - pooch

examples/README.md

@@ -1,10 +1,14 @@
 # Examples
 
 Examples for using the `micro_sam` annotation tools:
-- `annotator_2d.py`: run the interactive 2d annotation tool.
-- `annotator_3d.py`: run the interactive 3d annotation tool.
-- `annotator_tracking.py`: run the interactive tracking annotation tool.
-- `image_series_annotator.py`: run the annotation tool for a series of images.
+- `annotator_2d.py`: Run the interactive 2d annotation tool.
+- `annotator_3d.py`: Run the interactive 3d annotation tool.
+- `annotator_tracking.py`: Run the interactive tracking annotation tool.
+- `image_series_annotator.py`: Run the annotation tool for a series of images.
+
+And python scripts for automatic segmentation and tracking:
+- `automatic_segmentation.py`: Run automatic segmentation on 2d images.
+- `automatic_tracking.py`: Run automatic tracking on 2d timeseries images.
 
 And examples for using the `micro_sam` automatic segmentation feature:
 - `quick_start.py`: run the automatic segmentation feature of `micro_sam` on an example 2d image.

examples/automatic_tracking.py

@@ -0,0 +1,56 @@
+import os
+
+from elf.io import open_file
+
+from micro_sam.util import get_cache_directory
+from micro_sam.sample_data import fetch_tracking_example_data
+from micro_sam.automatic_segmentation import automatic_tracking, get_predictor_and_segmenter
+
+
+DATA_CACHE = os.path.join(get_cache_directory(), "sample_data")
+EMBEDDING_CACHE = os.path.join(get_cache_directory(), "embeddings")
+os.makedirs(EMBEDDING_CACHE, exist_ok=True)
+
+
+def example_automatic_tracking(use_finetuned_model):
+    """Run automatic tracking for data from the cell tracking challenge.
+    """
+    # Download the example tracking data.
+    example_data = fetch_tracking_example_data(DATA_CACHE)
+
+    # Load the example data (load the sequence of tif files as timeseries)
+    with open_file(example_data, mode="r") as f:
+        timeseries = f["*.tif"]
+
+    if use_finetuned_model:
+        embedding_path = os.path.join(EMBEDDING_CACHE, "embeddings-ctc-vit_b_lm.zarr")
+        model_type = "vit_b_lm"
+    else:
+        embedding_path = os.path.join(EMBEDDING_CACHE, "embeddings-ctc.zarr")
+        model_type = "vit_h"
+
+    predictor, segmenter = get_predictor_and_segmenter(model_type=model_type, amg=False)
+
+    masks_tracked, _ = automatic_tracking(
+        predictor=predictor,
+        segmenter=segmenter,
+        input_path=timeseries[:],
+        output_path="./hela_ctc",
+        embedding_path=embedding_path,
+    )
+
+    import napari
+    v = napari.Viewer()
+    v.add_image(timeseries)
+    v.add_labels(masks_tracked)
+    napari.run()
+
+
+def main():
+    # Whether to use the fine-tuned SAM model.
+    use_finetuned_model = True
+    example_automatic_tracking(use_finetuned_model)
+
+
+if __name__ == "__main__":
+    main()

examples/object_classifier.py

@@ -0,0 +1,276 @@
+import os
+
+import imageio.v3 as imageio
+import numpy as np
+
+from micro_sam.util import get_cache_directory
+from micro_sam.sample_data import fetch_livecell_example_data, fetch_wholeslide_example_data, fetch_3d_example_data
+
+from elf.io import open_file
+
+
+DATA_CACHE = os.path.join(get_cache_directory(), "sample_data")
+EMBEDDING_CACHE = os.path.join(get_cache_directory(), "embeddings")
+os.makedirs(EMBEDDING_CACHE, exist_ok=True)
+
+
+def livecell_annotator():
+    from micro_sam.sam_annotator.object_classifier import object_classifier
+
+    example_data = fetch_livecell_example_data(DATA_CACHE)
+    image = imageio.imread(example_data)
+
+    embedding_path = os.path.join(EMBEDDING_CACHE, "embeddings-livecell-vit_b_lm.zarr")
+    model_type = "vit_b_lm"
+
+    # This is the vit-b-lm segmentation
+    segmentation = imageio.imread("./clf-test-data/livecell-test-seg.tif")
+
+    object_classifier(image, segmentation, embedding_path=embedding_path, model_type=model_type)
+
+
+def wholeslide_annotator():
+    from micro_sam.sam_annotator.object_classifier import object_classifier
+
+    example_data = fetch_wholeslide_example_data(DATA_CACHE)
+    image = imageio.imread(example_data)
+
+    embedding_path = os.path.join(EMBEDDING_CACHE, "whole-slide-embeddings-vit_b_lm.zarr")
+    model_type = "vit_b_lm"
+
+    segmentation = imageio.imread("./clf-test-data/whole-slide-seg.tif")
+    object_classifier(
+        image, segmentation, embedding_path=embedding_path, model_type=model_type,
+        tile_shape=(1024, 1024), halo=(256, 256),
+    )
+
+
+def lucchi_annotator():
+    from micro_sam.sam_annotator.object_classifier import object_classifier
+
+    example_data = fetch_3d_example_data(DATA_CACHE)
+    with open_file(example_data) as f:
+        raw = f["*.png"][:]
+
+    embedding_path = os.path.join(EMBEDDING_CACHE, "embeddings-lucchi-vit_b_em_organelles.zarr")
+
+    model_type = "vit_b_lm"
+    segmentation = imageio.imread("./clf-test-data/lucchi-test-segmentation.tif")
+
+    object_classifier(raw, segmentation, embedding_path=embedding_path, model_type=model_type)
+
+
+def tiled_3d_annotator():
+    from micro_sam.sam_annotator.object_classifier import object_classifier
+    from skimage.data import cells3d
+
+    data = cells3d()[30:34, 1]
+    embed_path = "./clf-test-data/emebds-3d-tiled.zarr"
+
+    model_type = "vit_b_lm"
+    segmentation = imageio.imread("./clf-test-data/tiled-3d-segmentation.tif")
+
+    object_classifier(
+        data, segmentation, embedding_path=embed_path, model_type=model_type,
+        tile_shape=(128, 128), halo=(32, 32)
+    )
+
+
+def _get_livecell_data():
+    example_data = fetch_livecell_example_data(DATA_CACHE)
+    image = imageio.imread(example_data)
+
+    embedding_path = os.path.join(EMBEDDING_CACHE, "embeddings-livecell-vit_b_lm.zarr")
+
+    # This is the vit-b-lm segmentation and a test annotaiton.
+    segmentation = imageio.imread("./clf-test-data/livecell-test-seg.tif")
+    annotations = imageio.imread("./clf-test-data/livecell-test-annotations.tif")
+
+    model_type = "vit_b_lm"
+
+    return image, segmentation, annotations, model_type, embedding_path, None, None
+
+
+def _get_wholeslide_data():
+    example_data = fetch_wholeslide_example_data(DATA_CACHE)
+    image = imageio.imread(example_data)
+
+    embedding_path = os.path.join(EMBEDDING_CACHE, "whole-slide-embeddings-vit_b_lm.zarr")
+
+    # This is the vit-b-lm segmentation and a test annotaiton.
+    segmentation = imageio.imread("./clf-test-data/whole-slide-seg.tif")
+    annotations = imageio.imread("./clf-test-data/wholeslide-annotations.tif")
+
+    model_type = "vit_b_lm"
+
+    return image, segmentation, annotations, model_type, embedding_path, (1024, 1024), (256, 256)
+
+
+def _get_lucchi_data():
+    example_data = fetch_3d_example_data(DATA_CACHE)
+    with open_file(example_data) as f:
+        raw = f["*.png"][:]
+
+    embedding_path = os.path.join(EMBEDDING_CACHE, "embeddings-lucchi-vit_b_em_organelles.zarr")
+
+    segmentation = imageio.imread("./clf-test-data/lucchi-test-segmentation.tif")
+    annotations = imageio.imread("./clf-test-data/lucchi-annotations.tif")
+
+    model_type = "vit_b_em_organelles"
+
+    return raw, segmentation, annotations, model_type, embedding_path, None, None
+
+
+def _get_3d_tiled_data():
+    from skimage.data import cells3d
+
+    data = cells3d()[30:34, 1]
+    embed_path = "./clf-test-data/emebds-3d-tiled.zarr"
+    model_type = "vit_b_lm"
+
+    segmentation = imageio.imread("./clf-test-data/tiled-3d-segmentation.tif")
+    annotations = imageio.imread("./clf-test-data/tiled-3d-annotations.tif")
+
+    return data, segmentation, annotations, model_type, embed_path, (128, 128), (32, 32)
+
+
+def annotator_devel():
+    from micro_sam import object_classification as core_clf
+    from micro_sam.sam_annotator import object_classifier as clf
+    from micro_sam.util import precompute_image_embeddings, get_sam_model
+
+    # image, segmentation, annotations, model_type, embedding_path, tile_shape, halo = _get_livecell_data()
+    # image, segmentation, annotations, model_type, embedding_path, tile_shape, halo = _get_wholeslide_data()
+    # image, segmentation, annotations, model_type, embedding_path, tile_shape, halo = _get_lucchi_data()
+    image, segmentation, annotations, model_type, embedding_path, tile_shape, halo = _get_3d_tiled_data()
+
+    # 1. Get the SAM model
+    predictor = get_sam_model(model_type)
+    # 2. Precompute the image embeddings.
+    image_embeddings = precompute_image_embeddings(
+        predictor, image, save_path=embedding_path, tile_shape=tile_shape, halo=halo
+    )
+    # 3. Get the segmentation ids and the extracted features for the segmentations.
+    seg_ids, features = core_clf.compute_object_features(image_embeddings, segmentation)
+    # 4. Points to the objects we would like to select for training RF.
+    labels = clf._accumulate_labels(segmentation, annotations)
+    # 5. Traint the RF model.
+    rf = clf._train_rf(features, labels, n_estimators=200, max_depth=10)
+    # 6. Run the trained RF prediction on new images.
+    object_prediction = rf.predict(features)
+    # 7. Map the predictions back to the instance segmentation.
+    prediction = core_clf.project_prediction_to_segmentation(segmentation, object_prediction, seg_ids)
+
+    import napari
+    v = napari.Viewer()
+    v.add_image(image)
+    v.add_labels(annotations)
+    v.add_labels(prediction)
+    napari.run()
+
+
+def create_3d_data_with_tiling():
+    from skimage.data import cells3d
+    from micro_sam.automatic_segmentation import automatic_instance_segmentation, get_predictor_and_segmenter
+
+    predictor, segmenter = get_predictor_and_segmenter(model_type="vit_b_lm", is_tiled=True)
+    data = cells3d()[30:34, 1]
+
+    embed_path = "./clf-test-data/emebds-3d-tiled.zarr"
+    seg = automatic_instance_segmentation(
+        predictor, segmenter, data, embedding_path=embed_path, ndim=3, tile_shape=(128, 128), halo=(32, 32)
+    )
+
+    import napari
+    v = napari.Viewer()
+    v.add_image(data)
+    v.add_labels(seg)
+    # For annotations.
+    v.add_labels(np.zeros_like(seg))
+    napari.run()
+
+
+def histopathology_annotator():
+    from torch_em.data.datasets.histopathology.lynsec import get_lynsec_paths
+    from micro_sam.sam_annotator import object_classifier as clf
+    from micro_sam.automatic_segmentation import automatic_instance_segmentation, get_predictor_and_segmenter
+
+    predictor, segmenter = get_predictor_and_segmenter(model_type="vit_b_histopathology")
+
+    image_paths, _ = get_lynsec_paths(path="./clf-test-data/nuclick", choice="ihc", download=True)
+    image_paths = image_paths[:10]
+
+    images, segmentations = [], []
+    embedding_paths = []
+
+    for i, image_path in enumerate(image_paths):
+        image = imageio.imread(image_path)
+        embedding_path = f"./clf-test-data/embeddings_nuclick_{i}.zarr"
+        seg_path = f"./clf-test-data/seg-nuclick_{i}.tif"
+
+        if os.path.exists(seg_path):
+            segmentation = imageio.imread(seg_path)
+        else:
+            segmentation = automatic_instance_segmentation(
+                predictor, segmenter, embedding_path=embedding_path, input_path=image, ndim=2,
+            )
+            imageio.imwrite(seg_path, segmentation, compression="zlib")
+
+        images.append(image)
+        segmentations.append(segmentation)
+        embedding_paths.append(embedding_path)
+
+    clf.image_series_object_classifier(
+        images, segmentations, output_folder="./clf-test-data/histo-results",
+        embedding_paths=embedding_paths, model_type="vit_b_histopathology", ndim=2,
+    )
+
+
+def batch_prediction():
+    import napari
+    from torch_em.data.datasets.histopathology.lynsec import get_lynsec_paths
+    from micro_sam.automatic_segmentation import automatic_instance_segmentation, get_predictor_and_segmenter
+    from micro_sam.object_classification import run_prediction_with_object_classifier
+    from tqdm import tqdm
+
+    predictor, segmenter = get_predictor_and_segmenter(model_type="vit_b_histopathology")
+
+    image_paths, _ = get_lynsec_paths(path="./clf-test-data/nuclick", choice="ihc", download=True)
+    # Test the batch prediction on the next 5 images.
+    image_paths = image_paths[10:12]
+
+    images, segmentations = [], []
+    # Prepare images and segmentations
+    for image_path in tqdm(image_paths, desc="Segment images"):
+        image = imageio.imread(image_path)
+        segmentation = automatic_instance_segmentation(predictor, segmenter, input_path=image, ndim=2, verbose=False)
+        images.append(image)
+        segmentations.append(segmentation)
+
+    rf_path = "clf-test-data/histo-results/rf.joblib"
+    print("Start object clf")
+    predictions = run_prediction_with_object_classifier(images, segmentations, predictor, rf_path, ndim=2)
+
+    for im, seg, pred in zip(images, segmentations, predictions):
+        v = napari.Viewer()
+        v.add_image(im)
+        v.add_labels(seg)
+        v.add_labels(pred)
+        napari.run()
+
+
+def main():
+    # create_3d_data_with_tiling()
+
+    # livecell_annotator()
+    # wholeslide_annotator()
+    # lucchi_annotator()
+    # tiled_3d_annotator()
+    histopathology_annotator()
+    # batch_prediction()
+
+    # annotator_devel()
+
+
+if __name__ == "__main__":
+    main()