wsi-toolbox 0.1.0__py3-none-any.whl

wsi_toolbox/commands/clustering.py

@@ -0,0 +1,214 @@
+"""
+Clustering command for WSI features
+"""
+
+import h5py
+import numpy as np
+from pydantic import BaseModel
+from sklearn.preprocessing import StandardScaler
+
+from ..utils.analysis import leiden_cluster
+from . import _config, _get
+
+
+class ClusteringResult(BaseModel):
+    """Result of clustering operation"""
+    cluster_count: int
+    feature_count: int
+    target_path: str
+    skipped: bool = False
+
+
+class ClusteringCommand:
+    """
+    Perform clustering on extracted features
+
+    Usage:
+        # Set global config once
+        commands.set_default_model('gigapath')
+
+        # Create and run command
+        cmd = ClusteringCommand(resolution=1.0)
+        result = cmd(hdf5_paths=['data1.h5', 'data2.h5'])
+    """
+
+    def __init__(self,
+                 resolution: float = 1.0,
+                 cluster_name: str = '',
+                 cluster_filter: list[int] | None = None,
+                 use_umap: bool = False,
+                 overwrite: bool = False,
+                 model_name: str | None = None):
+        """
+        Initialize clustering command
+
+        Args:
+            resolution: Leiden clustering resolution
+            cluster_name: Name for multi-file clustering
+            cluster_filter: Filter to specific clusters for sub-clustering
+            use_umap: Whether to use UMAP embeddings for clustering
+            overwrite: Whether to overwrite existing clusters
+            model_name: Model name (None to use global default)
+        """
+        self.resolution = resolution
+        self.cluster_name = cluster_name
+        self.cluster_filter = cluster_filter or []
+        self.use_umap = use_umap
+        self.overwrite = overwrite
+        self.model_name = _get('model_name', model_name)
+
+        # Validate model
+        if self.model_name not in ['uni', 'gigapath', 'virchow2']:
+            raise ValueError(f'Invalid model: {self.model_name}')
+
+        self.sub_clustering = len(self.cluster_filter) > 0
+
+        # Internal state (initialized in __call__)
+        self.hdf5_paths = []
+        self.masks = []
+        self.features = None
+        self.total_clusters = None
+        self.umap_embeddings = None
+
+    def __call__(self, hdf5_paths: str | list[str]) -> ClusteringResult:
+        """
+        Execute clustering
+
+        Args:
+            hdf5_paths: Single HDF5 path or list of paths
+
+        Returns:
+            ClusteringResult: Result metadata (cluster_count, etc.)
+        """
+        # Normalize to list
+        if isinstance(hdf5_paths, str):
+            hdf5_paths = [hdf5_paths]
+
+        self.hdf5_paths = hdf5_paths
+        multi = len(hdf5_paths) > 1
+
+        # Validate multi-file clustering
+        if multi and not self.cluster_name:
+            raise RuntimeError('Multiple files provided but cluster_name was not specified.')
+
+        # Determine cluster path
+        if multi:
+            clusters_path = f'{self.model_name}/clusters_{self.cluster_name}'
+        else:
+            clusters_path = f'{self.model_name}/clusters'
+
+        # Load features
+        self._load_features(clusters_path)
+
+        # Check if already exists
+        if not self.sub_clustering and hasattr(self, 'has_clusters') and self.has_clusters and not self.overwrite:
+            if _config.verbose:
+                print('Skip clustering (already exists)')
+            return ClusteringResult(
+                cluster_count=len(np.unique(self.total_clusters)),
+                feature_count=len(self.features),
+                target_path=clusters_path,
+                skipped=True
+            )
+
+        # Perform clustering
+        self.total_clusters = leiden_cluster(
+            self.features,
+            umap_emb_func=self.get_umap_embeddings if self.use_umap else None,
+            resolution=self.resolution,
+            progress=_config.progress
+        )
+
+        # Write results
+        target_path = clusters_path
+        if self.sub_clustering:
+            suffix = '_sub' + '-'.join(map(str, self.cluster_filter))
+            target_path = target_path + suffix
+
+        if _config.verbose:
+            print(f'Writing to {target_path}')
+
+        cursor = 0
+        for hdf5_path, mask in zip(self.hdf5_paths, self.masks):
+            count = np.sum(mask)
+            clusters = self.total_clusters[cursor:cursor + count]
+            cursor += count
+
+            with h5py.File(hdf5_path, 'a') as f:
+                if target_path in f:
+                    del f[target_path]
+
+                # Fill with -1 for filtered patches
+                full_clusters = np.full(len(mask), -1, dtype=clusters.dtype)
+                full_clusters[mask] = clusters
+                f.create_dataset(target_path, data=full_clusters)
+
+        cluster_count = len(np.unique(self.total_clusters))
+
+        return ClusteringResult(
+            cluster_count=cluster_count,
+            feature_count=len(self.features),
+            target_path=target_path
+        )
+
+    def _load_features(self, clusters_path: str):
+        """Load features from HDF5 files"""
+        featuress = []
+        clusterss = []
+        self.masks = []
+
+        for hdf5_path in self.hdf5_paths:
+            with h5py.File(hdf5_path, 'r') as f:
+                patch_count = f['metadata/patch_count'][()]
+
+                # Check existing clusters
+                if clusters_path in f:
+                    clusters = f[clusters_path][:]
+                else:
+                    clusters = None
+
+                # Create mask
+                if self.cluster_filter:
+                    if clusters is None:
+                        raise RuntimeError('Sub-clustering requires pre-computed clusters')
+                    mask = np.isin(clusters, self.cluster_filter)
+                else:
+                    mask = np.ones(patch_count, dtype=bool)
+
+                self.masks.append(mask)
+
+                # Load features
+                feature_path = f'{self.model_name}/features'
+                features = f[feature_path][mask]
+                featuress.append(features)
+
+                # Store existing clusters
+                if clusters is not None:
+                    clusterss.append(clusters[mask])
+
+        # Concatenate and normalize
+        features = np.concatenate(featuress)
+        scaler = StandardScaler()
+        self.features = scaler.fit_transform(features)
+
+        # Store existing clusters state
+        if len(clusterss) == len(self.hdf5_paths):
+            self.has_clusters = True
+            self.total_clusters = np.concatenate(clusterss)
+        elif len(clusterss) == 0:
+            self.has_clusters = False
+            self.total_clusters = None
+        else:
+            raise RuntimeError(
+                f'Cluster count mismatch: {len(clusterss)} vs {len(self.hdf5_paths)}'
+            )
+
+    def get_umap_embeddings(self):
+        import umap
+        """Get UMAP embeddings (lazy evaluation)"""
+        if self.umap_embeddings is not None:
+            return self.umap_embeddings
+
+        reducer = umap.UMAP(n_components=2)
+        self.umap_embeddings = reducer.fit_transform(self.features)
+        return self.umap_embeddings

