Digital Staining with Knowledge Distillation: A Unified Framework for Unpaired and Paired-But-Misaligned Data

TL;DR

This work tackles the bottleneck of obtaining perfectly aligned stained-unstained image pairs for digital staining by introducing an unsupervised framework based on knowledge distillation. A two-stage teacher model performs light enhancement and colorization to provide references, while a student staining generator learns end-to-end under distillation losses, preserving structure and style without require accurate pairings. The approach is extended to paired-but-misaligned data through a Learning to Align module, enabling use of adjacent-section pixel information. Empirical results on the newly released DSD dataset show superior performance over unsupervised and paired baselines, with additional validation on a synthetic White Blood Cell task indicating clinical relevance and scalable inference for large images.

Abstract

Staining is essential in cell imaging and medical diagnostics but poses significant challenges, including high cost, time consumption, labor intensity, and irreversible tissue alterations. Recent advances in deep learning have enabled digital staining through supervised model training. However, collecting large-scale, perfectly aligned pairs of stained and unstained images remains difficult. In this work, we propose a novel unsupervised deep learning framework for digital cell staining that reduces the need for extensive paired data using knowledge distillation. We explore two training schemes: (1) unpaired and (2) paired-but-misaligned settings. For the unpaired case, we introduce a two-stage pipeline, comprising light enhancement followed by colorization, as a teacher model. Subsequently, we obtain a student staining generator through knowledge distillation with hybrid non-reference losses. To leverage the pixel-wise information between adjacent sections, we further extend to the paired-but-misaligned setting, adding the Learning to Align module to utilize pixel-level information. Experiment results on our dataset demonstrate that our proposed unsupervised deep staining method can generate stained images with more accurate positions and shapes of the cell targets in both settings. Compared with competing methods, our method achieves improved results both qualitatively and quantitatively (e.g., NIQE and PSNR).We applied our digital staining method to the White Blood Cell (WBC) dataset, investigating its potential for medical applications.

Figures (11)

• Figure 1: Examples of our unpaired dataset (DSD_U) and paired-but-misaligned dataset (DSD_P). For two images from the DSD_P dataset, by comparing the black part in the DF image and the white part in the BF image, we can see that cavity areas dilate obviously, which contributes to misalignment. We enhanced the dark-field images from DSD_U and added a color bar for improved visualization, the same for other examples of dark-field images from DSD_U in this paper.
• Figure 2: Illustrations of different digital staining categories. The samples for supervised digital staining are from rivenson2019virtual, and the samples for virtual staining are from liu2021unpaired. We uniformly refer to the transition from unstained to stained as "digital staining" and the transition from one type of staining to another as "virtual staining" throughout this paper.
• Figure 3: The overall structure of our method. In our two-stage pipeline, we begin with the light enhancement stage, where the dark-field image $\mathbf{x}$ is processed by $\mathcal{H}$ to obtain the enhanced image $\mathbf{z}$. Then, in the colorization stage, the teacher model produces $\mathcal{G}_t(\mathbf{z})$ for reference in staining. The student staining generator $\mathcal{G}$ generates the predicted image $\hat{\mathbf{y}}=\mathcal{G}(\mathbf{x} \textcircled{c} \mathbf{z})$. The blue, yellow, and green areas indicate Sections IV-A, IV-B, and IV-C, respectively.
• Figure 4: (a). Two example images from dark-field images and bright-field images of DSD_U. A, B, and C denote pairs of pixels located in the cavity area, pixels next to cells, and pixels inside cells respectively. A, B, and C's illuminance intensity values are shown in the figure as the values range from 0 to 255. (b). Histogram matching from dark-field to enhanced images, with $c=1-c_d(\mathbf{x}_i)=c_b(\mathbf{z}_i)$.
• Figure 5: Illustration of the light enhancement and colorization stages for an example from DSD_P.