K-Stain: Keypoint-Driven Correspondence for H&E-to-IHC Virtual Staining

TL;DR

K-Stain tackles misalignment in H&E-to-IHC virtual staining by embedding spatial correspondences through a three-module design: Hierarchical Spatial Keypoint Detector (HSKD), Keypoint-aware Enhancement Generator (KEG), and Keypoint Guided Discriminator (KGD). The framework aligns tissue structures using affine mappings derived from predicted keypoints and enforces alignment via a keypoint-guided reconstruction loss, a perceptual loss, and an adversarial loss, with all mathematical relationships expressed in $...$ notation. Empirical results on the BCI and HIT datasets show state-of-the-art performance in $SSIM$, $PSNR$, and $LPIPS$, while maintaining high inference efficiency (up to 67 FPS). The work demonstrates that incorporating sparse, robust structural anchors can yield accurate, visually plausible virtual staining, supporting practical deployment in digital pathology and biomarker analysis.

Abstract

Virtual staining offers a promising method for converting Hematoxylin and Eosin (H&E) images into Immunohistochemical (IHC) images, eliminating the need for costly chemical processes. However, existing methods often struggle to utilize spatial information effectively due to misalignment in tissue slices. To overcome this challenge, we leverage keypoints as robust indicators of spatial correspondence, enabling more precise alignment and integration of structural details in synthesized IHC images. We introduce K-Stain, a novel framework that employs keypoint-based spatial and semantic relationships to enhance synthesized IHC image fidelity. K-Stain comprises three main components: (1) a Hierarchical Spatial Keypoint Detector (HSKD) for identifying keypoints in stain images, (2) a Keypoint-aware Enhancement Generator (KEG) that integrates these keypoints during image generation, and (3) a Keypoint Guided Discriminator (KGD) that improves the discriminator's sensitivity to spatial details. Our approach leverages contextual information from adjacent slices, resulting in more accurate and visually consistent IHC images. Extensive experiments show that K-Stain outperforms state-of-the-art methods in quantitative metrics and visual quality.

Figures (5)

• Figure 1: Illustration of different strategies for handling misalignment in virtual staining. (a) Contrastive Learning, (b) External Supervision, (c) Registration Networks, and (d) our proposed K-Stain framework. Here, $X_h$ and $X_i$ denote the H&E and IHC images, respectively; $\hat{X}_i$ is the generated IHC; $\hat{P}_h$ and $\hat{P}_i$ are the predicted keypoints; $\mathcal{L}_{\text{kp}}$, $\mathcal{L}_{\text{perc}}$, and $\mathcal{L}_{\text{adv}}$ represent the keypoint-guided reconstruction, perceptual, and adversarial losses, respectively.
• Figure 2: The proposed K-Stain framework integrates keypoint-based spatial correspondence to address misalignment in virtual staining. It consists of (a) a Hierarchical Spatial Keypoint Detector (HSKD) that predicts consistent keypoints and estimates affine transformation, (b) a Keypoint-aware Enhancement Generator (KEG) that embeds keypoints into dense feature maps for IHC synthesis, and (c) a Keypoint Guided Discriminator (KGD) that enforces consistent adversarial supervision.
• Figure 3: Illustration of the Keypoint Spatial Embedding (KSE) module, which encodes keypoint information into feature maps for virtual staining.
• Figure 4: Comparison of model performance in terms of PSNR versus inference speed (FPS). The proposed K-Stain achieves the best trade-off between performance and efficiency compared to GAN- and diffusion-based baselines.
• Figure 5: Visual comparisons of proposed K-Stain and other methods.