OS2CR-Diff: A Self-Refining Diffusion Framework for CD8 Imputation from One-Step Inference to Conditional Representation

Stain imputation in multiplex immunofluorescence (mIF) imaging addresses the challenge of missing or damaged biomarker channels by reconstructing target biomarker images from a limited set of available stains. This approach offers a faster and more efficient alternative to full-panel staining, enabling detailed analysis of the tumour microenvironment. Existing One-Step Inference Models (OSIMs), primarily based on generative adversarial networks (GAN) or autoencoders, often generate suboptimal images with significant artifacts or reduced signal intensity. These limitations impair visual interpretability and reliability of the downstream immunotherapy response assessment. The challenge is further amplified when imputing cytoplasmic biomarkers such as CD8 from commonly used stains such as DAPI, due to the limited spatial correlation and the inherently complex structure of cytoplasmic signals. To address these limitations, we propose a self-refining diffusion model, OS2CR-Diff, which utilises the results from OSIMs as additional conditional representations. Unlike prior studies that rely on a single or limited conditional inputs, OS2CR-Diff incorporates three conditional inputs: the OSIM-imputed target biomarker image, OSIM-imputed complementary biomarker images, and non-antibody-stained images. Furthermore, we propose a feature fusion module that employs a cross-gated attention mechanism to effectively integrate these inputs, enabling context-aware feature refinement and improving the quality and reliability of imputed biomarker images. We evaluated OS2CR-Diff for CD8 biomarker imputation on mIF images of melanoma tissues. Our method out-performed state-of-the-art methods, achieving a 73.4% increase in the Structural Similarity Index Measure (SSIM), a 28.9% gain in the Peak Signal-to-Noise Ratio (PSNR), a 61.2% improvement in Mean Absolute Error (MAE), and significantly lower false positive rates compared to OSIM.