A stable diffusion–based framework for multi-class stroke lesion segmentation in non-contrast brain CT

Accurate segmentation of stroke subtypes in non-contrast computed tomography (CT) is essential for timely diagnosis and treatment planning. However, the inherently low tissue contrast in ischemic regions and the heterogeneous appearance of lesions across subtypes present significant challenges. This study introduces Stable Diffusion Multi-Class Segmentation (SDMC-Seg), a stable diffusion-based framework designed for multi-class stroke subtype segmentation in brain CT images. Unlike previous latent diffusion methods limited to binary tasks, SDMC-Seg addresses multi-class segmentation by encoding each stroke subtype into a distinct semantic channel, allowing the model to disentangle lesion-specific features in latent space for improved class-specific lesion delineation. The architecture integrates a fine-tuned mask autoencoder and a trainable CT image encoder, projecting both modalities into a shared, semantically meaningful latent space. Segmentation masks are reconstructed via a conditional denoising network operating on noisy latent representations conditioned on the CT image. In contrast to traditional diffusion methods requiring multiple iterative steps, SDMC-Seg performs single-step latent inference, enabling faster and more stable predictions while preserving structural detail. The model is evaluated on a curated, anonymized subset of a publicly available stroke CT dataset comprising 6,650 axial slices. Following common practice in lesion segmentation, quantitative training and evaluation are conducted on 2,223 stroke-positive slices (ischemic and hemorrhagic) to focus learning on lesion-specific patterns and to avoid dominance of background-only samples. It achieves an average Dice coefficient of 0.745 and an intersection over union (IoU) of 0.656, demonstrating consistently higher segmentation performance than U-Net, ResUNet, Attention U-Net, TransUNet, Swin-UNETR, and nnU-Net v2, particularly in ischemic stroke segmentation. Ablation results indicate that fine-tuning the mask autoencoder and jointly training the image encoder are both necessary; adapting both components yields the strongest gains. These results suggest that SDMC-Seg provides an efficient and scalable solution for stroke lesion delineation in CT, offering a computationally lightweight alternative to multi-step diffusion-based segmentation frameworks. Remaining limitations in detecting early ischemic changes highlight the need for future extensions incorporating multimodal imaging or temporal progression modeling.