Simplified SAM segmentation in TIAViz

.pre-commit-config.yaml

@@ -2,6 +2,8 @@ ci:
   autofix_prs: true
   autoupdate_commit_msg: ':technologist: pre-commit autoupdate'
   autoupdate_schedule: 'monthly'
+default_language_version:
+  python: python3.12
 repos:
   - repo: local
     hooks:

requirements/requirements.txt

@@ -33,6 +33,7 @@ timm>=1.0.3
 torch>=2.5.0
 torchvision>=0.15.0
 tqdm>=4.64.1
+transformers>=4.51.1
 umap-learn>=0.5.3
 wsidicom>=0.18.0
 zarr>=2.13.3, <3.0.0

tests/models/test_arch_sam.py

@@ -0,0 +1,63 @@
+"""Unit test package for SAM."""
+
+from collections.abc import Callable
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+
+from tiatoolbox.models.architecture.sam import SAM
+from tiatoolbox.utils import env_detection as toolbox_env
+from tiatoolbox.utils import imread
+from tiatoolbox.utils.misc import select_device
+
+ON_GPU = toolbox_env.has_gpu()
+
+# Test pretrained Model =============================
+
+
+def test_functional_sam(remote_sample: Callable) -> None:
+    """Test for SAM."""
+    # convert to pathlib Path to prevent wsireader complaint
+    tile_path = Path(remote_sample("patch-extraction-vf"))
+    img = imread(tile_path)
+
+    # test creation
+
+    model = SAM(device=select_device(on_gpu=ON_GPU))
+
+    # create image patch and prompts
+    patch = img[63:191, 750:878, :]
+
+    points = np.array([[[64, 64]]])
+    boxes = np.array([[[64, 64, 128, 128]]])
+
+    # test preproc
+    tensor = torch.from_numpy(img)
+    patch = np.expand_dims(model.preproc(tensor), axis=0)
+    patch = model.preproc(patch)
+
+    # test inference
+
+    mask_output, score_output = model.infer_batch(
+        model, patch, points, device=select_device(on_gpu=ON_GPU)
+    )
+
+    assert mask_output is not None, "Output should not be None"
+    assert len(mask_output) > 0, "Output should have at least one element"
+    assert len(score_output) > 0, "Output should have at least one element"
+
+    mask_output, score_output = model.infer_batch(
+        model, patch, box_coords=boxes, device=select_device(on_gpu=ON_GPU)
+    )
+
+    assert len(mask_output) > 0, "Output should have at least one element"
+    assert len(score_output) > 0, "Output should have at least one element"
+
+    # test error when no prompts provided
+    with pytest.raises(
+        ValueError,
+        match=r"At least one of point_coords or box_coords must be provided.",
+    ):
+        _ = model.infer_batch(model, patch, device=select_device(on_gpu=ON_GPU))

tests/models/test_prompt_segmentor.py

@@ -0,0 +1,53 @@
+"""Unit test package for prompt segmentor."""
+
+from pathlib import Path
+
+import cv2
+import numpy as np
+
+from tiatoolbox.annotation.storage import SQLiteStore
+from tiatoolbox.models.architecture.sam import SAM
+from tiatoolbox.models.engine.prompt_segmentor import PromptSegmentor
+
+
+def test_prompt_segmentor(track_tmp_path: Path) -> None:
+    """Test for Prompt Segmentor."""
+    # create dummy image patch 256x256x3 with small circle in middle
+    img = np.zeros((256, 256, 3), dtype=np.uint8)
+    cv2.circle(img, (128, 128), 50, (255, 255, 255), -1)
+
+    expected_area = np.pi * (50**2)
+
+    # prompt with pt in center of circle
+    points = np.array([[[128, 128]]], np.uint32)
+    boxes = None
+
+    # instantiate prompt segmentor with SAM model
+    sam_model = SAM()
+    prompt_segmentor = PromptSegmentor(model=sam_model)
+
+    # run prediction
+    output_paths = prompt_segmentor.run(
+        images=[img],
+        point_coords=points,
+        box_coords=boxes,
+        save_dir=track_tmp_path / "sam_test_output",
+        device="cpu",
+    )
+
+    assert isinstance(output_paths, list), "Output should be a list of paths"
+    assert len(output_paths) == 1, "Output list should contain one path"
+
+    # load the saved annotation db
+    store = SQLiteStore(output_paths[0])
+    ann = next(iter(store.values()))
+    # area should be close to expected area
+    assert abs(ann.geometry.area - expected_area) < expected_area * 0.1, (
+        "should segment circle area correctly"
+    )
+
+    # check with model=None
+    prompt_segmentor_default = PromptSegmentor(model=None)
+    assert isinstance(prompt_segmentor_default.model, SAM), (
+        "Default model should be SAM"
+    )