wsi_toolbox/commands/dzi_export.py

@@ -0,0 +1,202 @@
+"""
+DZI export command for Deep Zoom Image format
+"""
+
+import math
+import shutil
+from pathlib import Path
+
+import h5py
+import numpy as np
+from PIL import Image
+
+from . import _config, _progress
+
+
+class DziExportCommand:
+    """
+    Export HDF5 patches to DZI (Deep Zoom Image) format
+
+    Usage:
+        cmd = DziExportCommand(jpeg_quality=90, fill_empty=False)
+        cmd(hdf5_path='data.h5', output_dir='output', name='slide')
+    """
+
+    def __init__(self,
+                 jpeg_quality: int = 90,
+                 fill_empty: bool = False):
+        """
+        Initialize DZI export command
+
+        Args:
+            jpeg_quality: JPEG compression quality (0-100)
+            fill_empty: Fill missing tiles with black tiles
+        """
+        self.jpeg_quality = jpeg_quality
+        self.fill_empty = fill_empty
+
+    def __call__(self, hdf5_path: str, output_dir: str, name: str) -> dict:
+        """
+        Export to DZI format with full pyramid
+
+        Args:
+            hdf5_path: Path to HDF5 file
+            output_dir: Output directory
+            name: Base name for DZI files
+
+        Returns:
+            dict: Export metadata
+        """
+        output_dir = Path(output_dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        # Read HDF5
+        with h5py.File(hdf5_path, 'r') as f:
+            patches = f['patches'][:]
+            coords = f['coordinates'][:]
+            original_width = f['metadata/original_width'][()]
+            original_height = f['metadata/original_height'][()]
+            tile_size = f['metadata/patch_size'][()]
+
+        # Validate tile_size (256 or 512 only)
+        if tile_size not in [256, 512]:
+            raise ValueError(f'Unsupported patch_size: {tile_size}. Only 256 or 512 are supported.')
+
+        # Calculate grid and levels
+        cols = (original_width + tile_size - 1) // tile_size
+        rows = (original_height + tile_size - 1) // tile_size
+        max_dimension = max(original_width, original_height)
+        max_level = math.ceil(math.log2(max_dimension))
+
+        if _config.verbose:
+            print(f'Original size: {original_width}x{original_height}')
+            print(f'Tile size: {tile_size}')
+            print(f'Grid: {cols}x{rows}')
+            print(f'Total patches in HDF5: {len(patches)}')
+            print(f'Max zoom level: {max_level} (Level 0 = 1x1, Level {max_level} = original)')
+
+        coord_to_idx = {(int(x // tile_size), int(y // tile_size)): idx
+                        for idx, (x, y) in enumerate(coords)}
+
+        # Setup directories
+        dzi_path = output_dir / f'{name}.dzi'
+        files_dir = output_dir / f'{name}_files'
+        files_dir.mkdir(exist_ok=True)
+
+        # Create empty tile template for current tile_size
+        empty_tile_path = None
+        if self.fill_empty:
+            empty_tile_path = files_dir / '_empty.jpeg'
+            black_img = Image.fromarray(np.zeros((tile_size, tile_size, 3), dtype=np.uint8))
+            black_img.save(empty_tile_path, 'JPEG', quality=self.jpeg_quality)
+
+        # Export max level (original patches from HDF5)
+        level_dir = files_dir / str(max_level)
+        level_dir.mkdir(exist_ok=True)
+
+        tq = _progress(range(rows))
+        for row in tq:
+            tq.set_description(f'Exporting level {max_level}: row {row+1}/{rows}')
+            for col in range(cols):
+                tile_path = level_dir / f'{col}_{row}.jpeg'
+                if (col, row) in coord_to_idx:
+                    idx = coord_to_idx[(col, row)]
+                    patch = patches[idx]
+                    img = Image.fromarray(patch)
+                    img.save(tile_path, 'JPEG', quality=self.jpeg_quality)
+                elif self.fill_empty:
+                    shutil.copyfile(empty_tile_path, tile_path)
+
+        # Generate lower levels by downsampling
+        for level in range(max_level - 1, -1, -1):
+            if _config.verbose:
+                print(f'Generating level {level}...')
+            self._generate_zoom_level_down(
+                files_dir, level, max_level, original_width, original_height,
+                tile_size, empty_tile_path
+            )
+
+        # Generate DZI XML
+        self._generate_dzi_xml(dzi_path, original_width, original_height, tile_size)
+
+        if _config.verbose:
+            print(f'DZI export complete: {dzi_path}')
+
+        return {
+            'dzi_path': str(dzi_path),
+            'max_level': max_level,
+            'tile_size': tile_size,
+            'grid': f'{cols}x{rows}'
+        }
+
+    def _generate_zoom_level_down(self,
+                                   files_dir: Path,
+                                   curr_level: int,
+                                   max_level: int,
+                                   original_width: int,
+                                   original_height: int,
+                                   tile_size: int,
+                                   empty_tile_path: Path | None):
+        """Generate a zoom level by downsampling from the higher level"""
+        src_level = curr_level + 1
+        src_dir = files_dir / str(src_level)
+        curr_dir = files_dir / str(curr_level)
+        curr_dir.mkdir(exist_ok=True)
+
+        # Calculate dimensions at each level
+        curr_scale = 2 ** (max_level - curr_level)
+        curr_width = math.ceil(original_width / curr_scale)
+        curr_height = math.ceil(original_height / curr_scale)
+        curr_cols = math.ceil(curr_width / tile_size)
+        curr_rows = math.ceil(curr_height / tile_size)
+
+        src_scale = 2 ** (max_level - src_level)
+        src_width = math.ceil(original_width / src_scale)
+        src_height = math.ceil(original_height / src_scale)
+        src_cols = math.ceil(src_width / tile_size)
+        src_rows = math.ceil(src_height / tile_size)
+
+        tq = _progress(range(curr_rows))
+        for row in tq:
+            for col in range(curr_cols):
+                # Combine 4 tiles from source level
+                combined = np.zeros((tile_size * 2, tile_size * 2, 3), dtype=np.uint8)
+                has_any_tile = False
+
+                for dy in range(2):
+                    for dx in range(2):
+                        src_col = col * 2 + dx
+                        src_row = row * 2 + dy
+
+                        if src_col < src_cols and src_row < src_rows:
+                            src_path = src_dir / f'{src_col}_{src_row}.jpeg'
+                            if src_path.exists():
+                                src_img = Image.open(src_path)
+                                src_array = np.array(src_img)
+                                h, w = src_array.shape[:2]
+                                combined[dy*tile_size:dy*tile_size+h,
+                                        dx*tile_size:dx*tile_size+w] = src_array
+                                has_any_tile = True
+
+                tile_path = curr_dir / f'{col}_{row}.jpeg'
+                if has_any_tile:
+                    combined_img = Image.fromarray(combined)
+                    downsampled = combined_img.resize((tile_size, tile_size), Image.LANCZOS)
+                    downsampled.save(tile_path, 'JPEG', quality=self.jpeg_quality)
+                elif self.fill_empty and empty_tile_path:
+                    shutil.copyfile(empty_tile_path, tile_path)
+
+            tq.set_description(f'Generating level {curr_level}: row {row+1}/{curr_rows}')
+
+    def _generate_dzi_xml(self, dzi_path: Path, width: int, height: int, tile_size: int):
+        """Generate DZI XML file"""
+        dzi_content = f'''<?xml version="1.0" encoding="utf-8"?>
+<Image xmlns="http://schemas.microsoft.com/deepzoom/2008"
+       Format="jpeg"
+       Overlap="0"
+       TileSize="{tile_size}">
+    <Size Width="{width}" Height="{height}"/>
+</Image>
+'''
+        with open(dzi_path, 'w', encoding='utf-8') as f:
+            f.write(dzi_content)

wsi_toolbox/commands/patch_embedding.py

@@ -0,0 +1,199 @@
+"""
+Patch embedding extraction command
+"""
+
+import gc
+
+import h5py
+import numpy as np
+import torch
+from pydantic import BaseModel
+
+from ..models import create_model
+from ..utils.helpers import safe_del
+from . import _config, _get, _progress
+
+
+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('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.)
+        """
+        # Load model
+        model = create_model(self.model_name)
+        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 _config.verbose:
+                                print('Already extracted. Skipped.')
+                            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 _config.verbose:
+                                print('Already extracted. Skipped.')
+                            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 _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 _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 _config.verbose:
+                    print(f'ABORTED! Deleted {self.feature_name}')
+
+            # Cleanup
+            del model, mean, std
+            torch.cuda.empty_cache()
+            gc.collect()