🆕 Simplified SAM segmentation in TIAViz

requirements/requirements.txt

@@ -35,6 +35,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/engines/test_semantic_segmentor.py

@@ -157,7 +157,6 @@ def _test_qupath_output_patch(output: Path) -> None:
     assert "Polygon" in geometry_types
 
     # When class_dict is None, types are assigned as 0, 1, ...
-    assert 0 in class_values
     assert 1 in class_values
 
     # Basic sanity check
@@ -1071,7 +1070,6 @@ def _test_store_output_patch(output: Path) -> None:
             if "type" in probs:
                 annotation_types.add(probs.pop("type"))
     # When class_dict is none, types are assigned as 0, 1, ...
-    assert 0 in annotation_types
     assert 1 in annotation_types
 
     assert annotations_properties is not None

tests/models/test_arch_grandqc.py

@@ -105,7 +105,7 @@ def test_grandqc_with_semantic_segmentor(
     assert Path(output[sample_image]).exists()
 
     store = SQLiteStore.open(output[sample_image])
-    assert len(store) == 4
+    assert len(store) == 3
 
     unique_types = set()
 
@@ -116,7 +116,7 @@ def test_grandqc_with_semantic_segmentor(
             tissue_area_px += annotation.geometry.area
     assert 2950000 < tissue_area_px < 2960000
 
-    assert unique_types == {"background", "tissue"}
+    assert unique_types == {"tissue"}
     store.close()
 
 

tests/models/test_arch_sam.py

@@ -0,0 +1,61 @@
+"""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
+    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))