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,
+ )

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,
+ )

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,
+ )

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,
+ )

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,
+ )

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,
+ )

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,
+ )

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,
+ )