🆕 Define DeepFeatureExtractor (#963)

# 🚀 Summary

This PR introduces a new **`DeepFeatureExtractor` engine** to the TIAToolbox framework, enabling extraction of intermediate CNN feature representations from whole slide images (WSIs) or image patches. These features can be used for downstream tasks such as clustering, visualization, or training other models. The update also includes:

- A **command-line interface (CLI)** for the new engine.
- Extended **CLI utilities** for flexible input/output configurations.
- Comprehensive **unit tests** covering patch-based and WSI-based workflows, multi-GPU support, and CLI functionality.
- Integration with TIAToolbox’s model registry and CLI ecosystem.

---

## ✨ Key Features

### **New Engine: `DeepFeatureExtractor`**
- Extracts intermediate CNN features from WSIs or patches.
- Outputs feature embeddings and spatial coordinates in **Zarr** or **dict** format.
- Implements **memory-aware caching** for large-scale WSI processing.
- Compatible with:
  - TIAToolbox pretrained models.
  - Torchvision CNN backbones (e.g., ResNet, DenseNet, MobileNet).
  - **All timm architectures via `timm.list_models()`**, including HuggingFace-hosted models.
- Supports both **patch-mode** and **WSI-mode** workflows.

### **CLI Integration**
- Adds `deep-feature-extractor` command to TIAToolbox CLI.
- Supports options for:
  - Input/output paths and file types.
  - Model selection (`resnet18`, `efficientnet_b0`, timm-based backbones, etc.).
  - Patch extraction parameters (`patch_input_shape`, `stride_shape`, `input_resolutions`).
  - Batch size, device selection, memory threshold, overwrite behavior.
- Flexible JSON-based CLI options for resolutions and class mappings.

### **Extended CLI Utilities**
- New reusable options:
  - `--input-resolutions`, `--output-resolutions` (JSON list of dicts).
  - `--patch-input-shape`, `--stride-shape`, `--scale-factor`.
  - `--class-dict` for mapping class indices to names.
  - `--overwrite` and `--output-file` for fine-grained control.

### **Unit Tests**
- **Engine Tests**:
  - Patch-based and WSI-based feature extraction.
  - Validation of Zarr outputs (features and coordinates).
  - Multi-GPU functionality.
- **Model Compatibility**:
  - Tests with `CNNBackbone` and `TimmBackbone` models.
- **CLI Tests**:
  - Single-file and parameterized runs.
  - Validation of JSON parsing for CLI options.

### **Codebase Integration**
- Registers `DeepFeatureExtractor` in `tiatoolbox.models` and engine registry.
- Adds CLI command in `tiatoolbox.cli.__init__.py`.
- Updates architecture utilities to support timm-based backbones and HuggingFace models.
- Introduces dictionaries for Torch and timm backbones (`torch_cnn_backbone_dict`, `timm_arch_dict`).

Author: shaneahmed

tests/engines/test_engine_abc.py

@@ -313,6 +313,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")