CRAC-DM: class relation-aware categorical diffusion model for surgical scene segmentation

Accurate multi-class segmentation of surgical scenes remains challenging due to ambiguous anatomical boundaries and imaging artifacts. While diffusion-based segmentation methods have achieved good results, they rely on computationally heavy continuous diffusion processes. Recent discrete diffusion variants reduce computation but their performance is limited due to uniform noise, ignoring inter-class relationships that are crucial for generating semantically relevant training signals. To address this gap, we propose the class relation-aware categorical diffusion model (CRAC-DM).

CRAC-DM consists of three key components. In the forward process, we embed semantic class relationships for the first time when adding categorical noise via a class relation-aware transition matrix, biasing noise toward semantically similar categories to generate class-aware supervision signals. In the reverse process, we introduce a step-skipping categorical denoiser (S2D) tailored for discrete diffusion segmentation, enabling fast inference. To further boost inference, we propose a novel confidence-adaptive test time augmentation (TTA) that selectively refines regions of interest with low prediction confidence using entropy-weighted aggregation.

The proposed CRAC-DM was evaluated on the publicly available CholecSeg8k and EndoVis18 datasets. It consistently outperformed state-of-the-art U-Net-, transformer-, and diffusion-based baselines, particularly on tissue segmentation, even for small and under-represented classes while significantly reducing inference time compared to diffusion baselines.

By enhancing the forward process with inter-class similarity and improving the reverse process with a deterministic S2D and targeted TTA, CRAC-DM achieves superior segmentation accuracy, efficiency, and reliability, paving the way for practical deployment in computer-assisted surgery.