wsi-toolbox 0.2.0__py3-none-any.whl

wsi_toolbox/commands/patch_embedding.py

@@ -0,0 +1,196 @@
+"""
+Patch embedding extraction command
+"""
+
+import gc
+
+import h5py
+import numpy as np
+from pydantic import BaseModel
+
+from ..common import create_default_model
+from ..utils import safe_del
+from . import _get, _progress, get_config
+
+
+class PatchEmbeddingResult(BaseModel):
+    """Result of patch embedding extraction"""
+
+    feature_dim: int = 0
+    patch_count: int = 0
+    model: str = ""
+    with_latent: bool = False
+    skipped: bool = False
+
+
+class PatchEmbeddingCommand:
+    """
+    Extract embeddings from patches using foundation models
+
+    Usage:
+        # Set global config once
+        commands.set_default_model_preset('gigapath')
+        commands.set_default_device('cuda')
+
+        # Create and run command
+        cmd = PatchEmbeddingCommand(batch_size=256, with_latent=False)
+        result = cmd(hdf5_path='data.h5')
+    """
+
+    def __init__(
+        self,
+        batch_size: int = 256,
+        with_latent: bool = False,
+        overwrite: bool = False,
+        model_name: str | None = None,
+        device: str | None = None,
+    ):
+        """
+        Initialize patch embedding extractor
+
+        Args:
+            batch_size: Batch size for inference
+            with_latent: Whether to extract latent features
+            overwrite: Whether to overwrite existing features
+            model_name: Model name (None to use global default)
+            device: Device (None to use global default)
+
+        Note:
+            progress and verbose are controlled by global config
+        """
+        self.batch_size = batch_size
+        self.with_latent = with_latent
+        self.overwrite = overwrite
+        self.model_name = _get("model_name", model_name)
+        self.device = _get("device", device)
+
+        # Validate model
+        if self.model_name not in ["uni", "gigapath", "virchow2"]:
+            raise ValueError(f"Invalid model: {self.model_name}")
+
+        # Dataset paths
+        self.feature_name = f"{self.model_name}/features"
+        self.latent_feature_name = f"{self.model_name}/latent_features"
+
+    def __call__(self, hdf5_path: str) -> PatchEmbeddingResult:
+        """
+        Execute embedding extraction
+
+        Args:
+            hdf5_path: Path to HDF5 file
+
+        Returns:
+            PatchEmbeddingResult: Result metadata (feature_dim, patch_count, skipped, etc.)
+        """
+        # Lazy import: torch is slow to load (~800ms), defer until needed
+        import torch  # noqa: PLC0415
+
+        # Load model (uses globally registered model generator)
+        model = create_default_model()
+        model = model.eval().to(self.device)
+
+        # Normalization parameters
+        mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(self.device)
+        std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(self.device)
+
+        done = False
+
+        try:
+            with h5py.File(hdf5_path, "r+") as f:
+                latent_size = model.patch_embed.proj.kernel_size[0]
+
+                # Check if already exists
+                if not self.overwrite:
+                    if self.with_latent:
+                        if (self.feature_name in f) and (self.latent_feature_name in f):
+                            if get_config().verbose:
+                                print("Already extracted. Skipped.")
+                            done = True
+                            return PatchEmbeddingResult(skipped=True)
+                        if (self.feature_name in f) or (self.latent_feature_name in f):
+                            raise RuntimeError(f"Either {self.feature_name} or {self.latent_feature_name} exists.")
+                    else:
+                        if self.feature_name in f:
+                            if get_config().verbose:
+                                print("Already extracted. Skipped.")
+                            done = True
+                            return PatchEmbeddingResult(skipped=True)
+
+                # Delete if overwrite
+                if self.overwrite:
+                    safe_del(f, self.feature_name)
+                    safe_del(f, self.latent_feature_name)
+
+                # Get patch count
+                patch_count = f["metadata/patch_count"][()]
+
+                # Create batch indices
+                batch_idx = [(i, min(i + self.batch_size, patch_count)) for i in range(0, patch_count, self.batch_size)]
+
+                # Create datasets
+                f.create_dataset(self.feature_name, shape=(patch_count, model.num_features), dtype=np.float32)
+                if self.with_latent:
+                    f.create_dataset(
+                        self.latent_feature_name,
+                        shape=(patch_count, latent_size**2, model.num_features),
+                        dtype=np.float16,
+                    )
+
+                # Process batches
+                tq = _progress(batch_idx)
+                for i0, i1 in tq:
+                    # Load batch
+                    x = f["patches"][i0:i1]
+                    x = (torch.from_numpy(x) / 255).permute(0, 3, 1, 2)  # BHWC->BCHW
+                    x = x.to(self.device)
+                    x = (x - mean) / std
+
+                    # Forward pass
+                    with torch.inference_mode(), torch.autocast(device_type="cuda", dtype=torch.float16):
+                        h_tensor = model.forward_features(x)
+
+                    # Extract features
+                    h = h_tensor.cpu().detach().numpy()  # [B, T+L, H]
+                    latent_index = h.shape[1] - latent_size**2
+                    cls_feature = h[:, 0, ...]
+                    latent_feature = h[:, latent_index:, ...]
+
+                    # Save features
+                    f[self.feature_name][i0:i1] = cls_feature
+                    if self.with_latent:
+                        f[self.latent_feature_name][i0:i1] = latent_feature.astype(np.float16)
+
+                    # Cleanup
+                    del x, h_tensor
+                    torch.cuda.empty_cache()
+
+                    tq.set_description(f"Processing {i0}-{i1} (total={patch_count})")
+                    tq.refresh()
+
+                if get_config().verbose:
+                    print(f"Embeddings dimension: {f[self.feature_name].shape}")
+
+                done = True
+                return PatchEmbeddingResult(
+                    feature_dim=model.num_features,
+                    patch_count=patch_count,
+                    model=self.model_name,
+                    with_latent=self.with_latent,
+                )
+
+        finally:
+            if done and get_config().verbose:
+                print(f"Wrote {self.feature_name}")
+            elif not done:
+                # Cleanup on error
+                with h5py.File(hdf5_path, "a") as f:
+                    safe_del(f, self.feature_name)
+                    if self.with_latent:
+                        safe_del(f, self.latent_feature_name)
+                if get_config().verbose:
+                    print(f"ABORTED! Deleted {self.feature_name}")
+
+            # Cleanup
+            del model, mean, std
+            torch.cuda.empty_cache()
+            gc.collect()

