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

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 168×, 70×, and 48× for Reinhard, histogram matching, and Macenko respectively, and exceeds the fastest GPU peers by 7-8× (Reinhard) and 2× (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.