Enhanced stain normalization for histopathology images with batch processing support. Optimized for CPU, GPU (CUDA), and MPS (Apple Silicon) devices.
- Multiple algorithms: Histogram Matching, Reinhard, and Macenko normalization
- Automatic backend selection: PyTorch (CPU/GPU/MPS) or optimized CUDA kernels
- Batch processing: Enhanced normalization through efficient batch processing of multiple images
- Flexible device support: CPU, CUDA, MPS (Apple Silicon)
- Python >=3.11
- PyTorch >=2.0.0
- CUDA (optional, for GPU acceleration)
pip install stainxgit clone https://github.com/rendeirolab/stainx.git
cd stainx
pip install .CUDA extensions will be automatically built if CUDA is available.
import torch
from stainx import Reinhard, Macenko, HistogramMatching
# Load your images as torch.Tensor (N, C, H, W) or (N, H, W, C)
reference_image = torch.randn(1, 3, 512, 512) # Reference image
source_images = torch.randn(10, 3, 512, 512) # Images to normalize
# Reinhard normalization
normalizer = Reinhard(device="cuda") # or "cpu"
normalizer.fit(reference_image)
normalized = normalizer.transform(source_images)
# Macenko normalization
normalizer = Macenko(device="cuda")
normalizer.fit(reference_image)
normalized = normalizer.transform(source_images)
# Histogram Matching
normalizer = HistogramMatching(device="cuda", channel_axis=1)
normalizer.fit(reference_image)
normalized = normalizer.transform(source_images)All normalizers follow a scikit-learn-like interface:
fit(reference_images): Compute normalization parameters from reference image(s)transform(images): Apply normalization to imagesfit_transform(images): Fit and transform in one step
Reinhard: Reinhard color normalizationMacenko: Macenko stain separation and normalizationHistogramMatching: Histogram matching normalization
Backends are automatically selected based on device availability:
- CUDA: Used when CUDA is available and device is set to CUDA
- PyTorch: Fallback backend, works on CPU and GPU
You can explicitly specify a backend:
normalizer = Reinhard(device="cuda", backend="pytorch") # Force PyTorch backend- Python >=3.11
- PyTorch >=2.0.0
- CUDA Toolkit (optional, for CUDA backend)
This project is licensed under the GNU General Public License v3 (GPL-3.0-or-later).
See CONTRIBUTING.md for guidelines on contributing to StainX.