Diffusion Transformer Meets Heat Masks: An Advanced PET Reconstruction Framework

Deep learning has contributed to advances in PET image reconstruction, with reported improvements in image quality through direct training on sinogram or image data. Traditional methods often utilize masks for inpainting tasks, and their incorporation into PET reconstruction frameworks may also benefit PET reconstruction. In this study, we propose a PET reconstruction framework called Diffusion tRansformer mEets heAt Masks (DREAM), which integrates mask mechanisms into both the sino-gram domain and latent space. The framework employs a high-dimensional stacking approach to expand the solution space and designs a mask-driven latent space that is intended to accelerate the diffusion process by leveraging compact priors. Initially, this study introduces a random mask (DREAM-RM) strategy combined with hierarchical learning, which guides the model from focusing on fine-grained local details to capturing broader global patterns. Furthermore, guided by the physical characteristics of PET data, this study extends this framework through the heat mask (DREAM-HM) algorithm. By explicitly decoupling sino-gram data into hot and cold channels based on intensity thresholds and employing a hot-cold mutual mask (HCMM) strategy, DREAM-HM is designed to improve the information distribution between regions of interest and the background. Experimental results indicate that DREAM, particularly the HM variant, improves reconstructed-image quality and preserves clinically relevant details under the evaluated conditions, suggesting its potential utility for PET image reconstruction. The open-source code is available at: https://github.com/yqx7150/DREAM.