micro-sam/examples/annotator_2d.py at 4b7ac97a60b3ef23b659736cd0be189405b97f20 · computational-cell-analytics/micro-sam · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
import os

import imageio.v3 as imageio
from micro_sam.util import get_cache_directory
from micro_sam.sam_annotator import annotator
from micro_sam.sample_data import fetch_hela_2d_example_data, fetch_livecell_example_data, fetch_wholeslide_example_data


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(use_finetuned_model):
    """Run the 2d annotator for an example image from the LIVEcell dataset.

    See https://doi.org/10.1038/s41592-021-01249-6 for details on the data.
    """
    example_data = fetch_livecell_example_data(DATA_CACHE)
    image = imageio.imread(example_data)

    if use_finetuned_model:
        embedding_path = os.path.join(EMBEDDING_CACHE, "embeddings-livecell-vit_b_lm.zarr")
        model_type = "vit_b_lm"
    else:
        embedding_path = os.path.join(EMBEDDING_CACHE, "embeddings-livecell.zarr")
        model_type = "vit_b"

    annotator(
        image, embedding_path=embedding_path, model_type=model_type,
        precompute_autoseg_state=True,
    )


def hela_2d_annotator(use_finetuned_model):
    """Run the 2d annotator for an example image from the Cell Tracking Challenge (HeLa 2d) dataset.
    """
    example_data = fetch_hela_2d_example_data(DATA_CACHE)
    image = imageio.imread(example_data)

    if use_finetuned_model:
        embedding_path = os.path.join(EMBEDDING_CACHE, "embeddings-hela2d-vit_b_lm.zarr")
        model_type = "vit_b_lm"
    else:
        embedding_path = os.path.join(EMBEDDING_CACHE, "embeddings-hela2d.zarr")
        model_type = "vit_b"

    annotator(image, embedding_path=embedding_path, model_type=model_type)


def wholeslide_annotator(use_finetuned_model):
    """Run the 2d annotator with tiling for an example whole-slide image from the
    NeurIPS Cell Segmentation challenge.

    See https://neurips22-cellseg.grand-challenge.org/ for details on the data.
    """
    example_data = fetch_wholeslide_example_data(DATA_CACHE)
    image = imageio.imread(example_data)

    if use_finetuned_model:
        embedding_path = os.path.join(EMBEDDING_CACHE, "whole-slide-embeddings-vit_b_lm.zarr")
        model_type = "vit_b_lm"
    else:
        embedding_path = os.path.join(EMBEDDING_CACHE, "whole-slide-embeddings.zarr")
        model_type = "vit_b"

    annotator(
        image, embedding_path=embedding_path, tile_shape=(1024, 1024), halo=(256, 256), model_type=model_type
    )


def main():
    # Whether to use the fine-tuned SAM model for light microscopy data.
    use_finetuned_model = True

    # 2d annotator for livecell data
    livecell_annotator(use_finetuned_model)

    # 2d annotator for cell tracking challenge hela data
    # hela_2d_annotator(use_finetuned_model)

    # 2d annotator for a whole slide image
    # wholeslide_annotator(use_finetuned_model)


# The corresponding CLI call for hela_2d_annotator:
# (replace with cache directory on your machine)
# $ micro_sam.annotator -i /home/pape/.cache/micro_sam/sample_data/hela-2d-image.png -e /home/pape/.cache/micro_sam/embeddings/embeddings-hela2d.zarr  # noqa
if __name__ == "__main__":
    main()