StainX

Torch-first stain normalization for histopathology images with batch processing, training transforms, and optional CUDA kernels.

0.1.0 migration: CuPy backends are removed. Pin stainx<0.1 to stay on the CuPy stack, or switch to Torch tensors with backend="torch" / "torch_cuda".

Features

• Multiple algorithms: Histogram Matching, Reinhard, and Macenko normalization
• Torch backends: torch (CPU / CUDA / MPS) and optional torch_cuda compiled kernels
• Training-ready: StainNormalizerTransform for DataLoader / torchvision pipelines

Installation

Requirements

• Python >= 3.11
• PyTorch >= 2.0.0
• CUDA Toolkit + nvcc (optional; builds stainx_cuda_torch when available)

Supported platforms

Install from PyPI

pip install stainx

Install from source (recommended: Makefile)

git clone https://github.com/rendeirolab/stainx.git
cd stainx
make install          # editable + best-effort CUDA build
# or
make install-dev      # + test/docs tooling

Plain pip also works:

pip install .
# CUDA extension builds automatically when CUDA/nvcc are present; otherwise Torch-only.

Quick Start

import torch
from stainx import Reinhard, Macenko, HistogramMatching, StainNormalizerTransform

reference_image = torch.randn(1, 3, 512, 512)
source_images = torch.randn(10, 3, 512, 512)

normalizer = Reinhard(device="cuda")  # or "cpu" / "mps"
normalizer.fit(reference_image)
normalized = normalizer.transform(source_images)

# Training transform (fit once on a reference — preferred for supervised training)
transform = StainNormalizerTransform(
    method="macenko",
    mode="reference",
    reference=reference_image,
    device="cuda",
    # normalize_to_0_1 defaults to True for Macenko (float [0,1] pipelines)
)
batch_out = transform(source_images)

Modes

API

• fit(reference_images) / transform(images) / fit_transform(images)
• StainNormalizerTransform — nn.Module for pipelines
• Backends: "torch" (default) or "torch_cuda" when the extension is built

Documentation

See the documentation site for installation details, training usage, and examples.

License

GPL-3.0-or-later