Multi-head automated segmentation by incorporating detection head into the contextual layer neural network

TL;DR

This paper introduces a gated multi-head Transformer model that integrates detection and segmentation to improve the anatomical accuracy and robustness of auto-segmentation in radiotherapy, significantly reducing false positives.

Contribution

The novel architecture combines detection and segmentation in a Transformer framework, using gating to suppress false positives in slices lacking target structures, enhancing clinical reliability.

Findings

Significantly lower Dice loss with gating (0.013 vs. 0.732)

Detection probabilities correlate with anatomical presence

Eliminates false positives in invalid slices

Abstract

Deep learning based auto segmentation is increasingly used in radiotherapy, but conventional models often produce anatomically implausible false positives, or hallucinations, in slices lacking target structures. We propose a gated multi-head Transformer architecture based on Swin U-Net, augmented with inter-slice context integration and a parallel detection head, which jointly performs slice-level structure detection via a multi-layer perceptron and pixel-level segmentation through a context-enhanced stream. Detection outputs gate the segmentation predictions to suppress false positives in anatomically invalid slices, and training uses slice-wise Tversky loss to address class imbalance. Experiments on the Prostate-Anatomical-Edge-Cases dataset from The Cancer Imaging Archive demonstrate that the gated model substantially outperforms a non-gated segmentation-only baseline, achieving a mean Dice loss of $0.013 \pm 0.036$ versus $0.732 \pm 0.314$, with detection probabilities strongly correlated with anatomical presence, effectively eliminating spurious segmentations. In contrast, the non-gated model exhibited higher variability and persistent false positives across all slices. These results indicate that detection-based gating enhances robustness and anatomical plausibility in automated segmentation applications, reducing hallucinated predictions without compromising segmentation quality in valid slices, and offers a promising approach for improving the reliability of clinical radiotherapy auto-contouring workflows.