Add UI for Interactive Attention Exploration

pyproject.toml

@@ -98,6 +98,10 @@ virchow2 = [
     "huggingface-hub>=0.27.1",
     "torch>=2.0.0",
 ]
+attentionui =[
+    "napari>=0.6.0",
+    "pyqt5>=5.15.11",
+]
 cobra = [
     "stamp[flash-attention]",
     "causal-conv1d @ git+https://github.com/KatherLab/causal-conv1d.git@52ec902314b9eda800162c73502a89f3572fc522",

src/stamp/__main__.py

@@ -236,6 +236,27 @@ def _run_cli(args: argparse.Namespace) -> None:
                 default_slide_mpp=config.heatmaps.default_slide_mpp,
             )
 
+        case "attentionui":
+            from stamp.heatmaps import attention_ui_
+
+            if config.attentionui is None:
+                raise ValueError("no attention configuration supplied")
+
+            _add_file_handle_(_logger, output_dir=config.attentionui.output_dir)
+            _logger.info(
+                "using the following configuration:\n"
+                f"{yaml.dump(config.attentionui.model_dump(mode='json'))}"
+            )
+            attention_ui_(
+                feature_dir=config.attentionui.feature_dir,
+                wsi_dir=config.attentionui.wsi_dir,
+                checkpoint_path=config.attentionui.checkpoint_path,
+                output_dir=config.attentionui.output_dir,
+                slide_paths=config.attentionui.slide_paths,
+                device=config.attentionui.device,
+                default_slide_mpp=config.attentionui.default_slide_mpp,
+            )
+
         case _:
             raise RuntimeError(
                 "unreachable: the argparser should only allow valid commands"
@@ -295,6 +316,9 @@ def main() -> None:
     )
     commands.add_parser("config", help="Print the loaded configuration")
     commands.add_parser("heatmaps", help="Generate heatmaps for a trained model")
+    commands.add_parser(
+        "attentionui", help="Provides an interactive UI for exploring attention maps"
+    )
 
     args = parser.parse_args()
 

src/stamp/config.py

@@ -1,7 +1,7 @@
 from pydantic import BaseModel, ConfigDict
 
 from stamp.encoding.config import PatientEncodingConfig, SlideEncodingConfig
-from stamp.heatmaps.config import HeatmapConfig
+from stamp.heatmaps.config import HeatmapConfig, AttentionUIConfig
 from stamp.modeling.config import (
     AdvancedConfig,
     CrossvalConfig,
@@ -24,6 +24,7 @@ class StampConfig(BaseModel):
     statistics: StatsConfig | None = None
 
     heatmaps: HeatmapConfig | None = None
+    attentionui: AttentionUIConfig | None = None
 
     slide_encoding: SlideEncodingConfig | None = None
 

src/stamp/config.yaml

@@ -215,6 +215,25 @@ heatmaps:
   #bottomk: 5
 
 
+attentionui:
+  output_dir: "/path/to/save/files/to"
+
+  # Directory the extracted features are saved in.
+  feature_dir: "/path/your/extracted/features/are/stored/in"
+
+  wsi_dir: "/path/containing/whole/slide/images/to/extract/features/from"
+
+  # Path of the model to generate the attention maps with.
+  checkpoint_path: "/path/to/model.ckpt"
+
+  # Slides to generate the attention maps for.
+  # The slide paths have to be specified relative to `wsi_dir`.
+  # If not specified, stamp will allow processing for all slides in `wsi_dir`.
+  #slide_paths:
+  #- slide1.svs
+  #- slide2.mrxs
+
+
 slide_encoding:
   # Encoder to use for slide encoding. Possible options are "cobra",
   # "eagle", "titan", "gigapath", "chief", "prism", "madeleine".

src/stamp/heatmaps/__init__.py

@@ -314,3 +314,55 @@ def heatmaps_(
 
         fig.savefig(slide_output_dir / f"overview-{h5_path.stem}.png")
         plt.close(fig)
+
+
+def attention_ui_(
+    *,
+    feature_dir: Path,
+    wsi_dir: Path,
+    checkpoint_path: Path,
+    output_dir: Path,
+    slide_paths: Iterable[Path] | None,
+    device: DeviceLikeType,
+    default_slide_mpp: SlideMPP | None,
+) -> None:
+    try:
+        from stamp.heatmaps.attention_ui import show_attention_ui
+    except ImportError as e:
+        raise ImportError(
+            "Attention UI dependencies not installed. "
+            "Please reinstall stamp using `pip install 'stamp[attentionui]'`"
+        ) from e
+
+    with torch.no_grad():
+        # Collect slides to generate attention maps for
+        if slide_paths is not None:
+            wsis_to_process_all = (wsi_dir / slide for slide in slide_paths)
+        else:
+            wsis_to_process_all = (
+                p for ext in supported_extensions for p in wsi_dir.glob(f"**/*{ext}")
+            )
+
+        # Check of a corresponding feature file exists
+        wsis_to_process = []
+        for wsi_path in wsis_to_process_all:
+            h5_path = feature_dir / wsi_path.with_suffix(".h5").name
+
+            if not h5_path.exists():
+                _logger.info(
+                    f"could not find matching h5 file at {h5_path}. Skipping..."
+                )
+                continue
+
+            wsis_to_process.append(str(wsi_path))
+
+        # Launch the UI
+        show_attention_ui(
+            feature_dir,
+            wsis_to_process,
+            checkpoint_path,
+            output_dir,
+            slide_paths,
+            device,
+            default_slide_mpp,
+        )