Merge branch 'dev-define-engines-abc' into dev-define-nucleus-instance-segmentor

Author: shaneahmed

tests/engines/test_engine_abc.py

@@ -345,6 +345,7 @@ def test_engine_run() -> NoReturn:
             on_gpu=False,
         )
 
+    eng = TestEngineABC(model="alexnet-kather100k")
     with pytest.raises(
         ValueError,
         match=r".*The shape of the numpy array should be NHWC*",

tests/engines/test_feature_extractor.py

@@ -0,0 +1,241 @@
+"""Test for feature extractor."""
+
+import shutil
+from collections.abc import Callable
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+import zarr
+from click.testing import CliRunner
+
+from tiatoolbox import cli
+from tiatoolbox.models import IOPatchPredictorConfig
+from tiatoolbox.models.architecture.vanilla import CNNBackbone, TimmBackbone
+from tiatoolbox.models.engine.deep_feature_extractor import DeepFeatureExtractor
+from tiatoolbox.utils import env_detection as toolbox_env
+from tiatoolbox.wsicore.wsireader import WSIReader
+
+ON_GPU = not toolbox_env.running_on_ci() and toolbox_env.has_gpu()
+
+# -------------------------------------------------------------------------------------
+# Engine
+# -------------------------------------------------------------------------------------
+
+device = "cuda" if toolbox_env.has_gpu() else "cpu"
+
+
+def test_feature_extractor_patches(
+    remote_sample: Callable,
+) -> None:
+    """Tests DeepFeatureExtractor on image patches."""
+    extractor = DeepFeatureExtractor(
+        model="fcn-tissue_mask", batch_size=32, verbose=False, device=device
+    )
+
+    sample_image = remote_sample("thumbnail-1k-1k")
+
+    inputs = [sample_image, sample_image]
+
+    assert not extractor.patch_mode
+    output = extractor.run(
+        images=inputs,
+        return_probabilities=True,
+        return_labels=False,
+        device=device,
+        patch_mode=True,
+    )
+
+    assert 0.48 < np.mean(output["features"][:]) < 0.52
+
+    with pytest.raises(
+        ValueError,
+        match=r".*output_type: `annotationstore` is not supported "
+        r"for `DeepFeatureExtractor` engine",
+    ):
+        _ = extractor.run(
+            images=inputs,
+            return_probabilities=True,
+            return_labels=False,
+            device=device,
+            patch_mode=True,
+            output_type="annotationstore",
+        )
+
+
+def test_feature_extractor_wsi(remote_sample: Callable, track_tmp_path: Path) -> None:
+    """Test feature extraction with DeepFeatureExtractor engine."""
+    save_dir = track_tmp_path / "output"
+    # # convert to pathlib Path to prevent wsireader complaint
+    mini_wsi_svs = Path(remote_sample("wsi2_4k_4k_svs"))
+
+    # * test providing pretrained from torch vs pretrained_model.yaml
+    shutil.rmtree(save_dir, ignore_errors=True)  # default output dir test
+
+    extractor = DeepFeatureExtractor(batch_size=1, model="fcn-tissue_mask")
+    output = extractor.run(
+        images=[mini_wsi_svs],
+        return_probabilities=False,
+        return_labels=False,
+        device=device,
+        patch_mode=False,
+        save_dir=track_tmp_path / "wsi_out_check",
+        batch_size=1,
+        output_type="zarr",
+        memory_threshold=1,
+    )
+
+    output_ = zarr.open(output[mini_wsi_svs], mode="r")
+    assert len(output_["coordinates"].shape) == 2
+    assert len(output_["features"].shape) == 3
+
+
+@pytest.mark.parametrize(
+    "model",
+    [
+        CNNBackbone("resnet18"),
+        TimmBackbone("efficientnet_b0", pretrained=True),
+        "resnet18",
+        "efficientnet_b0",
+    ],
+)
+def test_full_inference(
+    remote_sample: Callable, track_tmp_path: Path, model: Callable
+) -> None:
+    """Test full inference with CNNBackbone and TimmBackbone models."""
+    save_dir = track_tmp_path / "output"
+    # pre-emptive clean up
+    shutil.rmtree(save_dir, ignore_errors=True)  # default output dir test
+
+    mini_wsi_svs = Path(remote_sample("wsi4_1k_1k_svs"))
+
+    ioconfig = IOPatchPredictorConfig(
+        input_resolutions=[
+            {"units": "mpp", "resolution": 0.25},
+        ],
+        patch_input_shape=[512, 512],
+        stride_shape=[256, 256],
+    )
+
+    extractor = DeepFeatureExtractor(batch_size=4, model=model)
+    output = extractor.run(
+        images=[mini_wsi_svs],
+        device=device,
+        save_dir=track_tmp_path / "wsi_out_check",
+        batch_size=4,
+        output_type="zarr",
+        ioconfig=ioconfig,
+        patch_mode=False,
+    )
+
+    output_ = zarr.open(output[mini_wsi_svs], mode="r")
+
+    positions = output_["coordinates"]
+    features = output_["features"]
+
+    reader = WSIReader.open(mini_wsi_svs)
+    patches = [
+        reader.read_bounds(
+            positions[patch_idx],
+            resolution=0.25,
+            units="mpp",
+            pad_constant_values=255,
+            coord_space="resolution",
+        )
+        for patch_idx in range(4)
+    ]
+    patches = np.array(patches)
+    patches = torch.from_numpy(patches)  # NHWC
+    patches = patches.permute(0, 3, 1, 2).contiguous()  # NCHW
+    patches = patches.to(device).type(torch.float32)
+    model = extractor.model
+    # Inference mode
+    model.eval()
+    with torch.inference_mode():
+        _features = model(patches).cpu().numpy()
+    # ! must maintain same batch size and likely same ordering
+    # ! else the output values will not exactly be the same (still < 1.0e-4
+    # ! of epsilon though)
+    assert np.mean(np.abs(features[:4] - _features)) < 1.0e-1
+
+
+@pytest.mark.skipif(
+    toolbox_env.running_on_ci() or not ON_GPU,
+    reason="Local test on machine with GPU.",
+)
+def test_multi_gpu_feature_extraction(
+    remote_sample: Callable, track_tmp_path: Path
+) -> None:
+    """Local functionality test for feature extraction using multiple GPUs."""
+    save_dir = track_tmp_path / "output"
+    mini_wsi_svs = Path(remote_sample("wsi4_1k_1k_svs"))
+    shutil.rmtree(save_dir, ignore_errors=True)
+
+    wsi_ioconfig = IOPatchPredictorConfig(
+        input_resolutions=[{"units": "mpp", "resolution": 0.5}],
+        patch_input_shape=[224, 224],
+        stride_shape=[224, 224],
+    )
+
+    extractor = DeepFeatureExtractor(
+        model="UNI",
+        batch_size=32,
+        num_workers=4,
+    )
+
+    output = extractor.run(
+        [mini_wsi_svs],
+        patch_mode=False,
+        device=device,
+        ioconfig=wsi_ioconfig,
+        save_dir=save_dir,
+        auto_get_mask=True,
+        output_type="zarr",
+    )
+    output_ = zarr.open(output[mini_wsi_svs], mode="r")
+
+    positions = output_["coordinates"]
+    features = output_["features"]
+    assert len(positions.shape) == 2
+    assert len(features.shape) == 2
+
+
+# -------------------------------------------------------------------------------------
+# Command Line Interface
+# -------------------------------------------------------------------------------------
+
+
+def test_cli_model_single_file(sample_svs: Path, track_tmp_path: Path) -> None:
+    """Test for feature extractor CLI single file."""
+    runner = CliRunner()
+
+    models_wsi_result = runner.invoke(
+        cli.main,
+        [
+            "deep-feature-extractor",
+            "--img-input",
+            str(sample_svs),
+            "--model",
+            "resnet18",
+            "--patch-mode",
+            "False",
+            "--output-path",
+            str(track_tmp_path / "output"),
+            "--patch-input-shape",
+            "224",
+            "224",
+            "--input-resolutions",
+            '[{"units": "mpp", "resolution": 0.25}]',
+        ],
+    )
+
+    assert models_wsi_result.exit_code == 0
+    assert (track_tmp_path / "output" / (sample_svs.stem + ".zarr")).exists()
+
+    output = zarr.open(
+        track_tmp_path / "output" / (sample_svs.stem + ".zarr"), mode="r"
+    )
+
+    # Output shape should be # of patches x feature size
+    assert output["features"].shape == (255, 512)

tests/engines/test_patch_predictor.py

@@ -268,17 +268,14 @@ def test_wsi_predictor_api(
     kwargs = {
         "patch_input_shape": patch_size,
         "stride_shape": patch_size,
-        "input_resolutions": [{"units": "baseline", "resolution": 1.0}],
         "save_dir": save_dir,
     }
-    # ! add this test back once the read at `baseline` is fixed
-    # sanity check, both output should be the same with same resolution read args
-    # remove previously generated data
 
     _kwargs = copy.deepcopy(kwargs)
     # test reading of multiple whole-slide images
     output = predictor.run(
         images=[mini_wsi_svs, str(mini_wsi_jpg)],
+        input_resolutions=[{"units": "baseline", "resolution": 1.0}],
         masks=[mini_wsi_msk, mini_wsi_msk],
         patch_mode=False,
         return_probabilities=True,
@@ -308,21 +305,7 @@ def test_patch_predictor_kather100k_output(
     pretrained_info = {
         "alexnet-kather100k": [1.0, 0.9999735355377197],
         "resnet18-kather100k": [1.0, 0.9999911785125732],
-        "resnet34-kather100k": [1.0, 0.9979840517044067],
-        "resnet50-kather100k": [1.0, 0.9999986886978149],
-        "resnet101-kather100k": [1.0, 0.9999932050704956],
-        "resnext50_32x4d-kather100k": [1.0, 0.9910059571266174],
-        "resnext101_32x8d-kather100k": [1.0, 0.9999971389770508],
-        "wide_resnet50_2-kather100k": [1.0, 0.9953408241271973],
-        "wide_resnet101_2-kather100k": [1.0, 0.9999831914901733],
-        "densenet121-kather100k": [1.0, 1.0],
-        "densenet161-kather100k": [1.0, 0.9999959468841553],
-        "densenet169-kather100k": [1.0, 0.9999934434890747],
-        "densenet201-kather100k": [1.0, 0.9999983310699463],
-        "mobilenet_v2-kather100k": [0.9999998807907104, 0.9999126195907593],
-        "mobilenet_v3_large-kather100k": [0.9999996423721313, 0.9999878406524658],
         "mobilenet_v3_small-kather100k": [0.9999998807907104, 0.9999997615814209],
-        "googlenet-kather100k": [1.0, 0.9999639987945557],
     }
     for model, expected_prob in pretrained_info.items():
         _test_predictor_output(

tests/engines/test_semantic_segmentor.py

@@ -498,7 +498,7 @@ def test_wsi_segmentor_annotationstore(
 
 
 def test_cli_model_single_file(sample_svs: Path, track_tmp_path: Path) -> None:
-    """Test for models CLI single file."""
+    """Test semantic segmentor CLI single file."""
     runner = CliRunner()
     models_wsi_result = runner.invoke(
         cli.main,

tests/models/test_models_abc.py

@@ -174,5 +174,5 @@ def test_get_pretrained_model_not_str() -> None:
 
 def test_get_pretrained_model_not_in_info() -> None:
     """Test ValueError is raised if input is not in info."""
-    with pytest.raises(ValueError, match=r"Pretrained model `alexnet` does not exist."):
-        _ = get_pretrained_model("alexnet")
+    with pytest.raises(ValueError, match=r"Pretrained model `random` does not exist."):
+        _ = get_pretrained_model("random")