StainX: GPU-accelerated batch stain normalization for computational pathology at scale

Abstract

Stain normalization reduces color variability in histopathology whole-slide images, but cohort-scale pipelines lack fused multi-image batch transforms for classical methods. We present StainX, a GPU-accelerated batch stain normalization framework built around a two-stage fit/transform interface. It implements histogram matching, Macenko, and Reinhard normalizers through a portable PyTorch backend and an optional CUDA backend that fuses per-pixel operations for batch throughput. On NVIDIA GPUs, the fused CUDA path outperforms the torch CPU backend by 168x, 70x, and 48x for Reinhard, histogram matching, and Macenko respectively, and exceeds the fastest GPU peers by 7-8x (Reinhard) and 2x (Macenko) at comparable accuracy. StainX also provides user-selectable precision modes, a documented Python API, continuous integration testing, and online documentation. Source code available at https://github.com/rendeirolab/stainx, and documentation at https://stainx.readthedocs.io. Implemented in Python. Runs on Linux, macOS, and Windows.