tests/test_app_bokeh.py

@@ -558,6 +558,52 @@ def test_hovernet_on_box(doc: Document, data_path: pytest.TempPathFactory) -> No
     assert len(main.UI["type_column"].children) == 1
 
 
+def test_sam_segment(doc: Document, data_path: pytest.TempPathFactory) -> None:
+    """Test running SAM on points and a box."""
+    slide_select = doc.get_model_by_name("slide_select0")
+    slide_select.value = [data_path["slide2"].name]
+    run_button = doc.get_model_by_name("to_model0")
+    assert len(main.UI["color_column"].children) == 0
+    slide_select.value = [data_path["slide1"].name]
+    # set up a box selection
+    main.UI["box_source"].data = {
+        "x": [1200],
+        "y": [-2000],
+        "width": [400],
+        "height": [400],
+    }
+
+    # select SAM model and run it on box
+    model_select = doc.get_model_by_name("model_drop0")
+    model_select.value = "SAM"
+
+    click = ButtonClick(run_button)
+    run_button._trigger_event(click)
+    assert len(main.UI["color_column"].children) > 0
+
+    # test save functionality
+    save_button = doc.get_model_by_name("save_button0")
+    click = ButtonClick(save_button)
+    save_button._trigger_event(click)
+    saved_path = (
+        data_path["base_path"]
+        / "overlays"
+        / (data_path["slide1"].stem + "_saved_anns.db")
+    )
+    assert saved_path.exists()
+
+    # load an overlay with different types
+    cprop_select = doc.get_model_by_name("cprop0")
+    cprop_select.value = ["prob"]
+    layer_drop = doc.get_model_by_name("layer_drop0")
+    click = MenuItemClick(layer_drop, str(data_path["dat_anns"]))
+    layer_drop._trigger_event(click)
+    assert main.UI["vstate"].types == ["annotation"]
+    # check the per-type ui controls have been updated
+    assert len(main.UI["color_column"].children) == 1
+    assert len(main.UI["type_column"].children) == 1
+
+
 def test_alpha_sliders(doc: Document) -> None:
     """Test sliders for adjusting slide and overlay alpha."""
     slide_alpha = doc.get_model_by_name("slide_alpha0")

tiatoolbox/models/__init__.py

@@ -8,6 +8,7 @@
 from .architecture.mapde import MapDe
 from .architecture.micronet import MicroNet
 from .architecture.nuclick import NuClick
+from .architecture.sam import SAM
 from .architecture.sccnn import SCCNN
 from .engine.multi_task_segmentor import MultiTaskSegmentor
 from .engine.nucleus_instance_segmentor import NucleusInstanceSegmentor
@@ -17,6 +18,7 @@
     PatchPredictor,
     WSIPatchDataset,
 )
+from .engine.prompt_segmentor import PromptSegmentor
 from .engine.semantic_segmentor import (
     DeepFeatureExtractor,
     IOSegmentorConfig,
@@ -25,6 +27,7 @@
 )
 
 __all__ = [
+    "SAM",
     "SCCNN",
     "HoVerNet",
     "HoVerNetPlus",
@@ -35,5 +38,6 @@
     "NuClick",
     "NucleusInstanceSegmentor",
     "PatchPredictor",
+    "PromptSegmentor",
     "SemanticSegmentor",
 ]