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

0.1.x migration (from 0.0.x): CuPy backends were removed in 0.1.0. Pin stainx<0.1 to stay on the CuPy stack, or switch to Torch tensors with backend="torch" / "torch_cuda". Current release: see stainx.__version__ / PyPI.

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
• Optional CUDA extension: CUDA GPU visible to PyTorch at build time and nvcc

Supported platforms

Platform	Support
Linux + CUDA	Primary (Torch + optional CUDA extension)
Linux CPU	Primary (Torch backend)
Windows	Torch path in CI (CUDA extension not guaranteed)
macOS (MPS / CPU)	Best-effort Torch path (no CUDA extension; not in CI)

Install from PyPI

pip install stainx

PyPI publishes an sdist. Torch backends work out of the box; torch_cuda compiles locally only when the CUDA build gates are met (no prebuilt CUDA wheels).

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 .
# Extension builds when torch.cuda.is_available() and nvcc are present; otherwise Torch-only.
# Prefer make install if you want compile failures to be skipped gracefully.

Quick Start

Use float tensors in [0, 1] (or uint8). Prefer torch.rand — Macenko does not accept negative pixels from torch.randn.

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

reference_image = torch.rand(1, 3, 512, 512)
source_images = torch.rand(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

Mode	Behavior	When to use
reference	Fit once on a fixed reference, then transform	Default for training
batch	Fit on the current batch every forward	Exploratory / domain-shift checks; usually unsafe for reproducible supervised training

API

• fit(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.

Citation

If you use StainX, please cite the preprint:

https://www.biorxiv.org/content/10.64898/2026.08.06.743198v1
DOI: 10.64898/2026.08.06.743198

License

GPL-3.0-or-later