wsi_toolbox/commands/pca.py

@@ -0,0 +1,206 @@
+"""
+PCA scoring command for feature analysis
+"""
+
+import h5py
+import numpy as np
+from pydantic import BaseModel
+
+from ..utils.hdf5_paths import build_cluster_path, build_namespace, ensure_groups
+from . import _get, _progress, get_config
+from .data_loader import MultipleContext
+
+
+def sigmoid(x):
+    """Apply sigmoid function"""
+    return 1 / (1 + np.exp(-x))
+
+
+class PCAResult(BaseModel):
+    """Result of PCA operation"""
+
+    n_samples: int
+    n_components: int
+    namespace: str
+    target_path: str
+    skipped: bool = False
+
+
+class PCACommand:
+    """
+    Compute PCA scores from features
+
+    Input:
+        - features (from <model>/features)
+        - namespace + filters (recursive hierarchy)
+        - n_components: 1, 2, or 3
+        - scaler: minmax or std
+
+    Output:
+        - PCA scores written to deepest level
+        - metadata saved as HDF5 attributes
+
+    Example hierarchy:
+        uni/default/pca2
+            ↑ with attributes: n_components=2, scaler="minmax"
+        uni/default/filter/1+2+3/pca1
+            ↑ filtered, with PCA scores
+
+    Usage:
+        # Basic PCA
+        cmd = PCACommand(n_components=2)
+        result = cmd('data.h5')  # → uni/default/pca2
+
+        # Filtered PCA
+        cmd = PCACommand(parent_filters=[[1,2,3]])
+        result = cmd('data.h5')  # → uni/default/filter/1+2+3/pca2
+    """
+
+    def __init__(
+        self,
+        n_components: int = 2,
+        namespace: str | None = None,
+        parent_filters: list[list[int]] | None = None,
+        scaler: str = "minmax",
+        overwrite: bool = False,
+        model_name: str | None = None,
+    ):
+        """
+        Args:
+            n_components: Number of PCA components (1, 2, or 3)
+            namespace: Explicit namespace (None = auto-generate)
+            parent_filters: Hierarchical filters, e.g., [[1,2,3], [4,5]]
+            scaler: Scaling method ("minmax" or "std")
+            overwrite: Overwrite existing PCA scores
+            model_name: Model name (None = use global default)
+        """
+        self.n_components = n_components
+        self.namespace = namespace
+        self.parent_filters = parent_filters or []
+        self.scaler = scaler
+        self.overwrite = overwrite
+        self.model_name = _get("model_name", model_name)
+
+        # Validate
+        if self.model_name not in ["uni", "gigapath", "virchow2"]:
+            raise ValueError(f"Invalid model: {self.model_name}")
+        if self.n_components not in [1, 2, 3]:
+            raise ValueError(f"Invalid n_components: {self.n_components}")
+        if self.scaler not in ["minmax", "std"]:
+            raise ValueError(f"Invalid scaler: {self.scaler}")
+
+        # Internal state
+        self.hdf5_paths = []
+        self.pca_scores = None
+
+    def __call__(self, hdf5_paths: str | list[str]) -> PCAResult:
+        """
+        Execute PCA computation
+
+        Args:
+            hdf5_paths: Single HDF5 path or list of paths
+
+        Returns:
+            PCAResult
+        """
+        # Normalize to list
+        if isinstance(hdf5_paths, str):
+            hdf5_paths = [hdf5_paths]
+        self.hdf5_paths = hdf5_paths
+
+        # Determine namespace
+        if self.namespace is None:
+            self.namespace = build_namespace(hdf5_paths)
+        elif "+" in self.namespace:
+            raise ValueError("Namespace cannot contain '+' (reserved for multi-file auto-generated namespaces)")
+
+        # Build target path
+        target_path = build_cluster_path(
+            self.model_name, self.namespace, filters=self.parent_filters, dataset=f"pca{self.n_components}"
+        )
+
+        # Check if already exists
+        if not self.overwrite:
+            with h5py.File(hdf5_paths[0], "r") as f:
+                if target_path in f:
+                    scores = f[target_path][:]
+                    n_samples = np.sum(~np.isnan(scores[:, 0]) if scores.ndim > 1 else ~np.isnan(scores))
+                    if get_config().verbose:
+                        print(f"PCA scores already exist at {target_path}")
+                    return PCAResult(
+                        n_samples=n_samples,
+                        n_components=self.n_components,
+                        namespace=self.namespace,
+                        target_path=target_path,
+                        skipped=True,
+                    )
+
+        # Execute with progress tracking
+        with _progress(total=3, desc="PCA") as pbar:
+            # Load data
+            pbar.set_description("Loading features")
+            ctx = MultipleContext(hdf5_paths, self.model_name, self.namespace, self.parent_filters)
+            features = ctx.load_features(source="features")
+            pbar.update(1)
+
+            # Compute PCA
+            pbar.set_description("Computing PCA")
+            self.pca_scores = self._compute_pca(features)
+            pbar.update(1)
+
+            # Write results
+            pbar.set_description("Writing results")
+            self._write_results(ctx, target_path)
+            pbar.update(1)
+
+        # Verbose output after progress bar closes
+        if get_config().verbose:
+            print(f"Computed PCA: {len(features)} samples → {self.n_components}D")
+            print(f"Wrote {target_path} to {len(hdf5_paths)} file(s)")
+
+        return PCAResult(
+            n_samples=len(features), n_components=self.n_components, namespace=self.namespace, target_path=target_path
+        )
+
+    def _compute_pca(self, features: np.ndarray) -> np.ndarray:
+        """Compute PCA and apply scaling"""
+        # Lazy import: sklearn is slow to load (~600ms), defer until needed
+        from sklearn.decomposition import PCA  # noqa: PLC0415
+        from sklearn.preprocessing import MinMaxScaler, StandardScaler  # noqa: PLC0415
+
+        pca = PCA(n_components=self.n_components)
+        pca_values = pca.fit_transform(features)
+
+        if self.scaler == "minmax":
+            scaler = MinMaxScaler()
+            pca_values = scaler.fit_transform(pca_values)
+        elif self.scaler == "std":
+            scaler = StandardScaler()
+            pca_values = scaler.fit_transform(pca_values)
+            pca_values = sigmoid(pca_values)
+
+        return pca_values
+
+    def _write_results(self, ctx: MultipleContext, target_path: str):
+        """Write PCA results to HDF5 files"""
+        for file_slice in ctx:
+            file_scores = file_slice.slice(self.pca_scores)
+
+            with h5py.File(file_slice.hdf5_path, "a") as f:
+                ensure_groups(f, target_path)
+
+                if target_path in f:
+                    del f[target_path]
+
+                # Fill with NaN for filtered patches
+                if self.n_components == 1:
+                    full_scores = np.full(len(file_slice.mask), np.nan)
+                    full_scores[file_slice.mask] = file_scores.flatten()
+                else:
+                    full_scores = np.full((len(file_slice.mask), self.n_components), np.nan)
+                    full_scores[file_slice.mask] = file_scores
+
+                ds = f.create_dataset(target_path, data=full_scores)
+                ds.attrs["n_components"] = self.n_components
+                ds.attrs["scaler"] = self.scaler
+                ds.attrs["model"] = self.model_name