diff --git a/README-commands.md b/README-commands.md index 46645b32..3c7c227a 100644 --- a/README-commands.md +++ b/README-commands.md @@ -247,6 +247,19 @@ The following models are currently supported, | MADELEINE | MADELEINE_SLIDE | CONCH1_5 | ๐Ÿ”’ gated | | CHIEF | CHIEF_SLIDE | CTRANSPATH | local ckpt | +`model_kwargs={...}` forwards extra constructor arguments to patch encoders, and +`slide_model_kwargs={...}` forwards them to slide encoders. `TITAN_SLIDE` applies +its GPU OOM patch by default (`patch_oom=true`), which also pins the validated +TITAN revision. Disable it only when testing upstream behavior: + +```bash +aggregate_slide_features \ + patch_features_h5_path=patch_features.h5 \ + output_h5_path=slide_features.h5 \ + slide_model_type=TITAN_SLIDE \ + slide_model_kwargs={patch_oom:false} +``` + OpenCLIP is used by default, with the default model being [QuiltNet-B-16-PMB](https://huggingface.co/wisdomik/QuiltNet-B-16-PMB). Use the `model_type` parameter to specify a different model. To use H-Optimus-0, for example, @@ -536,4 +549,3 @@ Each input file `.` produces `output_dir/.tiff`. Pass | `downscale_by` | `1` | Integer downsample factor (e.g. `2` converts a 40ร— slide to 20ร—). | | `num_workers` | `1` | Parallel workers for batch mode (`0` = all CPUs). | | `bigtiff` | `false` | Write BigTIFF format (required for files > ~4 GB). | - diff --git a/README.md b/README.md index e011008f..22fd3282 100644 --- a/README.md +++ b/README.md @@ -280,6 +280,47 @@ For running Mussel at scale on a compute cluster, see pipeline that wraps the Mussel CLI tools and handles job scheduling, parallelism, and output management across large slide cohorts. +## Development + +### Docker/Apptainer Containers with Flash Attention + +Mussel supports building Docker containers with **flash-attn 2.0** for accelerated attention in the CONCH1.5 patch encoder and the Prov-GigaPath slide encoder. Flash attention provides ~30-50% speedup on patch encoding (~20% overall TITAN pipeline improvement). + +**Building the flash-attn container:** + +```bash +# Build Docker image with flash-attn backend +make docker-build BACKEND=fastattn +# Or manually: +docker build --build-arg BACKEND=fastattn -t mussel:fastattn . + +# Convert to Apptainer SIF for HPC deployment +make sif +# Or manually: +apptainer build --force mussel-fastattn.sif docker-daemon://mussel:fastattn +``` + +**Key details:** +- The `[fastattn]` extra in `pyproject.toml` installs: + - torch 2.11.0+cu126 (CUDA 12.6) + - flash-attn 2.6.3 (custom manylinux_2_28 wheels for Rocky 8 compatibility) + - xformers 0.0.35 +- Flash attention is required for **CONCH1.5** (patch encoding) and **Prov-GigaPath** (slide encoding); both require CUDA compute capability โ‰ฅ 8.0 (A100, H100, etc.) +- For V100 (compute 7.0) or CPU, the code automatically falls back to PyTorch SDPA +- TITAN slide encoder uses `SDPBackend.EFFICIENT_ATTENTION` (Phase 1 optimization); flash-attn accelerates CONCH1.5 patch encoding and GigaPath slide encoding + +**Using with mussel-nf:** + +Copy the SIF to your mussel-nf repo and use the `apptainer_fastattn` profile: + +```bash +cp mussel-fastattn.sif /path/to/mussel-nf/ +cd /path/to/mussel-nf +nextflow run main.nf -profile cluster,slurm,apptainer_fastattn ... +``` + +The profile automatically uses the flash-attn container for `FEATURIZE` tasks. + ## License This code is made available under the GPLv3 License and is available for non-commercial academic purposes. Forked from CLAM, ยฉ [Mahmood Lab](http://www.mahmoodlab.org). diff --git a/mussel/cli/aggregate_slide_features.py b/mussel/cli/aggregate_slide_features.py index bd206ecd..409ffc25 100644 --- a/mussel/cli/aggregate_slide_features.py +++ b/mussel/cli/aggregate_slide_features.py @@ -1,11 +1,11 @@ import logging -from dataclasses import dataclass -from typing import List, Optional +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional import hydra from hydra.conf import HelpConf, HydraConf from hydra.core.config_store import ConfigStore -from omegaconf import MISSING +from omegaconf import MISSING, DictConfig from mussel.models import ModelType from mussel.utils import aggregate_slide_features, resolve_aggregation_method @@ -24,6 +24,9 @@ class AggregateSlideFeaturesConfig: aggregation_method (str): Method for aggregation: 'identity' (no aggregation), 'mean', 'max', or 'model'. slide_model_type (Optional[ModelType]): Type of slide encoder model (when aggregation_method="model"). slide_model_path (Optional[str]): Path to slide encoder model weights. + slide_model_kwargs (Dict[str, Any]): Extra keyword arguments forwarded to the slide model constructor. + TITAN_SLIDE applies its OOM patch by default via patch_oom=True; pass + slide_model_kwargs={patch_oom:false} to disable that patch and revision pin. use_gpu (bool): Whether to use GPU for computation. gpu_device_id (Optional[int]): Specific GPU device ID to use. gpu_device_ids (Optional[List[int]]): List of GPU device IDs for multi-GPU inference. @@ -37,12 +40,17 @@ class AggregateSlideFeaturesConfig: aggregation_method: str = "identity" slide_model_type: Optional[ModelType] = None slide_model_path: Optional[str] = None + slide_model_kwargs: Dict[str, Any] = field(default_factory=dict) use_gpu: bool = True gpu_device_id: Optional[int] = None gpu_device_ids: Optional[List[int]] = None ssl_verify: bool = True # Whether to verify SSL certificates for remote operations embedding_precision: str = "float32" + def __post_init__(self): + if isinstance(self.slide_model_kwargs, DictConfig): + self.slide_model_kwargs = dict(self.slide_model_kwargs) + desc_doc = """== ${hydra.help.app_name} == @@ -104,6 +112,13 @@ def main(cfg: AggregateSlideFeaturesConfig): patch_features_h5_path=patch_features.h5 \\ output_h5_path=slide_features.h5 \\ slide_model_type=TITAN_SLIDE + + # Disable TITAN's default OOM monkey-patch/revision pin, e.g. to test upstream + aggregate_slide_features \\ + patch_features_h5_path=patch_features.h5 \\ + output_h5_path=slide_features.h5 \\ + slide_model_type=TITAN_SLIDE \\ + slide_model_kwargs={patch_oom:false} """ logger.info("Starting slide feature aggregation") logger.info(f"Input: {cfg.patch_features_h5_path}") @@ -131,6 +146,7 @@ def main(cfg: AggregateSlideFeaturesConfig): gpu_device_id=cfg.gpu_device_id, gpu_device_ids=cfg.gpu_device_ids, embedding_precision=cfg.embedding_precision, + slide_model_kwargs=cfg.slide_model_kwargs, ) logger.info(f"Slide features saved to: {cfg.output_h5_path}") diff --git a/mussel/cli/extract_features.py b/mussel/cli/extract_features.py index 44d822d6..8a7e9fd3 100755 --- a/mussel/cli/extract_features.py +++ b/mussel/cli/extract_features.py @@ -1,9 +1,9 @@ import logging import os import ssl -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path -from typing import Dict, List, Optional +from typing import Any, Dict, List, Optional import h5py import hydra @@ -11,17 +11,20 @@ import torch from hydra.conf import HelpConf, HydraConf from hydra.core.config_store import ConfigStore -from omegaconf import MISSING +from omegaconf import MISSING, DictConfig from torch.utils.data import DataLoader from tqdm import tqdm from mussel.datasets import WholeSlideImageTileCoordDataset from mussel.models import ModelType, get_model_factory -from mussel.utils import (aggregate_slide_features_batch, - ensure_directory_exists, - extract_patch_features_batch, - get_slide_ids_from_paths, resolve_remote_paths, - save_features) +from mussel.utils import ( + aggregate_slide_features_batch, + ensure_directory_exists, + extract_patch_features_batch, + get_slide_ids_from_paths, + resolve_remote_paths, + save_features, +) from mussel.utils.feature_extract import _numpy_to_torch from mussel.utils.file import save_hdf5, save_torch_tensor from mussel.utils.ml import collate_features @@ -76,6 +79,10 @@ class ExtractFeaturesConfig: aggregation_method (str): Aggregation method: identity (single-step), mean/max/model (two-step). slide_model_type (Optional[ModelType]): Type of slide encoder model (when aggregation_method="model"). slide_model_path (Optional[str]): Path to slide encoder model weights. + model_kwargs (Dict[str, Any]): Extra keyword arguments forwarded to the patch model constructor. + slide_model_kwargs (Dict[str, Any]): Extra keyword arguments forwarded to the slide model constructor. + TITAN_SLIDE applies its OOM patch by default via patch_oom=True; pass + slide_model_kwargs={patch_oom:false} to disable that patch and revision pin. ssl_verify (bool): Whether to verify SSL certificates when downloading models or accessing remote resources (default: True). Processing Parameters: @@ -115,6 +122,8 @@ class ExtractFeaturesConfig: aggregation_method: str = "identity" slide_model_type: Optional[ModelType] = None slide_model_path: Optional[str] = None + model_kwargs: Dict[str, Any] = field(default_factory=dict) + slide_model_kwargs: Dict[str, Any] = field(default_factory=dict) ssl_verify: bool = True # Whether to verify SSL certificates for remote operations # Processing parameters batch_size: int = 64 @@ -125,6 +134,12 @@ class ExtractFeaturesConfig: is_test_run: bool = False embedding_precision: str = "float32" + def __post_init__(self): + if isinstance(self.model_kwargs, DictConfig): + self.model_kwargs = dict(self.model_kwargs) + if isinstance(self.slide_model_kwargs, DictConfig): + self.slide_model_kwargs = dict(self.slide_model_kwargs) + desc_doc = """== ${hydra.help.app_name} == @@ -223,6 +238,8 @@ def _main_single(cfg: ExtractFeaturesConfig): aggregation_method=cfg.aggregation_method, slide_model_type=cfg.slide_model_type, slide_model_path=cfg.slide_model_path, + model_kwargs=cfg.model_kwargs, + slide_model_kwargs=cfg.slide_model_kwargs, embedding_precision=cfg.embedding_precision, ) @@ -292,6 +309,7 @@ def _main_batch(cfg: ExtractFeaturesConfig): pin_memory=True, is_test_run=cfg.is_test_run, embedding_precision=cfg.embedding_precision, + model_kwargs=cfg.model_kwargs, ) # Aggregate to slide level using batch processing @@ -310,6 +328,7 @@ def _main_batch(cfg: ExtractFeaturesConfig): gpu_device_id=cfg.gpu_device_id, gpu_device_ids=cfg.gpu_device_ids, slide_batch_size=cfg.slide_batch_size, + slide_model_kwargs=cfg.slide_model_kwargs, ) # Defense-in-depth: verify at least one output file was created. @@ -340,6 +359,7 @@ def _main_batch(cfg: ExtractFeaturesConfig): pin_memory=True, is_test_run=cfg.is_test_run, embedding_precision=cfg.embedding_precision, + model_kwargs=cfg.model_kwargs, ) # Save as PT format for consistency diff --git a/mussel/cli/tessellate_extract_features.py b/mussel/cli/tessellate_extract_features.py index b179dae6..5f4c734f 100644 --- a/mussel/cli/tessellate_extract_features.py +++ b/mussel/cli/tessellate_extract_features.py @@ -153,6 +153,12 @@ class TessellateExtractFeaturesConfig: * GIGAPATH_SLIDE requires GIGAPATH patch encoder * TITAN_SLIDE requires CONCH1_5 patch encoder - The required patch encoder is automatically paired and run as needed + model_kwargs (Dict[str, Any]): Extra keyword arguments forwarded to patch model constructors. + slide_model_kwargs (Dict[str, Any]): Extra keyword arguments forwarded to slide model constructors. + TITAN_SLIDE applies GPU float16/expand-based OOM monkey-patches and a + validated revision pin by default (patch_oom=True). Pass + slide_model_kwargs={patch_oom:false} to disable that patch when testing + upstream TITAN behavior. model_dir (Optional[str]): Directory containing pre-downloaded models. - When specified, the system looks for model subdirectories named after each model type - For example: /mnt/batch_models/GIGAPATH_SLIDE, /mnt/batch_models/CONCH1_5 @@ -216,6 +222,8 @@ class TessellateExtractFeaturesConfig: prefilter_model_type: Optional[ModelType] = None # No prefilter by default prefilter_model_path: Optional[str] = None # Path to prefilter model file model_type: Any = None # Can be ModelType or List[ModelType] + model_kwargs: Dict[str, Any] = field(default_factory=dict) + slide_model_kwargs: Dict[str, Any] = field(default_factory=dict) model_dir: Optional[str] = None # Directory containing pre-downloaded models pre_download_models: bool = False # Whether to pre-download models to model_dir # Single mode visualization @@ -272,6 +280,10 @@ def __post_init__(self): # Convert DictConfig to regular dict for model_batch_sizes if isinstance(self.model_batch_sizes, DictConfig): self.model_batch_sizes = dict(self.model_batch_sizes) + if isinstance(self.model_kwargs, DictConfig): + self.model_kwargs = dict(self.model_kwargs) + if isinstance(self.slide_model_kwargs, DictConfig): + self.slide_model_kwargs = dict(self.slide_model_kwargs) # Convert DictConfig to regular dict for model_embedding_precision if isinstance(self.model_embedding_precision, DictConfig): @@ -910,6 +922,7 @@ def _main_batch( num_workers=cfg.num_workers, pin_memory=True, is_test_run=False, + model_kwargs=cfg.model_kwargs, ) except Exception as e: logger.error(f"Error during batch feature extraction: {e}") @@ -1079,6 +1092,7 @@ def _main_batch( slide_batch_size=cfg.slide_batch_size, max_slide_patches=cfg.max_slide_patches, embedding_precision=_resolve_precision(cfg, cfg.slide_model_type), + slide_model_kwargs=cfg.slide_model_kwargs, ) except Exception as e: logger.error(f"Error during batch aggregation: {e}") @@ -1163,6 +1177,7 @@ def _main_batch( num_workers=cfg.num_workers, pin_memory=True, is_test_run=False, + model_kwargs=cfg.model_kwargs, ) # Save PT files and coords-only H5 diff --git a/mussel/cli/tessellate_extract_features_common.py b/mussel/cli/tessellate_extract_features_common.py index b24db995..129bf08c 100644 --- a/mussel/cli/tessellate_extract_features_common.py +++ b/mussel/cli/tessellate_extract_features_common.py @@ -155,6 +155,7 @@ def _tessellate_and_filter( batch_size=cfg.batch_size, num_workers=cfg.num_workers, gpu_device_ids=cfg.gpu_device_ids, + model_kwargs=getattr(cfg, "model_kwargs", {}), ) logger.info(f"Filtering features: {slide_path}") @@ -317,6 +318,7 @@ def process_slide_tessellation_and_filtering( aggregation_method="identity", slide_model_type=None, slide_model_path=None, + model_kwargs=getattr(cfg, "model_kwargs", {}), ) return { "intermediate_h5_path": intermediate_h5_path, @@ -344,6 +346,8 @@ def process_slide_tessellation_and_filtering( aggregation_method=cfg.aggregation_method, slide_model_type=getattr(cfg, "slide_model_type", None), slide_model_path=slide_model_path, + model_kwargs=getattr(cfg, "model_kwargs", {}), + slide_model_kwargs=getattr(cfg, "slide_model_kwargs", {}), ) return { "final_coords": final_coords, diff --git a/mussel/models/conch.py b/mussel/models/conch.py index 39ad13ae..bca5e4a0 100644 --- a/mussel/models/conch.py +++ b/mussel/models/conch.py @@ -1,6 +1,9 @@ """CONCH v1.5 patch encoder and TITAN slide encoder from MahmoodLab.""" +import contextlib import logging +import math +import types from pathlib import Path from typing import Callable, List @@ -105,6 +108,141 @@ def save(self, save_path: str): ) +# --------------------------------------------------------------------------- +# TITAN monkey-patch helpers +# --------------------------------------------------------------------------- + +# Commit hash of MahmoodLab/TITAN on HuggingFace that the monkey-patches were +# written and verified against. Both get_alibi() and forward_features() are +# overridden; if the upstream implementation changes (new signature or logic) +# the patches must be re-validated before bumping this pin. +_TITAN_PINNED_REVISION = "dac6773d9961cfc75503440676ff157a2c6e8d2e" + +def _get_slopes(n: int) -> list: + """ALiBi attention slopes for ``n`` heads (from TITAN/vision_transformer.py).""" + if math.log2(n) == int(math.log2(n)): + p = 2 ** (-2 ** -(math.log2(n) - 3)) + return [p * (p ** i) for i in range(n)] + nearest = 2 ** math.floor(math.log2(n)) + base = _get_slopes(nearest) + if nearest == n: + return base + extra = _get_slopes(2 * nearest)[0::2][:n - nearest] + return base + extra + +def _titan_get_alibi_gpu_float16(self, w: int, h: int, bg_mask=None): + """GPU float16 replacement for VisionTransformer.get_alibi(). + + The original implementation creates O(Nยฒ) numpy float64 arrays on CPU + (17 GB for N=33k), causing SLURM OOM on large IMPACT resection specimens. + This version uses torch.cdist in float16 on the model's GPU, reducing peak + memory from ~82 GB CPU โ†’ ~26 GB GPU for N=33k. + + torch.cdist is fused in CUDA and does not create the intermediate (N, N, 2) + array that numpy broadcasting would require. + """ + device = next(self.parameters()).device + dtype = torch.float16 + + x_coords = torch.arange(w, device=device, dtype=dtype) + y_coords = torch.arange(h, device=device, dtype=dtype) + grid_x, grid_y = torch.meshgrid(x_coords, y_coords, indexing='ij') + + if bg_mask is not None: + # bg_mask shape is (1, H, W) bool โ€” squeeze to (H, W), index into 2D grid + mask_2d = bg_mask.to(device).squeeze(0).bool() # (H, W) or (W*H,) depending on caller + if mask_2d.dim() == 1: + # Already flattened (w*h,) + mask_flat = mask_2d + pts_x = grid_x.ravel()[mask_flat] + pts_y = grid_y.ravel()[mask_flat] + else: + # 2D mask (W, H) โ€” use as index into grid + pts_x = grid_x[mask_2d] + pts_y = grid_y[mask_2d] + else: + pts_x = grid_x.ravel() + pts_y = grid_y.ravel() + + points = torch.stack([pts_x, pts_y], dim=1) # (N, 2) float16 + + # Pairwise Euclidean distances โ€” fused CUDA, no (N, N, 2) intermediate + dists = torch.cdist(points.float(), points.float(), p=2).to(dtype) # (N, N) + + slopes = torch.tensor( + _get_slopes(self.num_heads), dtype=dtype, device=device + ).view(self.num_heads, 1, 1) + + n_patches = dists.shape[0] + bias_matrix = -dists.unsqueeze(0) * slopes # (H, N, N) + embed_len = n_patches + 1 + all_bias = torch.zeros( + 1, self.num_heads, embed_len, embed_len, dtype=dtype, device=device + ) + all_bias[:, :, 1:, 1:] = bias_matrix + return all_bias + + +def _titan_forward_features_efficient(self, x, coords=None, mask=None, bg_mask=None): + """Memory-efficient replacement for VisionTransformer.forward_features(). + + The original uses `attn_bias.repeat(B, 1, 1, 1)` which creates a full copy + of the (1, H, N, N) bias tensor โ€” 22 GB for N=30k on A100. This replacement + uses `expand()` (a zero-copy view) and avoids redundant dtype/device casts + when the bias is already in the correct format (float16 on GPU from the + get_alibi monkey-patch). + """ + B, nc, w, h = x.shape + x = x.flatten(2, 3).transpose(1, 2) + + if self.pos_encode_type == 'alibi': + if w * h == 36 and B != 1: + if not self.local_alibi_status: + self.prepare_tensor(x, 'local', 'alibi') + attn_bias = self.local_alibi + elif w * h == 196 and B != 1: + if not self.global_alibi_status: + self.prepare_tensor(x, 'global', 'alibi') + attn_bias = self.global_alibi + else: + # Calls our monkey-patched get_alibi (returns float16 GPU tensor) + attn_bias = self.get_alibi(w, h, bg_mask) if B == 1 else self.get_alibi(w, h) + # Use expand instead of repeat: zero-copy view for the B dimension + attn_bias = attn_bias.expand(x.shape[0], -1, -1, -1) + # Only cast if dtype/device differ (avoids copy when already float16 on GPU) + if attn_bias.dtype != x.dtype or attn_bias.device != x.device: + attn_bias = attn_bias.to(dtype=x.dtype, device=x.device) + else: + attn_bias = None + + if self.masked_im_modeling: + assert mask is not None + x = self.patch_embed(x) + x = self.mask_model(x, mask) + else: + x = self.patch_embed(x) + + x = self._pos_embed(x, coords, w, h) + x = self.norm_pre(x) + + # Mask background tokens when evaluating (B=1) + if bg_mask is not None and B == 1: + bg_mask_cat = torch.cat( + (torch.ones((1, 1), dtype=torch.bool, device=x.device), bg_mask.view(1, -1)), + dim=1, + ) + x = x[bg_mask_cat].unsqueeze(0) + + if self.grad_checkpointing and not torch.jit.is_scripting(): + from timm.models._manipulate import checkpoint_seq + x = checkpoint_seq(self.blocks, x, attn_bias, bg_mask) + else: + x = self.blocks(x, attn_bias, bg_mask) + + x = self.norm(x) + return x + + @register_model(ModelType.TITAN_SLIDE) class TitanSlideEncoderModel(TorchModel): def __init__( @@ -112,6 +250,7 @@ def __init__( model_path, use_gpu: bool = True, gpu_device_id: int | List[int] | None = None, + patch_oom: bool = True, ): """Initialize TITAN slide encoder model. @@ -122,7 +261,13 @@ def __init__( model_path: Path to slide encoder model directory or HuggingFace repo ID. use_gpu: Whether to use GPU (default: True). gpu_device_id: GPU device ID or list of IDs for multi-GPU (default: None). + patch_oom: Apply GPU float16 / expand() monkey-patches that fix CPU/GPU OOM + on large slides (>25k patches) and pin the model to the validated + HuggingFace revision (default: True). Set to False to load the + latest unpinned TITAN code without any monkey-patches โ€” useful when + testing upstream changes or on small slides where OOM is not a concern. """ + self._patch_oom = patch_oom if model_path is None: model_path = ModelType.TITAN_SLIDE.path @@ -139,30 +284,66 @@ def __init__( # Load the TITAN model from HuggingFace or saved directory # TITAN doesn't support Flash Attention 2.0, so we use eager mode # Use locking when downloading from HuggingFace + revision = _TITAN_PINNED_REVISION if self._patch_oom else None with model_download_lock(model_path) as should_download: try: model_obj = AutoModel.from_pretrained( model_path, trust_remote_code=True, + revision=revision, attn_implementation="eager", ) except TypeError: # Fallback for older transformers that don't support attn_implementation model_obj = AutoModel.from_pretrained( - model_path, trust_remote_code=True + model_path, + trust_remote_code=True, + revision=revision, ) super().__init__(model_path, model_obj, use_gpu, gpu_device_id) def get_model_fun(self) -> Callable: """Get model inference function for TITAN slide encoder. - The TITAN slide encoder uses encode_slide_from_patch_features method - which requires patch features, coordinates, and patch size at level 0. + Applies monkey-patches to the TITAN vision encoder to fix CPU/GPU RAM OOM + on large IMPACT slides (>25k patches): - Returns: - Callable that takes patch features, coords, and patch_size, returns slide-level features - with shape (768,) matching the GIGAPATH slide encoder output format. + 1. ``get_alibi`` โ†’ GPU float16 via torch.cdist + Eliminates ~82 GB CPU RAM peak for N=30k patches. + 2. ``forward_features`` โ†’ uses expand() instead of repeat() + Avoids a 22 GB copy of the bias tensor for N=30k. + + Memory budget on A100 (80 GB): bias (22 GB) + model (2 GB) + QK^T (22 GB) + + intermediates (~3 GB) โ‰ˆ 49 GB โ†’ fits with headroom. + To further reduce allocator fragmentation set + ``PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`` in the process + environment *before* Python starts (the CUDA allocator reads this env var + once at initialization, before any model is loaded). """ + # Apply monkey-patches to the vision encoder + if getattr(self, "_patch_oom", True): + vision_enc = self.obj.vision_encoder + vision_enc.get_alibi = types.MethodType(_titan_get_alibi_gpu_float16, vision_enc) + vision_enc.forward_features = types.MethodType(_titan_forward_features_efficient, vision_enc) + logger.debug( + "TITAN: applied GPU float16 get_alibi + expand-based forward_features monkey-patches" + ) + else: + logger.debug("TITAN: patch_oom=False, running with unpatched upstream code") + + # Use SDPBackend.EFFICIENT_ATTENTION in model_fun to prevent the math kernel + # from materializing the full QK^T matrix (~22 GB for N=18k), which would OOM. + # EFFICIENT_ATTENTION requires CUDA compute >= 8.0 (A100+); fall back to the + # default SDPA kernel selection on older hardware (P40, V100, etc.). + _efficient_ctx = contextlib.nullcontext + try: + from torch.nn.attention import sdpa_kernel, SDPBackend + if self.device.type == 'cuda': + major, _ = torch.cuda.get_device_capability(self.device) + if major >= 8: + _efficient_ctx = lambda: sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION) + except (ImportError, RuntimeError): + pass def model_fun(patch_features, coords, patch_size): """Run TITAN slide encoder on patch features with coordinates and patch size.""" @@ -170,6 +351,7 @@ def model_fun(patch_features, coords, patch_size): torch.no_grad(), torch.inference_mode(), torch.autocast(device_type=self.device.type, dtype=torch.float16), + _efficient_ctx(), ): patch_features = patch_features.to(self.device, non_blocking=True) coords = coords.to(self.device, non_blocking=True) diff --git a/mussel/models/model_factory.py b/mussel/models/model_factory.py index ad16a371..685885d7 100644 --- a/mussel/models/model_factory.py +++ b/mussel/models/model_factory.py @@ -239,13 +239,14 @@ def decorator(cls): class ModelFactory(ABC): @abstractmethod - def get_model(self, model_path, use_gpu, gpu_device_id) -> "Model": + def get_model(self, model_path, use_gpu, gpu_device_id, **kwargs) -> "Model": """Get a model instance. Args: model_path: Path to model weights or config. use_gpu: Whether to use GPU. gpu_device_id: GPU device ID or list of IDs. + **kwargs: Extra keyword arguments forwarded to the model constructor. Returns: Model instance. @@ -259,8 +260,8 @@ class _SimpleModelFactory(ModelFactory): def __init__(self, model_cls): self._cls = model_cls - def get_model(self, model_path=None, use_gpu=True, gpu_device_id=None) -> "Model": - return self._cls(model_path, use_gpu, gpu_device_id) + def get_model(self, model_path=None, use_gpu=True, gpu_device_id=None, **kwargs) -> "Model": + return self._cls(model_path, use_gpu, gpu_device_id, **kwargs) def get_model_factory( diff --git a/mussel/utils/__init__.py b/mussel/utils/__init__.py index c47a3b0f..355f594b 100644 --- a/mussel/utils/__init__.py +++ b/mussel/utils/__init__.py @@ -1,19 +1,5 @@ from .artifact_removal import GrandQCArtifactRemover from .converter import AnyToTiffConverter -from .feature_extract import (DatasetProcessor, FeatureExtractionResult, - H5DatasetProcessor, ImageFolderProcessor, - TileCoordProcessor, aggregate_sample_features, - aggregate_slide_features, - aggregate_slide_features_batch, - extract_patch_features, - extract_patch_features_batch, filter_features, - get_batch_size_for_model, - get_classifier_pkl_from_model_dir, - get_dataset_processor, get_features, - get_model_path_from_dir, process_dataset, - resolve_aggregation_method, - resolve_patch_encoder, save_features, - subsample_tiles) from .file import (WSI_EXTENSIONS, collect_wsi_paths, download_model_path, ensure_directory_exists, get_slide_id_from_path, get_slide_ids_from_paths, is_remote_path, load_classifier, @@ -27,3 +13,68 @@ from .timer import timed from .visualization import visualize_heatmap from .wsi_backend import open_slide + +_FEATURE_EXTRACT_EXPORTS = { + "DatasetProcessor", + "FeatureExtractionResult", + "H5DatasetProcessor", + "ImageFolderProcessor", + "TileCoordProcessor", + "aggregate_sample_features", + "aggregate_slide_features", + "aggregate_slide_features_batch", + "extract_patch_features", + "extract_patch_features_batch", + "filter_features", + "get_batch_size_for_model", + "get_classifier_pkl_from_model_dir", + "get_dataset_processor", + "get_features", + "get_model_path_from_dir", + "process_dataset", + "resolve_aggregation_method", + "resolve_patch_encoder", + "save_features", + "subsample_tiles", +} + + +def __getattr__(name): + if name in _FEATURE_EXTRACT_EXPORTS: + from . import feature_extract + + value = getattr(feature_extract, name) + globals()[name] = value + return value + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +__all__ = [ + "GrandQCArtifactRemover", + "AnyToTiffConverter", + "WSI_EXTENSIONS", + "collect_wsi_paths", + "download_model_path", + "ensure_directory_exists", + "get_slide_id_from_path", + "get_slide_ids_from_paths", + "is_remote_path", + "load_classifier", + "load_features_from_h5", + "resolve_remote_paths", + "safe_path_join", + "save_hdf5", + "save_torch_tensor", + "collate_features", + "contours_to_polygon", + "draw_slide_mask", + "get_level_for_magnification", + "get_slide_mpp", + "save_patches_png", + "segment_tissue", + "export_tiles", + "timed", + "visualize_heatmap", + "open_slide", + *_FEATURE_EXTRACT_EXPORTS, +] diff --git a/mussel/utils/feature_extract.py b/mussel/utils/feature_extract.py index d6fe1cdd..3f86248a 100644 --- a/mussel/utils/feature_extract.py +++ b/mussel/utils/feature_extract.py @@ -767,6 +767,7 @@ def _apply_slide_aggregation( coords: Optional[np.ndarray] = None, patch_size: Optional[int] = None, slide_model=None, + slide_model_kwargs: Optional[dict] = None, ) -> np.ndarray: """Apply slide-level aggregation to patch features. @@ -786,6 +787,9 @@ def _apply_slide_aggregation( If not provided, will be extracted from h5 file 'coords' attributes or default to 256. slide_model: Optional pre-loaded slide encoder model instance. If provided, slide_model_type and slide_model_path are ignored. + slide_model_kwargs: Extra keyword arguments forwarded to the slide model + constructor when loading the model. For TITAN_SLIDE, patch_oom=True is + the default OOM fix; set {"patch_oom": False} to disable it. Returns: Numpy array of aggregated features. @@ -834,7 +838,7 @@ def _apply_slide_aggregation( if model_factory is None: raise ValueError(f"Slide model type {slide_model_type} not recognized") slide_model = model_factory.get_model( - slide_model_path, use_gpu, gpu_device_id + slide_model_path, use_gpu, gpu_device_id, **(slide_model_kwargs or {}) ) model_fun = slide_model.get_model_fun() @@ -932,6 +936,8 @@ def get_features( aggregation_method: str = "identity", model=None, slide_model=None, + model_kwargs: Optional[dict] = None, + slide_model_kwargs: Optional[dict] = None, ) -> tuple[np.ndarray, np.ndarray]: """Extract features from whole slide image tiles. @@ -961,6 +967,10 @@ def get_features( model_type and model_path are ignored. slide_model: Optional pre-loaded slide encoder model instance. If provided, slide_model_type and slide_model_path are ignored for slide encoding. + model_kwargs: Extra keyword arguments forwarded to the patch model constructor. + slide_model_kwargs: Extra keyword arguments forwarded to the slide model + constructor. For TITAN_SLIDE, patch_oom=True is the default OOM fix; + set {"patch_oom": False} to disable it. Returns: Tuple of (features array, labels array). @@ -1016,7 +1026,9 @@ def get_features( model_factory = get_model_factory(model_type) if model_factory is None: raise ValueError("model not recognized") - model = model_factory.get_model(model_path, use_gpu, gpu_device_id) + model = model_factory.get_model( + model_path, use_gpu, gpu_device_id, **(model_kwargs or {}) + ) else: logger.info("using pre-loaded model") preprocessing = model.get_preprocessing_fun() @@ -1067,6 +1079,7 @@ def get_features( coords=coords, patch_size=patch_size, slide_model=slide_model, + slide_model_kwargs=slide_model_kwargs, ) return features, labels @@ -1089,6 +1102,7 @@ def extract_patch_features( pin_memory: bool = True, is_test_run: bool = False, embedding_precision: str = "float32", + model_kwargs: Optional[dict] = None, ) -> str: """Extract patch-level features from whole slide image (Step 1: Patch Encoding). @@ -1111,6 +1125,7 @@ def extract_patch_features( num_workers: Number of worker processes for data loading (default: 16). pin_memory: Whether to pin memory for data loading (default: True). is_test_run: If True, only process first 3 batches (default: False). + model_kwargs: Extra keyword arguments forwarded to the patch model constructor. embedding_precision: Numeric precision for saved patch embeddings. "float32" (default) preserves full model precision; "float16" halves storage size; "bfloat16" uses brain-float format. @@ -1129,7 +1144,9 @@ def extract_patch_features( model_factory = get_model_factory(model_type) if model_factory is None: raise ValueError("model not recognized") - model = model_factory.get_model(model_path, use_gpu, gpu_device_id) + model = model_factory.get_model( + model_path, use_gpu, gpu_device_id, **(model_kwargs or {}) + ) if model_save_path is not None: Path(model_save_path).parent.mkdir(parents=True, exist_ok=True) logger.info(f"saving model to {model_save_path}") @@ -1216,6 +1233,7 @@ def extract_patch_features_batch( pin_memory=True, is_test_run=False, embedding_precision="float32", + model_kwargs=None, ): """Extract patch-level features from multiple slides in batch mode. @@ -1240,6 +1258,7 @@ def extract_patch_features_batch( num_workers: Number of worker processes for data loading (default: 16). pin_memory: Whether to pin memory for data loading (default: True). is_test_run: If True, only process first 3 batches per slide (default: False). + model_kwargs: Extra keyword arguments forwarded to the patch model constructor. embedding_precision: Numeric precision for saved patch embeddings. "float32" (default) preserves full model precision; "float16" halves storage size; "bfloat16" uses brain-float format. @@ -1278,7 +1297,9 @@ def extract_patch_features_batch( model_factory = get_model_factory(model_type) if model_factory is None: raise ValueError("model not recognized") - model = model_factory.get_model(model_path, use_gpu, gpu_device_id) + model = model_factory.get_model( + model_path, use_gpu, gpu_device_id, **(model_kwargs or {}) + ) preprocessing = model.get_preprocessing_fun() model_fun = model.get_model_fun() @@ -1349,6 +1370,7 @@ def aggregate_slide_features_batch( slide_batch_size=8, max_slide_patches=None, embedding_precision="float32", + slide_model_kwargs=None, ): """Aggregate patch-level features to slide-level for multiple slides (Step 2: Batch Slide Encoding). @@ -1379,6 +1401,9 @@ def aggregate_slide_features_batch( When a slide exceeds this limit, patches are randomly subsampled before encoding. Useful for TITAN whose alibi attention is O(Nยฒ) and OOMs on very large slides. None (default) disables subsampling. + slide_model_kwargs: Extra keyword arguments forwarded to the slide model constructor. + For TITAN_SLIDE, patch_oom=True is the default OOM fix; set + {"patch_oom": False} to disable it. embedding_precision: Numeric precision for saved slide embeddings ("float32", "float16", or "bfloat16"). Default "float32". Applied to the aggregated output before saving; input patch features are always read as-is. @@ -1432,6 +1457,7 @@ def aggregate_slide_features_batch( gpu_device_ids=gpu_device_ids, coords=coords, patch_size=patch_size, + slide_model_kwargs=slide_model_kwargs, ) feature_dtype = _parse_feature_dtype(embedding_precision) @@ -1515,7 +1541,9 @@ def aggregate_slide_features_batch( model_factory = get_model_factory(model_type) if model_factory is None: raise ValueError(f"Slide model type {model_type} not recognized") - model = model_factory.get_model(model_path, use_gpu, gpu_device_id) + model = model_factory.get_model( + model_path, use_gpu, gpu_device_id, **(slide_model_kwargs or {}) + ) model_fun = model.get_model_fun() successful_slides = [] @@ -1744,6 +1772,7 @@ def aggregate_slide_features( gpu_device_id: Optional[Union[int, List[int]]] = None, gpu_device_ids: Optional[List[int]] = None, embedding_precision: str = "float32", + slide_model_kwargs: Optional[dict] = None, ) -> Union[tuple[Optional[str], Optional[str]], np.ndarray]: """Aggregate patch-level features to slide-level (Step 2: Slide Encoding). @@ -1768,6 +1797,9 @@ def aggregate_slide_features( embedding_precision: Numeric precision for saved slide embeddings ("float32", "float16", or "bfloat16"). Default "float32". Cast is applied to the aggregated output before saving; input patch features are read as-is. + slide_model_kwargs: Extra keyword arguments forwarded to the slide model + constructor. For TITAN_SLIDE, patch_oom=True is the default OOM fix; + set {"patch_oom": False} to disable it. Returns: Tuple of (output_h5_path, output_pt_path) if saving, otherwise features tensor. @@ -1798,6 +1830,7 @@ def aggregate_slide_features( gpu_device_ids=gpu_device_ids, coords=coords, patch_size=patch_size, + slide_model_kwargs=slide_model_kwargs, ) feature_dtype = _parse_feature_dtype(embedding_precision) @@ -1849,6 +1882,8 @@ def save_features( slide_model_type: Optional[ModelType] = None, slide_model_path: Optional[str] = None, embedding_precision: str = "float32", + model_kwargs: Optional[dict] = None, + slide_model_kwargs: Optional[dict] = None, ) -> tuple[str, Optional[str]]: """Extract features from whole slide image and save to HDF5 and PyTorch formats. @@ -1881,6 +1916,10 @@ def save_features( The required patch encoder will be automatically inferred and used. For example, specifying GIGAPATH_SLIDE will automatically use GIGAPATH as the patch encoder. slide_model_path: Optional path to slide encoder model weights. + model_kwargs: Extra keyword arguments forwarded to the patch model constructor. + slide_model_kwargs: Extra keyword arguments forwarded to the slide model + constructor. For TITAN_SLIDE, patch_oom=True is the default OOM fix; + set {"patch_oom": False} to disable it. embedding_precision: Numeric precision for saved patch embeddings. "float32" (default) preserves full model precision; "float16" halves storage size; "bfloat16" uses brain-float format. @@ -1938,6 +1977,7 @@ def save_features( pin_memory=pin_memory, is_test_run=is_test_run, embedding_precision="float32", + model_kwargs=model_kwargs, ) # Step 2: Aggregate to slide level โ€” apply embedding_precision to final output @@ -1952,6 +1992,7 @@ def save_features( gpu_device_id=gpu_device_id, gpu_device_ids=gpu_device_ids, embedding_precision=embedding_precision, + slide_model_kwargs=slide_model_kwargs, ) else: # Single-step process (backward compatible) @@ -1964,7 +2005,9 @@ def save_features( model_factory = get_model_factory(model_type) if model_factory is None: raise ValueError("model not recognized") - model = model_factory.get_model(model_path, use_gpu, gpu_device_id) + model = model_factory.get_model( + model_path, use_gpu, gpu_device_id, **(model_kwargs or {}) + ) if model_save_path is not None: Path(model_save_path).parent.mkdir(parents=True, exist_ok=True) logger.info(f"saving model to {model_save_path}") diff --git a/pyproject.toml b/pyproject.toml index 531c8042..272c283f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -170,7 +170,7 @@ requires = ["setuptools"] build-backend = "setuptools.build_meta" [tool.pytest.ini_options] -addopts = "-v --tb=short --import-mode=importlib" +addopts = "-v --tb=short --import-mode=importlib -m \"not integration\"" markers = [ "slow: marks tests as slow and require significant time to run (deselect with '-m \"not slow\"')", "integration: marks integration tests that perform actual model inference", diff --git a/tests/mussel/cli/test_create_class_embeddings.py b/tests/mussel/cli/test_create_class_embeddings.py index e17c3ce4..5430561a 100644 --- a/tests/mussel/cli/test_create_class_embeddings.py +++ b/tests/mussel/cli/test_create_class_embeddings.py @@ -1,12 +1,10 @@ -import os - -from omegaconf import OmegaConf +import torch import mussel.cli.create_class_embeddings from mussel.cli.create_class_embeddings import ClassEmbeddingConfig -def test_create_class_embeddings(tmp_path): +def test_create_class_embeddings(tmp_path, monkeypatch): annotation_classes = [ "carcinoma in situ", "invasive carcinoma", @@ -18,8 +16,32 @@ def test_create_class_embeddings(tmp_path): "sarcoma", ] output_pt_path = tmp_path / "test.pt" + + class FakeModel: + def encode_text(self, text): + return torch.arange(4, dtype=torch.float32).unsqueeze(0) + text.float() + + def fake_create_model_and_transforms(model_path): + return FakeModel(), None, None + + def fake_get_tokenizer(model_path): + return lambda class_text: torch.tensor([[len(class_text)]], dtype=torch.float32) + + monkeypatch.setattr( + mussel.cli.create_class_embeddings.open_clip, + "create_model_and_transforms", + fake_create_model_and_transforms, + ) + monkeypatch.setattr( + mussel.cli.create_class_embeddings.open_clip, + "get_tokenizer", + fake_get_tokenizer, + ) + cfg = ClassEmbeddingConfig( classes=annotation_classes, output_pt_path=output_pt_path ) mussel.cli.create_class_embeddings.main(cfg) - assert os.path.exists(output_pt_path) + assert output_pt_path.exists() + class_emb = torch.load(output_pt_path, weights_only=True) + assert class_emb.shape == (len(annotation_classes), 4) diff --git a/tests/mussel/models/test_encoder_integration.py b/tests/mussel/models/test_encoder_integration.py index 90c8cce5..2cc633d5 100644 --- a/tests/mussel/models/test_encoder_integration.py +++ b/tests/mussel/models/test_encoder_integration.py @@ -36,6 +36,11 @@ _SLIDE_PATH = str(_TESTDATA / "948176.svs") _PATCH_H5 = str(_TESTDATA / "948176.patch.h5") +_skip_if_no_testdata = pytest.mark.skipif( + not (Path(_SLIDE_PATH).exists() and Path(_PATCH_H5).exists()), + reason="Test slide data not available (948176.svs / 948176.patch.h5)", +) + # Slide encoders that are *encoder-agnostic* (not listed in SLIDE_ENCODER_COMPATIBILITY # because they work with any patch encoder โ€” e.g. ABMIL). _AGNOSTIC_SLIDE_ENCODERS = {ModelType.ABMIL_SLIDE} @@ -137,6 +142,7 @@ def wrapper(*args, **kwargs): @pytest.mark.slow @pytest.mark.integration +@_skip_if_no_testdata @pytest.mark.timeout(600) @pytest.mark.parametrize("model_type", _PATCH_ENCODER_TYPES, ids=lambda m: m.name) def test_patch_encoder_extracts_features(tmp_path, model_type, use_gpu): @@ -372,6 +378,7 @@ def run_slide(): @pytest.mark.slow @pytest.mark.integration +@_skip_if_no_testdata @pytest.mark.timeout(600) @pytest.mark.parametrize("model_type", _PATCH_ENCODER_TYPES, ids=lambda m: m.name) def test_patch_encoder_is_deterministic(tmp_path, model_type, use_gpu): @@ -421,6 +428,7 @@ def run(): @pytest.mark.slow @pytest.mark.integration +@_skip_if_no_testdata @pytest.mark.timeout(600) @pytest.mark.parametrize("model_type", _PATCH_ENCODER_TYPES, ids=lambda m: m.name) def test_patch_encoder_matches_snapshot( @@ -477,6 +485,81 @@ def run(): run() +@pytest.mark.slow +@pytest.mark.integration +@pytest.mark.timeout(300) +@pytest.mark.parametrize("slide_model_type", _SLIDE_ENCODER_TYPES, ids=lambda m: m.name) +def test_slide_encoder_matches_snapshot( + tmp_path, slide_model_type, use_gpu, update_snapshots +): + """Slide encoder embeddings match a previously saved golden snapshot (regression test). + + On first run (or with ``--update-snapshots``) the current output is saved + to ``tests/testdata/snapshots/.npy`` and the test is skipped. + On subsequent runs the saved snapshot is compared with ``np.allclose``. + + Generate / refresh snapshots:: + + uv run pytest tests/mussel/models/test_encoder_integration.py \\ + -k test_slide_encoder_matches_snapshot --use-gpu --update-snapshots + """ + snapshot_path = _SNAPSHOT_DIR / f"{slide_model_type.name}.npy" + + required_patch_enc = get_required_patch_encoder(slide_model_type) + patch_dim = _SLIDE_ENCODER_INPUT_DIM.get( + slide_model_type + ) or _PATCH_ENCODER_DIM.get(required_patch_enc) + if patch_dim is None: + pytest.skip( + f"Feature dim for {required_patch_enc.name} not in _PATCH_ENCODER_DIM" + ) + + n_patches = 32 + rng = np.random.default_rng(42) + fake_features = rng.standard_normal((n_patches, patch_dim)).astype(np.float32) + fake_features /= np.linalg.norm(fake_features, axis=1, keepdims=True) + 1e-8 + + patch_size_native = 512 + fake_coords = np.stack( + [ + np.arange(n_patches) * patch_size_native, + np.zeros(n_patches, dtype=np.int64), + ], + axis=1, + ).astype(np.int64) + + @_skip_on_load_failure + def run(): + return _apply_slide_aggregation( + features=fake_features, + aggregation_method="model", + slide_model_type=slide_model_type, + use_gpu=use_gpu, + coords=fake_coords, + patch_size=patch_size_native, + ) + + result = run() + + if update_snapshots or not snapshot_path.exists(): + _SNAPSHOT_DIR.mkdir(parents=True, exist_ok=True) + np.save(snapshot_path, result) + if not update_snapshots: + pytest.skip( + f"Snapshot saved to {snapshot_path.name}; re-run to compare." + ) + return + + golden = np.load(snapshot_path) + assert ( + result.shape == golden.shape + ), f"{slide_model_type.name}: shape {result.shape} != snapshot {golden.shape}" + assert np.allclose(result, golden, rtol=1e-3, atol=1e-4), ( + f"{slide_model_type.name}: embedding differs from snapshot " + "(model weights, attention backend, or preprocessing changed?)" + ) + + # --------------------------------------------------------------------------- # Encoder-agnostic slide encoder integration tests (ABMIL) # --------------------------------------------------------------------------- diff --git a/tests/mussel/models/test_model_classes.py b/tests/mussel/models/test_model_classes.py index bf4520f0..ac160700 100644 --- a/tests/mussel/models/test_model_classes.py +++ b/tests/mussel/models/test_model_classes.py @@ -473,6 +473,84 @@ def test_calls_model_and_squeezes(self): assert result.shape == torch.Size([embed_dim]) +class TestTitanSlideEncoderModelFun: + def test_calls_encode_slide_from_patch_features_and_squeezes(self): + embed_dim = 768 + batch_size = 1 + num_patches = 100 + patch_dim = 768 + patch_size = 512 + + mock_model = MagicMock() + mock_model.encode_slide_from_patch_features = MagicMock( + return_value=torch.rand(batch_size, embed_dim) + ) + + m = _make_model(TitanSlideEncoderModel, mock_model) + m._patch_oom = True + model_fun = m.get_model_fun() + + patch_features = torch.rand(batch_size, num_patches, patch_dim) + coords = torch.rand(batch_size, num_patches, 2) + + result = model_fun(patch_features, coords, patch_size) + + mock_model.encode_slide_from_patch_features.assert_called_once() + assert result.device.type == "cpu" + assert result.shape == torch.Size([embed_dim]) + + def test_patch_oom_false_skips_monkey_patches(self): + """When patch_oom=False, get_model_fun must NOT monkey-patch the vision encoder.""" + embed_dim = 768 + batch_size = 1 + num_patches = 100 + patch_dim = 768 + patch_size = 512 + + mock_model = MagicMock() + mock_model.encode_slide_from_patch_features = MagicMock( + return_value=torch.rand(batch_size, embed_dim) + ) + original_get_alibi = mock_model.vision_encoder.get_alibi + original_forward_features = mock_model.vision_encoder.forward_features + + m = _make_model(TitanSlideEncoderModel, mock_model) + m._patch_oom = False + model_fun = m.get_model_fun() + + # Patches must NOT have been applied + assert mock_model.vision_encoder.get_alibi == original_get_alibi + assert mock_model.vision_encoder.forward_features == original_forward_features + + patch_features = torch.rand(batch_size, num_patches, patch_dim) + coords = torch.rand(batch_size, num_patches, 2) + result = model_fun(patch_features, coords, patch_size) + + mock_model.encode_slide_from_patch_features.assert_called_once() + assert result.shape == torch.Size([embed_dim]) + + +class TestGigapathSlideEncoderModelFun: + def test_calls_model_and_squeezes(self): + embed_dim = 768 + batch_size = 1 + num_patches = 256 + patch_dim = 1536 + + mock_model = MagicMock(return_value=[torch.rand(batch_size, embed_dim)]) + m = _make_model(GigapathSlideEncoderModel, mock_model) + model_fun = m.get_model_fun() + + patch_features = torch.rand(batch_size, num_patches, patch_dim) + coords = torch.rand(batch_size, num_patches, 2) + + result = model_fun(patch_features, coords) + + mock_model.assert_called_once() + assert result.device.type == "cpu" + assert result.shape == torch.Size([embed_dim]) + + # --------------------------------------------------------------------------- # Slide encoders โ€“ save() raises ValueError for file paths # --------------------------------------------------------------------------- diff --git a/tests/mussel/models/test_titan_get_alibi_patch.py b/tests/mussel/models/test_titan_get_alibi_patch.py new file mode 100644 index 00000000..92b7b744 --- /dev/null +++ b/tests/mussel/models/test_titan_get_alibi_patch.py @@ -0,0 +1,156 @@ +"""Tests for the TITAN get_alibi GPU float16 monkey-patch. + +These tests verify that the patch: +1. Produces numerically close results to the original numpy float64 implementation +2. Stays within memory bounds for large N +3. Does not regress on model output shape/type +""" +import math + +import numpy as np +import pytest +import torch + +from mussel.models.conch import _get_slopes + + +# --------------------------------------------------------------------------- +# Helpers: reference implementation (original numpy float64 from TITAN) +# --------------------------------------------------------------------------- + +def _get_alibi_original_numpy(w: int, h: int, num_heads: int = 12, bg_mask=None): + """Original numpy float64 implementation from TITAN.""" + x, y = np.meshgrid(np.arange(w), np.arange(h), indexing='ij') + if bg_mask is not None: + x = x[bg_mask.cpu().squeeze(0)] + y = y[bg_mask.cpu().squeeze(0)] + points = np.stack([x.ravel(), y.ravel()], axis=1) + diffs = points[:, None, :] - points[None, :, :] + dists = np.sqrt(np.sum(diffs ** 2, axis=-1)) + slopes = torch.tensor(_get_slopes(num_heads), dtype=torch.float32).view(num_heads, 1, 1) + n_patches = dists.shape[-1] + dists_tensor = torch.tensor(dists, dtype=torch.float32).view(1, n_patches, n_patches) + bias_matrix = dists_tensor * slopes * -1 + embed_len = n_patches + 1 + all_bias = torch.zeros(1, num_heads, embed_len, embed_len) + all_bias[:, :, 1:, 1:] = bias_matrix + return all_bias + + +def _get_alibi_gpu_float16_standalone(w: int, h: int, num_heads: int = 12, bg_mask=None, + device: str = 'cpu'): + """Standalone version of the GPU float16 patch for testing without loading TITAN.""" + dtype = torch.float16 + dev = torch.device(device) + x_c = torch.arange(w, device=dev, dtype=dtype) + y_c = torch.arange(h, device=dev, dtype=dtype) + gx, gy = torch.meshgrid(x_c, y_c, indexing='ij') + if bg_mask is not None: + if bg_mask.dim() == 3: + mf = bg_mask.to(dev).squeeze(0).bool() # (W, H) + pts_x, pts_y = gx[mf], gy[mf] + else: + mf = bg_mask.to(dev).squeeze(0).bool() # flat (W*H,) + pts_x = gx.ravel()[mf] + pts_y = gy.ravel()[mf] + else: + pts_x, pts_y = gx.ravel(), gy.ravel() + points = torch.stack([pts_x, pts_y], dim=1) + dists = torch.cdist(points.float(), points.float(), p=2).to(dtype) + slopes = torch.tensor( + _get_slopes(num_heads), dtype=dtype, device=dev + ).view(num_heads, 1, 1) + n_patches = dists.shape[0] + bias_matrix = -dists.unsqueeze(0) * slopes + embed_len = n_patches + 1 + all_bias = torch.zeros(1, num_heads, embed_len, embed_len, dtype=dtype, device=dev) + all_bias[:, :, 1:, 1:] = bias_matrix + return all_bias + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + +class TestGetAlibiGpuFloat16: + """Test the GPU float16 get_alibi monkey-patch.""" + + @pytest.mark.parametrize("w,h", [(6, 6), (14, 14), (30, 40), (100, 80)]) + def test_output_shape(self, w, h): + """Output shape matches original.""" + num_heads = 12 + ref = _get_alibi_original_numpy(w, h, num_heads) + patched = _get_alibi_gpu_float16_standalone(w, h, num_heads) + assert patched.shape == ref.shape, f"Shape mismatch: {patched.shape} vs {ref.shape}" + + @pytest.mark.parametrize("w,h", [(6, 6), (14, 14), (30, 40)]) + def test_numerical_closeness(self, w, h): + """Patched output is numerically close to reference (float16 vs float64).""" + num_heads = 12 + ref = _get_alibi_original_numpy(w, h, num_heads).float() + patched = _get_alibi_gpu_float16_standalone(w, h, num_heads).float() + # float16 has ~3 significant digits; allow relative tolerance of 1e-2 + assert torch.allclose(ref, patched, rtol=1e-2, atol=1e-3), ( + f"Output too different: max abs diff = {(ref - patched).abs().max():.4f}" + ) + + def test_with_bg_mask(self): + """Mask-filtered version has correct shape and values.""" + w, h = 20, 20 + # bg_mask shape is (1, H, W) bool as TITAN uses it + bg_mask = torch.zeros(1, w, h, dtype=torch.bool) + bg_mask[0, ::2, ::2] = True # every other cell + n_fg = bg_mask.sum().item() + + patched = _get_alibi_gpu_float16_standalone(w, h, bg_mask=bg_mask) + expected_size = (1, 12, n_fg + 1, n_fg + 1) + assert patched.shape == expected_size, f"Shape: {patched.shape} vs {expected_size}" + + def test_large_n_no_oom(self): + """Large N (simulating a 33k-patch slide) doesn't OOM on CPU.""" + # Use CPU to test logic without needing GPU + # N=1000 is enough to verify the pattern; real OOM tests need GPU + w, h = 50, 50 # 2500 patches (manageable on CPU) + patched = _get_alibi_gpu_float16_standalone(w, h, num_heads=12) + assert patched.shape == (1, 12, 2501, 2501) + assert patched.dtype == torch.float16 + assert torch.isfinite(patched).all(), "Non-finite values in output" + + def test_output_dtype_and_device(self): + """Output is float16 on correct device.""" + patched = _get_alibi_gpu_float16_standalone(10, 10) + assert patched.dtype == torch.float16 + assert patched.device.type == 'cpu' + + def test_diagonal_is_zero(self): + """Self-distance (diagonal) should produce maximum bias (distance=0).""" + w, h = 4, 4 + patched = _get_alibi_gpu_float16_standalone(w, h, num_heads=12) + # bias[head, i+1, i+1] = -slope * 0 = 0 for all i (self-distance = 0) + for head in range(12): + diag = torch.diagonal(patched[0, head, 1:, 1:]) + assert (diag == 0).all(), f"Non-zero diagonal for head {head}" + + def test_cosine_similarity_with_reference(self): + """Flattened output has cosine similarity > 0.99 with reference.""" + w, h = 20, 20 + ref = _get_alibi_original_numpy(w, h).float().flatten() + patched = _get_alibi_gpu_float16_standalone(w, h).float().flatten() + cos_sim = torch.nn.functional.cosine_similarity(ref.unsqueeze(0), patched.unsqueeze(0)) + assert cos_sim.item() > 0.99, f"Cosine similarity too low: {cos_sim.item():.4f}" + + +class TestMonkeyPatchApplied: + """Test that the patch functions exist in the conch module.""" + + def test_import(self): + """The patch functions exist as module-level callables in conch.py.""" + import ast + from pathlib import Path + + conch_path = Path(__file__).parents[3] / "mussel" / "models" / "conch.py" + src = conch_path.read_text() + tree = ast.parse(src) + fn_names = {n.name for n in ast.walk(tree) if isinstance(n, ast.FunctionDef)} + assert "_titan_get_alibi_gpu_float16" in fn_names + assert "_titan_forward_features_efficient" in fn_names diff --git a/tests/mussel/test_import_order.py b/tests/mussel/test_import_order.py new file mode 100644 index 00000000..a052b9d0 --- /dev/null +++ b/tests/mussel/test_import_order.py @@ -0,0 +1,32 @@ +import subprocess +import sys + + +def run_clean_import(script: str) -> None: + result = subprocess.run( + [sys.executable, "-c", script], + check=False, + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr + + +def test_extract_features_cli_imports_in_clean_interpreter(): + run_clean_import("import mussel.cli.extract_features") + + +def test_extract_features_cli_imports_after_datasets(): + run_clean_import( + "import mussel.datasets; " + "import mussel.cli.extract_features; " + "from mussel.utils import aggregate_slide_features_batch" + ) + + +def test_feature_extract_exports_are_lazy_importable(): + run_clean_import( + "from mussel.utils import (" + "DatasetProcessor, FeatureExtractionResult, aggregate_slide_features_batch" + ")" + ) diff --git a/tests/mussel/test_wsi_pipeline_comparison.py b/tests/mussel/test_wsi_pipeline_comparison.py deleted file mode 100644 index 597de15b..00000000 --- a/tests/mussel/test_wsi_pipeline_comparison.py +++ /dev/null @@ -1,424 +0,0 @@ -"""Comparison tests between Mussel tessellation and an external WSI patching pipeline. - -These tests verify that Mussel produces comparable results to a reference pipeline -when using the same input slide and equivalent parameters (Otsu segmentation, -same target magnification/MPP, same patch size, no overlap). - -Requires the reference pipeline to be installed at: - /gpfs/cdsi_ess/home/limr/ess/repos/wsi-patching-ref - -Run with: - uv run pytest tests/mussel/test_wsi_pipeline_comparison.py -m slow -v -""" - -from __future__ import annotations - -import os -import subprocess -import tempfile -import textwrap -from pathlib import Path -from typing import Any, Dict - -import h5py -import numpy as np -import pytest - -# --------------------------------------------------------------------------- -# Paths -# --------------------------------------------------------------------------- - -REF_PIPELINE_DIR = os.environ.get( - "WSI_REF_PIPELINE_DIR", - "/gpfs/cdsi_ess/home/limr/ess/repos/wsi-reference-pipeline", -) -REF_PIPELINE_PYTHON = f"{REF_PIPELINE_DIR}/venv/bin/python" -MUSSEL_TEST_WSI = str(Path(__file__).parent.parent / "testdata" / "948176.svs") - -REF_PIPELINE_AVAILABLE = ( - os.path.isfile(REF_PIPELINE_PYTHON) - and os.path.isdir(REF_PIPELINE_DIR) - and os.path.isfile(MUSSEL_TEST_WSI) -) - -ref_pipeline_required = pytest.mark.skipif( - not REF_PIPELINE_AVAILABLE, - reason=( - f"Reference WSI pipeline not available at {REF_PIPELINE_DIR} or test WSI missing at {MUSSEL_TEST_WSI}" - ), -) - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _run_ref_patching( - wsi_path: str, - tmpdir: str, - target_mag: int = 20, - patch_size: int = 256, - overlap: int = 0, - min_tissue_proportion: float = 0.0, - timeout: int = 300, -) -> str: - """Run reference pipeline Otsu segmentation + patching via subprocess. - - Returns the path to the resulting coordinates H5 file. - """ - script = textwrap.dedent( - f""" - import sys, os - sys.path.insert(0, {REF_PIPELINE_DIR!r}) - from trident import load_wsi - from trident.segmentation_models import segmentation_model_factory - - wsi = load_wsi({wsi_path!r}) - seg_model = segmentation_model_factory('otsu') - wsi.segment_tissue( - seg_model, - target_mag=10, - job_dir={tmpdir!r}, - device='cpu', - verbose=False, - ) - h5_path = wsi.extract_tissue_coords( - target_mag={target_mag}, - patch_size={patch_size}, - save_coords={tmpdir!r}, - overlap={overlap}, - min_tissue_proportion={min_tissue_proportion}, - ) - print(h5_path) - """ - ) - result = subprocess.run( - [REF_PIPELINE_PYTHON, "-c", script], - capture_output=True, - text=True, - timeout=timeout, - ) - if result.returncode != 0: - raise RuntimeError( - f"Reference pipeline subprocess failed (rc={result.returncode}):\n{result.stderr}" - ) - h5_path = result.stdout.strip().splitlines()[-1] - assert os.path.isfile( - h5_path - ), f"Expected reference H5 at {h5_path!r}, got stdout: {result.stdout!r}" - return h5_path - - -def _read_coords_h5(path: str) -> tuple[np.ndarray, Dict[str, Any]]: - """Return (coords array, attrs dict) from an H5 coords file.""" - with h5py.File(path, "r") as f: - coords = f["coords"][:] - attrs = dict(f["coords"].attrs) - return coords, attrs - - -def _run_mussel_patching( - wsi_path: str, - output_h5_path: str, - patch_size: int = 256, - mpp: float = 0.5, - overlap: int = 0, - min_tissue_proportion: float = 0.0, -) -> str: - """Run Mussel segment_tissue() with HSV segmentation. - - Uses tissue_area_threshold=1 to disable area filtering โ€” the default - threshold (100) is scaled by the segmentation level downsample factor - and can exceed actual contour areas on small or coarse-resolution slides, - producing zero contours. For these comparison tests we want all tissue - regions to participate in the patch grid, which is consistent with - reference pipeline default behaviour (no minimum contour area filter). - - Returns *output_h5_path*. - """ - from mussel.utils.segment import segment_tissue - - segment_tissue( - slide_path=wsi_path, - patch_size=patch_size, - mpp=mpp, - tissue_area_threshold=1, - output_h5_path=output_h5_path, - overlap=overlap, - min_tissue_proportion=min_tissue_proportion, - ) - return output_h5_path - - -# --------------------------------------------------------------------------- -# Tests: H5 format validation (no external pipeline needed) -# --------------------------------------------------------------------------- - - -class TestMusselH5Format: - """Validate that Mussel's H5 output has the expected structure.""" - - def test_existing_patch_h5_has_coords_key(self): - """The pre-generated test fixture has a 'coords' dataset.""" - fixture = Path(MUSSEL_TEST_WSI).parent / "948176.patch.h5" - assert fixture.is_file(), f"Fixture missing: {fixture}" - with h5py.File(fixture, "r") as f: - assert "coords" in f, "Missing 'coords' key in Mussel H5" - - def test_existing_patch_h5_coords_shape(self): - """Coords have shape (N, 2) with int64 dtype.""" - fixture = Path(MUSSEL_TEST_WSI).parent / "948176.patch.h5" - coords, _ = _read_coords_h5(str(fixture)) - assert coords.ndim == 2 - assert coords.shape[1] == 2 - assert coords.dtype == np.int64 or np.issubdtype(coords.dtype, np.integer) - - def test_existing_patch_h5_attrs(self): - """Required metadata attributes are present.""" - fixture = Path(MUSSEL_TEST_WSI).parent / "948176.patch.h5" - _, attrs = _read_coords_h5(str(fixture)) - required = {"name", "patch_size", "mpp", "native_mpp", "patch_level"} - missing = required - set(attrs.keys()) - assert not missing, f"Missing attrs: {missing}" - - def test_existing_patch_h5_coords_within_slide_bounds(self): - """All coordinates are within the slide's level-0 dimensions.""" - fixture = Path(MUSSEL_TEST_WSI).parent / "948176.patch.h5" - coords, attrs = _read_coords_h5(str(fixture)) - level_dim = attrs["level_dim"] # [width, height] - assert np.all(coords[:, 0] >= 0), "Negative x coordinates" - assert np.all(coords[:, 1] >= 0), "Negative y coordinates" - assert np.all(coords[:, 0] < level_dim[0]), "x coords exceed slide width" - assert np.all(coords[:, 1] < level_dim[1]), "y coords exceed slide height" - - def test_mussel_tessellate_produces_valid_h5(self, tmp_path): - """segment_tissue() writes a valid H5 with expected structure. - - Uses tissue_area_threshold=1 to bypass the default threshold - scaling, which can filter all contours on small/coarse slides. - """ - from mussel.utils.segment import segment_tissue - - out = str(tmp_path / "test.h5") - segment_tissue( - slide_path=MUSSEL_TEST_WSI, - patch_size=256, - mpp=0.5, - tissue_area_threshold=1, - output_h5_path=out, - ) - assert os.path.isfile(out) - coords, attrs = _read_coords_h5(out) - assert coords.ndim == 2 - assert coords.shape[1] == 2 - assert coords.shape[0] > 0 - assert "name" in attrs - assert "patch_size" in attrs - - -# --------------------------------------------------------------------------- -# Tests: Reference pipeline H5 format validation -# --------------------------------------------------------------------------- - - -@ref_pipeline_required -@pytest.mark.slow -class TestTridentH5Format: - """Validate that the reference pipeline H5 output has the expected structure.""" - - @pytest.fixture(scope="class") - def ref_h5(self, tmp_path_factory): - """Run reference pipeline once for the class; cache result path.""" - tmpdir = str(tmp_path_factory.mktemp("ref_run")) - return _run_ref_patching(MUSSEL_TEST_WSI, tmpdir) - - def test_has_coords_key(self, ref_h5): - with h5py.File(ref_h5, "r") as f: - assert "coords" in f, "Missing 'coords' key in reference H5" - - def test_coords_shape(self, ref_h5): - coords, _ = _read_coords_h5(ref_h5) - assert coords.ndim == 2 - assert coords.shape[1] == 2 - assert coords.shape[0] > 0 - - def test_coords_dtype_integer(self, ref_h5): - coords, _ = _read_coords_h5(ref_h5) - assert np.issubdtype( - coords.dtype, np.integer - ), f"Expected int dtype, got {coords.dtype}" - - def test_required_attrs_present(self, ref_h5): - _, attrs = _read_coords_h5(ref_h5) - required = { - "patch_size", - "patch_size_level0", - "level0_magnification", - "target_magnification", - "overlap", - "name", - } - missing = required - set(attrs.keys()) - assert not missing, f"Missing reference pipeline attrs: {missing}" - - def test_coords_within_slide_bounds(self, ref_h5): - coords, attrs = _read_coords_h5(ref_h5) - w = attrs["level0_width"] - h = attrs["level0_height"] - assert np.all(coords[:, 0] >= 0) - assert np.all(coords[:, 1] >= 0) - assert np.all(coords[:, 0] < w), f"x coords exceed slide width {w}" - assert np.all(coords[:, 1] < h), f"y coords exceed slide height {h}" - - -# --------------------------------------------------------------------------- -# Tests: Mussel vs reference pipeline comparison -# --------------------------------------------------------------------------- - - -@ref_pipeline_required -@pytest.mark.slow -class TestTridentMusselComparison: - """Compare Mussel tessellation output against a reference WSI patching pipeline. - - Both pipelines use Otsu segmentation on the same WSI with matching - parameters (20x / 0.5 MPP, 256px patches, no overlap). - """ - - @pytest.fixture(scope="class") - def both_h5(self, tmp_path_factory): - """Run both pipelines and return (mussel_coords, mussel_attrs, ref_coords, ref_attrs).""" - td = str(tmp_path_factory.mktemp("ref_run")) - md = str(tmp_path_factory.mktemp("mussel_run")) - - ref_h5 = _run_ref_patching(MUSSEL_TEST_WSI, td, target_mag=20, patch_size=256) - mussel_h5 = _run_mussel_patching( - MUSSEL_TEST_WSI, - os.path.join(md, "mussel.h5"), - patch_size=256, - mpp=0.5, - ) - - tc, ta = _read_coords_h5(ref_h5) - mc, ma = _read_coords_h5(mussel_h5) - return mc, ma, tc, ta - - def test_patch_count_within_20_percent(self, both_h5): - """Both pipelines should produce similar patch counts (within 20%). - - Mussel uses HSV-based segmentation; the reference pipeline uses Otsu on the saturation - channel. Both operate on the same slide at approximately 0.5 MPP, but - at different internal segmentation resolutions (Mussel: level 3 / ~32x - downsample; reference pipeline: 10x thumbnail). A 20% tolerance accommodates the - resulting minor differences in detected tissue area. - """ - mc, _, tc, _ = both_h5 - n_mussel = len(mc) - n_ref = len(tc) - ratio = abs(n_mussel - n_ref) / max(n_mussel, n_ref) - assert ratio <= 0.20, ( - f"Patch count divergence too large: Mussel={n_mussel}, reference={n_ref} " - f"({ratio:.1%} difference, threshold 20%)" - ) - - def test_both_use_level0_coordinate_space(self, both_h5): - """Both pipelines should produce coordinates in level-0 pixel space. - - The test slide is 85656 ร— 19917 at level 0. Valid coordinates - must be within these bounds. - """ - mc, _, tc, _ = both_h5 - slide_w, slide_h = 85656, 19917 - # Mussel - assert np.all(mc[:, 0] < slide_w), "Mussel x coords exceed slide width" - assert np.all(mc[:, 1] < slide_h), "Mussel y coords exceed slide height" - # reference pipeline - assert np.all(tc[:, 0] < slide_w), "Reference x coords exceed slide width" - assert np.all(tc[:, 1] < slide_h), "Reference y coords exceed slide height" - - def test_coordinate_range_similar(self, both_h5): - """The spatial extent (span) of patch grids should be roughly similar. - - Both pipelines segment the same tissue, so the width and height of - the bounding box that contains all patches should agree to within 20% - of the slide dimensions. We compare spans rather than absolute min/max - because the two segmenters may disagree on whether narrow slide margins - count as tissue. - """ - mc, _, tc, _ = both_h5 - slide_w, slide_h = 85656, 19917 - tol_x = slide_w * 0.20 # 20% of slide width - tol_y = slide_h * 0.20 - - mussel_span_x = int(mc[:, 0].max()) - int(mc[:, 0].min()) - ref_span_x = int(tc[:, 0].max()) - int(tc[:, 0].min()) - mussel_span_y = int(mc[:, 1].max()) - int(mc[:, 1].min()) - ref_span_y = int(tc[:, 1].max()) - int(tc[:, 1].min()) - - assert abs(mussel_span_x - ref_span_x) < tol_x, ( - f"x span differs by more than 20% of slide width: " - f"Mussel={mussel_span_x}, reference={ref_span_x}" - ) - assert abs(mussel_span_y - ref_span_y) < tol_y, ( - f"y span differs by more than 20% of slide height: " - f"Mussel={mussel_span_y}, reference={ref_span_y}" - ) - - def test_mussel_patch_size_attr_set(self, both_h5): - """Mussel H5 should record patch_size attribute.""" - _, ma, _, _ = both_h5 - assert "patch_size" in ma, "Mussel H5 missing 'patch_size' attr" - - def test_ref_patch_size_attr_matches_input(self, both_h5): - """Reference H5 should record patch_size=256 and target_magnification=20.""" - _, _, _, ta = both_h5 - assert ta["patch_size"] == 256 - assert ta["target_magnification"] == 20 - - def test_no_duplicate_coords_mussel(self, both_h5): - """Mussel should not produce duplicate patch coordinates.""" - mc, _, _, _ = both_h5 - unique = np.unique(mc, axis=0) - assert len(unique) == len( - mc - ), f"Mussel has {len(mc) - len(unique)} duplicate coordinates" - - def test_no_duplicate_coords_ref(self, both_h5): - """Reference pipeline should not produce duplicate patch coordinates.""" - _, _, tc, _ = both_h5 - unique = np.unique(tc, axis=0) - assert len(unique) == len( - tc - ), f"Reference pipeline has {len(tc) - len(unique)} duplicate coordinates" - - def test_overlap_zero_produces_no_overlap_mussel(self, tmp_path): - """With overlap=0, Mussel patches should not overlap each other. - - For non-overlapping patches, x-coordinates should step by at least - patch_size pixels between consecutive sorted patches in a row. - Uses tissue_area_threshold=1 to ensure tissue is found. - """ - from mussel.utils.segment import segment_tissue - - out = str(tmp_path / "nooverlap.h5") - segment_tissue( - slide_path=MUSSEL_TEST_WSI, - patch_size=256, - mpp=0.5, - tissue_area_threshold=1, - output_h5_path=out, - overlap=0, - ) - coords, attrs = _read_coords_h5(out) - patch_size_px = int(attrs["patch_size"]) - - # Group by y, check x spacing - ys = np.unique(coords[:, 1]) - for y in ys: - row = np.sort(coords[coords[:, 1] == y][:, 0]) - if len(row) > 1: - steps = np.diff(row) - assert np.all( - steps >= patch_size_px - ), f"Overlapping patches in row y={y}: steps={steps[steps < patch_size_px]}" diff --git a/tests/mussel/utils/test_feature_extract.py b/tests/mussel/utils/test_feature_extract.py index c3516472..8253151b 100644 --- a/tests/mussel/utils/test_feature_extract.py +++ b/tests/mussel/utils/test_feature_extract.py @@ -23,7 +23,7 @@ def _make_mock_model(feature_dim=384): mock = MagicMock() mock.get_preprocessing_fun.return_value = None mock.get_model_fun.return_value = MagicMock( - side_effect=lambda x: __import__('torch').randn(len(x), feature_dim) + side_effect=lambda x: __import__("torch").randn(len(x), feature_dim) ) return mock @@ -40,14 +40,24 @@ def _base_patches(): # -- model= parameter ------------------------------------------------------- + def test_model_factory_not_called_when_model_provided(): """get_model_factory must not be called when a pre-loaded model is given.""" coords = np.zeros((10, 2), dtype=np.int32) - attrs = {"patch_size": 256, "patch_level": 0, "mpp": 0.5, - "patch_size_to_resize_to_for_desired_mpp": 224} + attrs = { + "patch_size": 256, + "patch_level": 0, + "mpp": 0.5, + "patch_size_to_resize_to_for_desired_mpp": 224, + } mock_model = _make_mock_model() - with patch("mussel.utils.feature_extract.get_model_factory") as mock_factory, patch("mussel.utils.feature_extract.WholeSlideImageTileCoordDataset"), patch("mussel.utils.feature_extract._make_dataloader"), patch("mussel.utils.feature_extract.process_dataset") as mock_proc: + with ( + patch("mussel.utils.feature_extract.get_model_factory") as mock_factory, + patch("mussel.utils.feature_extract.WholeSlideImageTileCoordDataset"), + patch("mussel.utils.feature_extract._make_dataloader"), + patch("mussel.utils.feature_extract.process_dataset") as mock_proc, + ): mock_proc.return_value = MagicMock( features=np.zeros((10, 384)), labels=np.zeros(10) ) @@ -59,12 +69,23 @@ def test_model_factory_not_called_when_model_provided(): def test_model_factory_called_when_model_not_provided(): """get_model_factory must be called when no pre-loaded model is given.""" coords = np.zeros((10, 2), dtype=np.int32) - attrs = {"patch_size": 256, "patch_level": 0, "mpp": 0.5, - "patch_size_to_resize_to_for_desired_mpp": 224} - - with patch("mussel.utils.feature_extract.get_model_factory") as mock_factory, patch("mussel.utils.feature_extract.WholeSlideImageTileCoordDataset"), patch("mussel.utils.feature_extract._make_dataloader"), patch("mussel.utils.feature_extract.process_dataset") as mock_proc: + attrs = { + "patch_size": 256, + "patch_level": 0, + "mpp": 0.5, + "patch_size_to_resize_to_for_desired_mpp": 224, + } + + with ( + patch("mussel.utils.feature_extract.get_model_factory") as mock_factory, + patch("mussel.utils.feature_extract.WholeSlideImageTileCoordDataset"), + patch("mussel.utils.feature_extract._make_dataloader"), + patch("mussel.utils.feature_extract.process_dataset") as mock_proc, + ): mock_model = _make_mock_model() - mock_factory.return_value = MagicMock(get_model=MagicMock(return_value=mock_model)) + mock_factory.return_value = MagicMock( + get_model=MagicMock(return_value=mock_model) + ) mock_proc.return_value = MagicMock( features=np.zeros((10, 384)), labels=np.zeros(10) ) @@ -76,12 +97,23 @@ def test_model_factory_called_when_model_not_provided(): def test_positional_args_unchanged(): """Existing positional call pattern must still work after adding model= at end.""" coords = np.zeros((5, 2), dtype=np.int32) - attrs = {"patch_size": 256, "patch_level": 0, "mpp": 0.5, - "patch_size_to_resize_to_for_desired_mpp": 224} - - with patch("mussel.utils.feature_extract.get_model_factory") as mock_factory, patch("mussel.utils.feature_extract.WholeSlideImageTileCoordDataset"), patch("mussel.utils.feature_extract._make_dataloader"), patch("mussel.utils.feature_extract.process_dataset") as mock_proc: + attrs = { + "patch_size": 256, + "patch_level": 0, + "mpp": 0.5, + "patch_size_to_resize_to_for_desired_mpp": 224, + } + + with ( + patch("mussel.utils.feature_extract.get_model_factory") as mock_factory, + patch("mussel.utils.feature_extract.WholeSlideImageTileCoordDataset"), + patch("mussel.utils.feature_extract._make_dataloader"), + patch("mussel.utils.feature_extract.process_dataset") as mock_proc, + ): mock_model = _make_mock_model() - mock_factory.return_value = MagicMock(get_model=MagicMock(return_value=mock_model)) + mock_factory.return_value = MagicMock( + get_model=MagicMock(return_value=mock_model) + ) mock_proc.return_value = MagicMock( features=np.zeros((5, 384)), labels=np.zeros(5) ) @@ -96,8 +128,12 @@ def test_positional_args_unchanged(): def test_model_invalid_interface_raises_type_error(): """Passing an object without the required methods must raise TypeError immediately.""" coords = np.zeros((5, 2), dtype=np.int32) - attrs = {"patch_size": 256, "patch_level": 0, "mpp": 0.5, - "patch_size_to_resize_to_for_desired_mpp": 224} + attrs = { + "patch_size": 256, + "patch_level": 0, + "mpp": 0.5, + "patch_size_to_resize_to_for_desired_mpp": 224, + } bad_model = object() with pytest.raises(TypeError, match="get_preprocessing_fun"): @@ -106,30 +142,47 @@ def test_model_invalid_interface_raises_type_error(): # -- slide_model= parameter ------------------------------------------------ + def test_slide_model_factory_not_called_when_slide_model_provided(): """get_model_factory must not be loaded for the slide encoder when slide_model is given.""" coords = np.zeros((10, 2), dtype=np.int32) - attrs = {"patch_size": 256, "patch_level": 0, "mpp": 0.5, - "patch_size_to_resize_to_for_desired_mpp": 224} + attrs = { + "patch_size": 256, + "patch_level": 0, + "mpp": 0.5, + "patch_size_to_resize_to_for_desired_mpp": 224, + } mock_patch_model = _make_mock_model() mock_slide_model = MagicMock() mock_slide_model.get_model_fun.return_value = MagicMock( - return_value=__import__('torch').zeros(1, 512) + return_value=__import__("torch").zeros(1, 512) ) call_log = [] + def factory_side_effect(model_type): call_log.append(model_type) m = MagicMock() m.get_model.return_value = mock_patch_model return m - with patch("mussel.utils.feature_extract.get_model_factory", side_effect=factory_side_effect), patch("mussel.utils.feature_extract.validate_slide_encoder_compatibility"), patch("mussel.utils.feature_extract.WholeSlideImageTileCoordDataset"), patch("mussel.utils.feature_extract._make_dataloader"), patch("mussel.utils.feature_extract.process_dataset") as mock_proc: + with ( + patch( + "mussel.utils.feature_extract.get_model_factory", + side_effect=factory_side_effect, + ), + patch("mussel.utils.feature_extract.validate_slide_encoder_compatibility"), + patch("mussel.utils.feature_extract.WholeSlideImageTileCoordDataset"), + patch("mussel.utils.feature_extract._make_dataloader"), + patch("mussel.utils.feature_extract.process_dataset") as mock_proc, + ): mock_proc.return_value = MagicMock( features=np.zeros((10, 384)), labels=np.zeros(10) ) get_features( - coords, "slide.svs", attrs, + coords, + "slide.svs", + attrs, model=mock_patch_model, use_slide_encoder=True, slide_model_type=ModelType.GIGAPATH_SLIDE, @@ -139,14 +192,20 @@ def factory_side_effect(model_type): # Only the patch encoder factory may be called (for auto-infer check), not the slide encoder slide_encoder_calls = [t for t in call_log if t == ModelType.GIGAPATH_SLIDE] - assert len(slide_encoder_calls) == 0, "Slide encoder factory must not be called when slide_model is provided" + assert ( + len(slide_encoder_calls) == 0 + ), "Slide encoder factory must not be called when slide_model is provided" def test_slide_model_invalid_interface_raises_type_error(): """Passing a slide_model without get_model_fun must raise TypeError.""" coords = np.zeros((5, 2), dtype=np.int32) - attrs = {"patch_size": 256, "patch_level": 0, "mpp": 0.5, - "patch_size_to_resize_to_for_desired_mpp": 224} + attrs = { + "patch_size": 256, + "patch_level": 0, + "mpp": 0.5, + "patch_size_to_resize_to_for_desired_mpp": 224, + } bad_slide_model = object() with pytest.raises(TypeError, match="get_model_fun"): @@ -155,21 +214,37 @@ def test_slide_model_invalid_interface_raises_type_error(): # -- compatibility validation ----b??----------------------------------------- + def test_compatibility_validated_with_preloaded_patch_model(): """validate_slide_encoder_compatibility must be called even with pre-loaded patch model.""" coords = np.zeros((5, 2), dtype=np.int32) - attrs = {"patch_size": 256, "patch_level": 0, "mpp": 0.5, - "patch_size_to_resize_to_for_desired_mpp": 224} + attrs = { + "patch_size": 256, + "patch_level": 0, + "mpp": 0.5, + "patch_size_to_resize_to_for_desired_mpp": 224, + } mock_patch_model = _make_mock_model() - with patch("mussel.utils.feature_extract.validate_slide_encoder_compatibility") as mock_validate, patch("mussel.utils.feature_extract.WholeSlideImageTileCoordDataset"), patch("mussel.utils.feature_extract._make_dataloader"), patch("mussel.utils.feature_extract.process_dataset") as mock_proc: + with ( + patch( + "mussel.utils.feature_extract.validate_slide_encoder_compatibility" + ) as mock_validate, + patch("mussel.utils.feature_extract.WholeSlideImageTileCoordDataset"), + patch("mussel.utils.feature_extract._make_dataloader"), + patch("mussel.utils.feature_extract.process_dataset") as mock_proc, + ): mock_proc.return_value = MagicMock( features=np.zeros((5, 384)), labels=np.zeros(5) ) mock_slide = MagicMock() - mock_slide.get_model_fun.return_value = MagicMock(return_value=__import__('torch').zeros(1, 512)) + mock_slide.get_model_fun.return_value = MagicMock( + return_value=__import__("torch").zeros(1, 512) + ) get_features( - coords, "slide.svs", attrs, + coords, + "slide.svs", + attrs, model_type=ModelType.GIGAPATH, model=mock_patch_model, use_slide_encoder=True, @@ -334,14 +409,17 @@ def mock_model_fun(batch): def test_extract_features_config_has_embedding_precision_field(): - """ExtractFeaturesConfig must include embedding_precision defaulting to float32.""" + """ExtractFeaturesConfig must include precision and model kwargs defaults.""" cfg = ExtractFeaturesConfig() assert hasattr(cfg, "embedding_precision") assert cfg.embedding_precision == "float32" + assert cfg.model_kwargs == {} + assert cfg.slide_model_kwargs == {} # -- slide pipeline precision semantics ---------------------------------------- + def test_aggregate_slide_features_precision(): """aggregate_slide_features saves output at the requested precision.""" import h5py @@ -369,13 +447,13 @@ def test_aggregate_slide_features_precision(): ) with h5py.File(out_h5, "r") as f: dtype = f["features"].dtype - assert dtype.itemsize == expected_itemsize, ( - f"precision={precision}: expected itemsize {expected_itemsize}, got {dtype.itemsize}" - ) + assert ( + dtype.itemsize == expected_itemsize + ), f"precision={precision}: expected itemsize {expected_itemsize}, got {dtype.itemsize}" if expected_kind != "V": - assert dtype.kind == expected_kind, ( - f"precision={precision}: expected kind {expected_kind!r}, got {dtype.kind!r}" - ) + assert ( + dtype.kind == expected_kind + ), f"precision={precision}: expected kind {expected_kind!r}, got {dtype.kind!r}" def test_save_features_two_step_keeps_intermediate_float32(): @@ -398,13 +476,27 @@ def test_save_features_two_step_keeps_intermediate_float32(): intermediate_calls = [] - def fake_extract_patch(patch_h5_path, slide_path, output_h5_path, embedding_precision="float32", **kwargs): + def fake_extract_patch( + patch_h5_path, + slide_path, + output_h5_path, + embedding_precision="float32", + **kwargs, + ): intermediate_calls.append(embedding_precision) - with h5py.File(patch_h5_path, "r") as src, h5py.File(output_h5_path, "w") as dst: + with ( + h5py.File(patch_h5_path, "r") as src, + h5py.File(output_h5_path, "w") as dst, + ): dst.create_dataset("features", data=src["features"][:]) dst.create_dataset("coords", data=src["coords"][:]) - def fake_aggregate(patch_features_h5_path, output_h5_path, embedding_precision="float32", **kwargs): + def fake_aggregate( + patch_features_h5_path, + output_h5_path, + embedding_precision="float32", + **kwargs, + ): feature_dtype = _parse_feature_dtype(embedding_precision) with h5py.File(patch_features_h5_path, "r") as src: data = src["features"][:] @@ -417,8 +509,16 @@ def fake_aggregate(patch_features_h5_path, output_h5_path, embedding_precision=" out_h5 = os.path.join(tmpdir, "slide_out.h5") - with patch("mussel.utils.feature_extract.extract_patch_features", side_effect=fake_extract_patch), \ - patch("mussel.utils.feature_extract.aggregate_slide_features", side_effect=fake_aggregate): + with ( + patch( + "mussel.utils.feature_extract.extract_patch_features", + side_effect=fake_extract_patch, + ), + patch( + "mussel.utils.feature_extract.aggregate_slide_features", + side_effect=fake_aggregate, + ), + ): save_features( patch_h5_path=patch_h5, slide_path="dummy.svs", @@ -427,14 +527,120 @@ def fake_aggregate(patch_features_h5_path, output_h5_path, embedding_precision=" embedding_precision="float16", ) - assert intermediate_calls == ["float32"], ( - f"Intermediate tile extraction must use float32, got {intermediate_calls}" - ) + assert intermediate_calls == [ + "float32" + ], f"Intermediate tile extraction must use float32, got {intermediate_calls}" with h5py.File(out_h5, "r") as f: assert f["features"].dtype.itemsize == 2 assert f["features"].dtype.kind == "f" # float16 +def test_save_features_forwards_model_kwargs_to_correct_steps(): + """save_features forwards patch and slide model kwargs to their respective stages.""" + import h5py + import tempfile + from mussel.utils.feature_extract import save_features + + features = np.random.rand(4, 8).astype(np.float32) + coords = np.zeros((4, 2), dtype=np.int32) + + with tempfile.TemporaryDirectory() as tmpdir: + patch_h5 = os.path.join(tmpdir, "patches.h5") + with h5py.File(patch_h5, "w") as f: + f.create_dataset("features", data=features) + f.create_dataset("coords", data=coords) + f["features"].attrs["patch_size"] = 256 + + seen_model_kwargs = [] + seen_slide_model_kwargs = [] + + def fake_extract_patch(patch_h5_path, slide_path, output_h5_path, **kwargs): + seen_model_kwargs.append(kwargs["model_kwargs"]) + with ( + h5py.File(patch_h5_path, "r") as src, + h5py.File(output_h5_path, "w") as dst, + ): + dst.create_dataset("features", data=src["features"][:]) + dst.create_dataset("coords", data=src["coords"][:]) + dst["features"].attrs["patch_size"] = src["features"].attrs[ + "patch_size" + ] + + def fake_aggregate(patch_features_h5_path, output_h5_path, **kwargs): + seen_slide_model_kwargs.append(kwargs["slide_model_kwargs"]) + with ( + h5py.File(patch_features_h5_path, "r") as src, + h5py.File(output_h5_path, "w") as dst, + ): + dst.create_dataset("features", data=src["features"][:]) + return output_h5_path, None + + out_h5 = os.path.join(tmpdir, "slide_out.h5") + with ( + patch( + "mussel.utils.feature_extract.extract_patch_features", + side_effect=fake_extract_patch, + ), + patch( + "mussel.utils.feature_extract.aggregate_slide_features", + side_effect=fake_aggregate, + ), + ): + save_features( + patch_h5_path=patch_h5, + slide_path="dummy.svs", + output_h5_path=out_h5, + aggregation_method="model", + model_kwargs={"patch_arg": "value"}, + slide_model_kwargs={"patch_oom": False}, + ) + + assert seen_model_kwargs == [{"patch_arg": "value"}] + assert seen_slide_model_kwargs == [{"patch_oom": False}] + + +def test_aggregate_slide_features_forwards_slide_model_kwargs(): + """aggregate_slide_features forwards slide_model_kwargs to the model factory.""" + import h5py + import tempfile + import torch + from mussel.utils.feature_extract import aggregate_slide_features + + features = np.random.rand(4, 8).astype(np.float32) + coords = np.zeros((4, 2), dtype=np.int32) + + with tempfile.TemporaryDirectory() as tmpdir: + patch_h5 = os.path.join(tmpdir, "patches.h5") + out_h5 = os.path.join(tmpdir, "slide_out.h5") + with h5py.File(patch_h5, "w") as f: + f.create_dataset("features", data=features) + f.create_dataset("coords", data=coords) + f["features"].attrs["patch_size"] = 256 + + fake_model = MagicMock() + fake_model.get_model_fun.return_value = ( + lambda features, coords, patch_size: torch.zeros(1, 3) + ) + fake_factory = MagicMock() + fake_factory.get_model.return_value = fake_model + + with patch( + "mussel.utils.feature_extract.get_model_factory", return_value=fake_factory + ): + aggregate_slide_features( + patch_features_h5_path=patch_h5, + output_h5_path=out_h5, + aggregation_method="model", + model_type=ModelType.TITAN_SLIDE, + use_gpu=False, + slide_model_kwargs={"patch_oom": False}, + ) + + fake_factory.get_model.assert_called_once_with( + None, False, None, patch_oom=False + ) + + def test_aggregate_slide_features_batch_precision(): """aggregate_slide_features_batch casts output to the requested precision.""" import h5py @@ -463,21 +669,23 @@ def test_aggregate_slide_features_batch_precision(): def test_aggregate_slide_features_config_has_embedding_precision(): - """AggregateSlideFeaturesConfig must expose embedding_precision.""" + """AggregateSlideFeaturesConfig must expose precision and slide model kwargs.""" from mussel.cli.aggregate_slide_features import AggregateSlideFeaturesConfig from omegaconf import OmegaConf cfg = OmegaConf.structured(AggregateSlideFeaturesConfig) assert hasattr(cfg, "embedding_precision") assert cfg.embedding_precision == "float32" + assert cfg.slide_model_kwargs == {} def test_tessellate_extract_features_config_has_embedding_precision(): - """TessellateExtractFeaturesConfig must expose embedding_precision.""" + """TessellateExtractFeaturesConfig must expose precision and model kwargs.""" from mussel.cli.tessellate_extract_features import TessellateExtractFeaturesConfig from mussel.cli.tessellate import SegConfig - from omegaconf import OmegaConf cfg = TessellateExtractFeaturesConfig(seg_config=SegConfig()) assert hasattr(cfg, "embedding_precision") assert cfg.embedding_precision == "float32" + assert cfg.model_kwargs == {} + assert cfg.slide_model_kwargs == {} diff --git a/tests/regression/regression_full_pipeline.py b/tests/regression/regression_full_pipeline.py deleted file mode 100644 index 661af7e4..00000000 --- a/tests/regression/regression_full_pipeline.py +++ /dev/null @@ -1,211 +0,0 @@ -"""Full-pipeline regression: tessellate โ†’ CTransPath โ†’ filter vs REEF reference. - -Pipeline: - 1. Tessellate 948176.svs with same params as reference (patch_size=224, mpp=0.5) - 2. Extract CTransPath features for all tiles - 3. Filter with classifier at threshold 0.75 - 4. Compare filtered coords + features to REEF reference - -Reference: - Filter tiles: /gpfs/cdsi_ess/foundation/reef/filter_tiles/9481/948176.patch.h5 - Features: /gpfs/cdsi_ess/foundation/reef/features/ctranspath/9481/948176.features.pt - -Usage (from repo root, on a GPU node): - uv run python tests/regression/regression_full_pipeline.py -""" - -import pickle -import sys -import tempfile -from pathlib import Path - -import h5py -import numpy as np -import torch - -REPO = Path(__file__).resolve().parents[2] -sys.path.insert(0, str(REPO)) - -from mussel.models.model_factory import ModelType -from mussel.utils import load_classifier -from mussel.utils.feature_extract import (extract_patch_features, - filter_features) - -SLIDE_PATH = REPO / "tests/testdata/948176.svs" -CLASSIFIER_PKL = Path("/gpfs/mskmind_ess/limr/repos/Mussel/model-1727990346535.pkl") -CLASSIFIER_THR = 0.75 -REF_FILTER_H5 = Path("/gpfs/cdsi_ess/foundation/reef/filter_tiles/9481/948176.patch.h5") -REF_FEATURES_PT = Path( - "/gpfs/cdsi_ess/foundation/reef/features/ctranspath/9481/948176.features.pt" -) - - -def tessellate(slide_path: Path, out_h5: str) -> int: - """Tessellate slide with parameters matching the reference pipeline.""" - from mussel.cli.tessellate import SegConfig - from mussel.utils.segment import segment_tissue - - seg_cfg = SegConfig(patch_size=224) # matches CTransPath default / reference H5 - - result = segment_tissue( - slide_path=str(slide_path), - output_h5_path=out_h5, - **{k: v for k, v in vars(seg_cfg).items()}, - ) - if result is None: - raise RuntimeError("segment_tissue returned None โ€” tessellation failed") - _, _, coords, _ = result - return len(coords) - - -def main(): - print(f"Slide: {SLIDE_PATH}") - print(f"Classifier: {CLASSIFIER_PKL} (threshold={CLASSIFIER_THR})") - print(f"Reference: {REF_FEATURES_PT}") - print() - - with tempfile.TemporaryDirectory() as tmpdir: - tmp = Path(tmpdir) - tess_h5 = str(tmp / "tessellate.h5") - feats_h5 = str(tmp / "features.h5") - - # Step 1: Tessellate - print("Step 1/3: Tessellating...") - n_tiles = tessellate(SLIDE_PATH, tess_h5) - with h5py.File(tess_h5, "r") as f: - attrs = dict(f["coords"].attrs) - print( - f" {n_tiles} tiles patch_size={attrs.get('patch_size')} mpp={attrs.get('mpp')}" - ) - - # Step 2: Extract CTransPath features - print("Step 2/3: Extracting CTransPath features...") - extract_patch_features( - patch_h5_path=tess_h5, - slide_path=str(SLIDE_PATH), - output_h5_path=feats_h5, - model_type=ModelType.CTRANSPATH, - batch_size=64, - use_gpu=True, - num_workers=0, - pin_memory=False, - ) - with h5py.File(feats_h5, "r") as f: - all_feats = f["features"][:] - all_coords = f["coords"][:] - print(f" Features: {all_feats.shape}") - - # Step 3: Filter - print(f"Step 3/3: Filtering (threshold={CLASSIFIER_THR})...") - classifier = load_classifier(str(CLASSIFIER_PKL)) - feats_t = torch.from_numpy(all_feats) - filt_feats_t, filt_coords = filter_features( - feats_t, all_coords, classifier, CLASSIFIER_THR - ) - filt_feats = filt_feats_t.numpy() - # filt_coords is already np.ndarray - print( - f" After filter: {len(filt_coords)} tiles (removed {n_tiles - len(filt_coords)})" - ) - - # --- Load reference --- - with h5py.File(REF_FILTER_H5, "r") as f: - ref_coords = f["coords"][:] - ref_feats = torch.load( - REF_FEATURES_PT, map_location="cpu", weights_only=False - ).numpy() - - print() - print("=== Comparison ===") - print(f" Mussel filtered: {filt_feats.shape} coords {filt_coords.shape}") - print(f" Reference: {ref_feats.shape} coords {ref_coords.shape}") - - # --- Coordinate grid analysis --- - mussel_set = set(map(tuple, filt_coords)) - ref_set = set(map(tuple, ref_coords)) - exact_overlap = len(mussel_set & ref_set) - - y0_mussel = filt_coords[:, 1].min() - y0_ref = ref_coords[:, 1].min() - x0_mussel = filt_coords[:, 0].min() - x0_ref = ref_coords[:, 0].min() - print( - f" Grid origin: Mussel xโ‚€={x0_mussel} yโ‚€={y0_mussel} | Ref xโ‚€={x0_ref} yโ‚€={y0_ref} (ฮ”y={y0_ref - y0_mussel})" - ) - print( - f" Exact coord overlap: {exact_overlap} / {len(ref_coords)} reference patches" - ) - - # Bounding-box IoU - m_xmin, m_ymin = filt_coords.min(axis=0) - m_xmax, m_ymax = filt_coords.max(axis=0) - r_xmin, r_ymin = ref_coords.min(axis=0) - r_xmax, r_ymax = ref_coords.max(axis=0) - ix1, iy1 = max(m_xmin, r_xmin), max(m_ymin, r_ymin) - ix2, iy2 = min(m_xmax, r_xmax), min(m_ymax, r_ymax) - inter = max(0, ix2 - ix1) * max(0, iy2 - iy1) - union = ( - (m_xmax - m_xmin) * (m_ymax - m_ymin) - + (r_xmax - r_xmin) * (r_ymax - r_ymin) - - inter - ) - bb_iou = inter / union if union > 0 else 0.0 - print(f" Bounding-box IoU: {bb_iou:.3f}") - - # --- Feature distribution comparison --- - print() - print(" Feature distribution (all filtered patches):") - print( - f" Mussel mean={filt_feats.mean():.5f} std={filt_feats.std():.5f} " - f"min={filt_feats.min():.4f} max={filt_feats.max():.4f}" - ) - print( - f" Ref mean={ref_feats.mean():.5f} std={ref_feats.std():.5f} " - f"min={ref_feats.min():.4f} max={ref_feats.max():.4f}" - ) - - # --- Feature comparison on exactly matching patches --- - if exact_overlap > 0: - mussel_idx = {tuple(c): i for i, c in enumerate(filt_coords)} - ref_idx = {tuple(c): i for i, c in enumerate(ref_coords)} - common = sorted(mussel_set & ref_set) - mi = [mussel_idx[c] for c in common] - ri = [ref_idx[c] for c in common] - am = filt_feats[mi] - ar = ref_feats[ri] - dot = (am * ar).sum(axis=1) - nrms = np.linalg.norm(am, axis=1) * np.linalg.norm(ar, axis=1) - cos = dot / np.clip(nrms, 1e-8, None) - print(f"\n Exact-overlap patches ({exact_overlap}):") - print(f" Cosine sim: mean={cos.mean():.6f} min={cos.min():.6f}") - print(f" Max abs diff: {np.abs(am - ar).max():.6f}") - - # --- Verdict --- - tile_ratio = len(filt_coords) / len(ref_coords) - coord_note = ( - "exact" - if exact_overlap == len(ref_coords) - else f"{exact_overlap}/{len(ref_coords)} patches align" - ) - status = ( - "PASS" - if (0.95 <= tile_ratio <= 1.05 and bb_iou > 0.9) - else ("WARN" if (0.90 <= tile_ratio <= 1.10 and bb_iou > 0.8) else "FAIL") - ) - sym = {"PASS": "โœ…", "WARN": "โš ๏ธ", "FAIL": "โŒ"}[status] - print() - print( - f" โ†’ {sym} {status} tiles={len(filt_coords)}/{len(ref_coords)} ({tile_ratio:.1%})" - f" bb_iou={bb_iou:.3f} coords: {coord_note}" - ) - print() - print(" NOTE: Tile sets may differ due to segmentation differences between Mussel") - print(" and REEF. Feature accuracy for matching patches is validated separately") - print( - " by tests/regression/regression_vs_reference.py (cos=1.000000 for CTransPath)." - ) - sys.exit(0 if status in ("PASS", "WARN") else 1) - - -if __name__ == "__main__": - main() diff --git a/tests/regression/regression_vs_reference.py b/tests/regression/regression_vs_reference.py deleted file mode 100644 index 3605e4fb..00000000 --- a/tests/regression/regression_vs_reference.py +++ /dev/null @@ -1,153 +0,0 @@ -"""Regression check: Mussel features vs reference pipeline output. - -Compares Mussel's OPTIMUS and CTRANSPATH feature extraction against -pre-computed reference features from the REEF pipeline, using the same -reference patch H5 (1675 patches, 223 px at 0.5 ยตm/px โ†’ resized to 224). - -Usage (from repo root, on a GPU node): - uv run python tests/regression/regression_vs_reference.py -""" - -import sys -import tempfile -from pathlib import Path - -import h5py -import numpy as np -import torch - -REPO = Path(__file__).resolve().parents[2] -sys.path.insert(0, str(REPO)) - -from mussel.models.model_factory import ModelType -from mussel.utils.feature_extract import extract_patch_features - -REF_PATCH_H5 = Path("/gpfs/cdsi_ess/foundation/reef/filter_tiles/9481/948176.patch.h5") -SLIDE_PATH = REPO / "tests/testdata/948176.svs" - -MODELS = [ - ( - ModelType.OPTIMUS, - Path("/gpfs/cdsi_ess/foundation/reef/features/optimus/9481/948176.features.pt"), - ), - ( - ModelType.CTRANSPATH, - Path( - "/gpfs/cdsi_ess/foundation/reef/features/ctranspath/9481/948176.features.pt" - ), - ), -] - - -def _run_model(model_type: ModelType, ref_feat_path: Path) -> dict: - print(f"\n{'='*60}") - print(f"Model: {model_type.name}") - print(f"Ref: {ref_feat_path}") - - with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as tmp: - out_h5 = tmp.name - - extract_patch_features( - patch_h5_path=str(REF_PATCH_H5), - slide_path=str(SLIDE_PATH), - output_h5_path=out_h5, - model_type=model_type, - batch_size=64, - use_gpu=True, - num_workers=0, - pin_memory=False, - ) - - with h5py.File(out_h5, "r") as f: - mussel_feats = f["features"][:] - mussel_coords = f["coords"][:] - - ref_feats = torch.load(ref_feat_path, map_location="cpu", weights_only=False) - if isinstance(ref_feats, dict): - ref_feats = next(iter(ref_feats.values())) - ref_feats = ref_feats.numpy() - - with h5py.File(REF_PATCH_H5, "r") as f: - ref_coords = f["coords"][:] - - print(f" Mussel output: {mussel_feats.shape} {mussel_feats.dtype}") - print(f" Reference: {ref_feats.shape} {ref_feats.dtype}") - - if mussel_feats.shape != ref_feats.shape: - print(f" SHAPE MISMATCH โŒ") - return {"model": model_type.name, "status": "SHAPE_MISMATCH"} - - if not np.array_equal(mussel_coords, ref_coords): - print(f" COORD MISMATCH โŒ โ€” patches not aligned") - return {"model": model_type.name, "status": "COORD_MISMATCH"} - - # --- Metrics --- - dot = (mussel_feats * ref_feats).sum(axis=1) - norms = np.linalg.norm(mussel_feats, axis=1) * np.linalg.norm(ref_feats, axis=1) - cos = dot / np.clip(norms, 1e-8, None) - - l2 = np.linalg.norm(mussel_feats - ref_feats, axis=1) - max_absdiff = np.abs(mussel_feats - ref_feats).max() - - close_tight = np.allclose(mussel_feats, ref_feats, rtol=1e-3, atol=1e-4) - close_loose = np.allclose(mussel_feats, ref_feats, rtol=1e-2, atol=1e-3) - - print( - f" Cosine sim: mean={cos.mean():.6f} min={cos.min():.6f} p5={np.percentile(cos, 5):.6f}" - ) - print(f" L2 distance: mean={l2.mean():.5f} max={l2.max():.5f}") - print(f" Max abs diff: {max_absdiff:.6f}") - print(f" allclose(rtol=1e-3, atol=1e-4): {close_tight}") - print(f" allclose(rtol=1e-2, atol=1e-3): {close_loose}") - - if cos.mean() > 0.999: - status = "PASS" - print(f" โ†’ PASS โœ… mean cosine={cos.mean():.6f} > 0.999") - elif cos.mean() > 0.99: - status = "WARN" - print(f" โ†’ WARN โš ๏ธ mean cosine={cos.mean():.6f} in [0.99, 0.999]") - else: - status = "FAIL" - print(f" โ†’ FAIL โŒ mean cosine={cos.mean():.6f} < 0.99") - - return { - "model": model_type.name, - "status": status, - "n": len(cos), - "cos_mean": float(cos.mean()), - "cos_min": float(cos.min()), - "cos_p5": float(np.percentile(cos, 5)), - "l2_mean": float(l2.mean()), - "l2_max": float(l2.max()), - "max_absdiff": float(max_absdiff), - } - - -def main(): - print(f"Slide: {SLIDE_PATH}") - print(f"Patch H5: {REF_PATCH_H5}") - - results = [] - for model_type, ref_path in MODELS: - r = _run_model(model_type, ref_path) - results.append(r) - - print(f"\n{'='*60}") - print("SUMMARY") - print(f"{'='*60}") - all_pass = True - for r in results: - status_sym = {"PASS": "โœ…", "WARN": "โš ๏ธ", "FAIL": "โŒ"}.get(r["status"], "โ“") - cos = r.get("cos_mean", float("nan")) - print( - f" {status_sym} {r['model']:20s} cos_mean={cos:.6f} status={r['status']}" - ) - if r["status"] not in ("PASS", "WARN"): - all_pass = False - - print() - sys.exit(0 if all_pass else 1) - - -if __name__ == "__main__": - main() diff --git a/tests/testdata/snapshots/CONCH1_5.npy b/tests/testdata/snapshots/CONCH1_5.npy index 68069674..1ea6e6f1 100644 Binary files a/tests/testdata/snapshots/CONCH1_5.npy and b/tests/testdata/snapshots/CONCH1_5.npy differ diff --git a/tests/testdata/snapshots/TITAN_SLIDE.npy b/tests/testdata/snapshots/TITAN_SLIDE.npy new file mode 100644 index 00000000..cfb120c4 Binary files /dev/null and b/tests/testdata/snapshots/TITAN_SLIDE.npy differ