OAR-UNet: Enhancing Long-Distance Dependencies for Head and Neck OAR Segmentation

Abstract

Accurate segmentation of organs at risk (OARs) is a crucial step in the precise planning of radiotherapy for head and neck tumors. However, manual segmentation methods using CT images, which are still predominantly applied in clinical settings, are inefficient and expensive. Additionally, existing segmentation methods struggle with small organs and have difficulty managing the complex interdependencies between organs. To address these issues, this study proposed an OAR-UNet segmentation method based on a U-shaped architecture with two key designs. To tackle the challenge of segmenting small organs, a Local Feature Perception Module (LFPM) is developed to enhance the sensitivity of the method to subtle structures. Furthermore, a Cross-shaped Transformer Block (CSTB) with a cross-shaped attention mechanism is introduced to improve the ability of the model to capture and process long-distance dependency information. To accelerate the convergence of the Transformer, we designed a Local Encoding Module (LEM) based on depthwise separable convolutions. In our experimental evaluation, we utilized two publicly available datasets, SegRap2023 and PDDCA, achieving Dice coefficients of 78.22% and 89.42%, respectively. These results demonstrate that our method outperforms both previous classic methods and state-of-the-art (SOTA) methods.

1. Introduction

Head and neck cancer is undergoing profound changes in terms of therapeutic strategies. The introduction of Intensity-modulated Radiotherapy (IMRT) in the 1990s [1] marks a major leap forward in therapeutic technology. Because of its ability to accurately respond to highly radiation-sensitive organs in the complex anatomical structure of the head and neck region, IMRT has emerged as a pivotal therapeutic modality [2]. Through optimization of the radiation dose distribution, this technique ensures highly tailored and uniform irradiation of the tumor target area and minimizes radiation exposure to the surrounding normal tissues and critical organs at risk (OARs) [3], which can significantly reduce the side effects of treatment. In radiation therapy (RT), the precise modulation of radiation dose levels to OARs is a central strategy for minimizing complications [4]. Simultaneously, optimization of the dose distribution within the Planned Target Volume (PTV) to achieve the best therapeutic outcome is indispensable. Especially in the framework of Image-guided Radiotherapy (IGRT) [5], the realization of this goal is based on the precise segmentation of OARs in Radiotherapy Computed Tomography (RTCT) [6] or Cone-beam Computed Tomography (CBCT) [7] images. However, the current segmentation of OARs relies heavily on manual work by radiation oncologists. This process is not only extremely time-consuming (e.g., it often takes more than two hours to fully segment nine OARs) but also introduces significant operational variability due to individual differences. Given the significant volume variation of OARs, especially for labeling small structures, it is even more difficult and time-consuming. As treatment protocols become more complex and the number of OARs to be considered increases, the time cost of manual segmentation increases dramatically, which in turn limits patient access to timely radiotherapy. This situation strongly drives the medical community to explore and develop efficient and accurate automated segmentation techniques to meet the increasing treatment challenges and improve the efficiency and effectiveness of treatment.

Owing to the rapid development of deep learning technologies, deep convolutional models have achieved significant success in biomedical image segmentation. The UNet method, proposed by Ronneberger et al. [8], is a classic medical image segmentation method that consists of an encoder and decoder. It combines the features of convolutional neural networks (CNNs) and uses skip connections to make CNNs suitable for medical image segmentation. However, UNet tends to struggle with insufficient feature extraction, resulting in a lower accuracy. Zhou et al. [9] introduced UNet++, an extended and improved version of UNet that incorporates multi-scale skip connections and dense connections to fully extract feature information from medical images. Furthermore, Diakogiannis et al. [10] proposed ResUNet to highlight the benefit of rich gradient flow for segmenting smaller targets and designed a hybrid model based on residual connections and the UNet structure. Limited by the size of convolutional kernels, convolutional neural networks are inadequate in capturing long-range dependencies, which may prevent the model from fully understanding the spatial relationships between different organs, leading to inaccurate segmentation results.

To address this challenge, Transformer architecture [11] has been widely employed in medical image segmentation models. For example, Chen et al. [12] introduced TransUNet, wherein a Vision Transformer block [13] is incorporated into the bottleneck of the UNet architecture to improve the model’s ability to capture long-range dependencies. Cao et al. [14] proposed SwinUNet, which replaces the convolutional blocks in UNet with Swin Transformer modules [15] to enhance the model’s ability to learn global information. Li et al. [16] proposed UTAC-Net, a model that utilizes a dual-stream encoder integrating convolutional layers with Transformer modules. This encoder simultaneously captures global and local information from the target images for segmentation, thereby improving the model’s precision and contextual awareness in segmentation tasks. Although these Transformer-based models have an advantage in capturing global dependencies, they may not be precise enough when dealing with the fine details of smaller tissue organs. This is because the Transformer’s self-attention mechanism inherently has a global receptive field, which can lead to the “dilution” of local detail information when processing fine structures. Furthermore, Transformer-based models are computationally complex and have a slower convergence rate. Therefore, for the task of multi-organ segmentation in the head area, Transformer-based models need further optimization.

In addition to classic medical image segmentation methods, several methods that specifically target head and neck OAR segmentation have drawn considerable attention. Gao et al. [17] proposed FocusNetv2 to simulate how radiologists delineate OARs. Their solution consists of two main components: a segmentation network and an adversarial autoencoder (AAE) for organ-shape constraints. AAE employs prior shapes of OARs to constrain the segmentation results, effectively addressing the challenge of detecting small targets. Furthermore, Zhong et al. [18] designed a “human-in-the-loop” strategy to improve the OARs segmentation accuracy. They trained a two-dimensional (2D) U-Net-like network, which was subjectively evaluated by 2–3 oncologists to further retrain the model. Moreover, Zhang et al. [19] proposed SegReg, which used standard templates to register training images and labels before segmentation, thereby reducing the training difficulty. Although SegReg and FocusNetv2 have achieved acceptable segmentation accuracy, these methods face challenges in practical clinical application. For example, the “human-in-the-loop” strategy requires radiologists to continuously refine the model’s segmentation predictions. Although this iterative refinement process gradually enhances the model’s precision, it does not simplify the radiologists’ workflow but rather introduces additional complexity into the procedures. On the other hand, FocusNetv2 relies heavily on shape constraints as prior information which is frequently unavailable. This limitation hinders the application of FocusNetv2 in clinical application.

To address the issues mentioned above, we have developed the OAR-UNet method for the segmentation of head and neck organs (OARs). This method does not rely on extra prior information; instead, it converts image features into coordinate information and dynamically generates weight maps to highlight key positions, thereby improving the recognition ability of small head and neck organs. Additionally, we redesigned the Transformer block to allow the model to capture more comprehensive long-range dependencies while enhancing convergence speed. The main contributions of this paper are as follows:

To address the challenge of recognizing small organs in medical images of the head and neck, we developed the Local Feature Perception Module (LFPM). This module converts image features into coordinate information and employs a coordinate weight calculation network to dynamically generate weight maps during feature extraction. By leveraging these maps, the LFPM selectively enhances key gradient information at specific coordinate positions in the feature map through long skip-distance connections, thereby improving the detection of small organs.

To enhance the model’s capability in handling long-distance dependencies, we have developed a novel Cross-shaped Transformer Block (CSTB). This block incorporates a cross-shaped self-attention mechanism within the Transformer architecture, effectively reducing the computational complexity associated with traditional self-attention mechanisms.

To accelerate convergence, we introduced a lightweight deep separable convolution module named Local Encoding Module (LEM). It serves as a pre-processing module in front of the CSTB, utilizing feature information obtained through local induction to enhance the training efficiency of the Transformer.

2. Methodology

2.1. Overview of OAR-UNet

As a model specifically designed for head and neck OAR segmentation, the OAR-UNet architecture effectively addresses the limitations of traditional U-shaped networks by incorporating the LFPM and CSTB. LFPM processes feature maps in the early stages of the U-shaped network, computes the importance of local features for coordinate information in both directions, captures detailed information in high-resolution feature maps, and efficiently transfers gradient information to the decoder. In the deeper layers of the network, where abundant semantic information is present, CSTB employs Cross-Shaped Self-Attention for self-attention computations. This method is more efficient than traditional patch-based methods and significantly enhances a network’s ability to perceive long-distance dependencies. Consequently, OAR-UNet can thoroughly interpret OAR images and generate highly accurate segmentation maps, thereby improving the accuracy of the OAR image segmentation tasks. The architecture of the OAR-UNet is illustrated in Figure 1.

2.2. Designing OAR-UNet Model for OAR Segmentation

2.2.1. Local Feature Perception Module

From the perspective of image information, smaller organs within at-risk regions have less feature information and a weaker gradient flow, requiring more intensive perception [20]. Providing sufficient local information to the decoder during training can improve the performance of the U-shaped network and enable a better understanding of the features of smaller organs. Given that critical organ regions are often small, and their features are less apparent, the U-shaped network may lose crucial feature information of these organs during the downsampling process. In the U-shaped network, the early layers produce larger feature maps with smaller receptive fields, thereby allowing for more localized information extraction [21]. To enhance the network’s focus on OAR regions and improve local information extraction, we designed LFPM for OAR-UNet. This module is integrated into the long skip connections in the first three stages of the network, as shown in Figure 1, with a detailed diagram in Figure 2. The LFPM design incorporates coordinate decomposition, which splits input features into components along both the width (W) and height (H) directions. It uses a four-branch fully connected layer to apply attention weighting to these coordinate positions. The resulting attention map, obtained by multiplying the attention layer components, is then multiplied by the input feature map to pass the important local feature gradient information to the decoder. This process strengthens the network’s ability to perceive local information and enhances attention to smaller organs.

The LFPM is integrated into the long skip connections of the first three stages of OAR-UNet. After feature extraction by Encoder1 to Encoder3, the three feature maps fed into the LFPM module had resolutions of the original resolution, half of the original resolution, and a quarter of the original resolution, respectively. It is assumed that the input features of the LFPM are $\mathrm{X}\in {\mathbb{R}}^{\left(\mathrm{H}\times \mathrm{W}\right)\times \mathrm{C}}$, then the input feature map undergoes the following steps. To encourage long-distance interactions for capturing accurate spatial location information, the local perceptual attention is first divided into two branches, with one pooled in the H direction and the other pooled in the W direction. To obtain feature vectors of sizes of (C,1,W) and (C,H,1), the decomposed vectors represent the horizontal and vertical coordinate components on the original feature map, with each coordinate indicating the information of a specific pixel. Considering input X, the output of the C-th channel in the H direction can be expressed as follows:

$${\mathrm{Z}}_{\mathrm{c}}^{\mathrm{h}}\left(\mathrm{h}\right)={\displaystyle \frac{1}{\mathrm{W}}}\sum _{0\le \mathrm{m}<\mathrm{W}}{\mathrm{X}}_{\mathrm{c}}(\mathrm{h},\mathrm{m}).$$

(1)

Similarly, the output of the C-th channel in the W direction can be expressed as follows:

$${\mathrm{Z}}_{\mathrm{c}}^{\mathrm{w}}\left(\mathrm{w}\right)={\displaystyle \frac{1}{\mathrm{H}}}\sum _{0\le \mathrm{n}<\mathrm{H}}{\mathrm{X}}_{\mathrm{c}}\left(\mathrm{n},\mathrm{w}\right).$$

(2)

The two feature maps were reshaped to obtain feature vectors with sizes of (C, H) and (C, W). When H = W, these two feature maps are spliced to obtain feature vectors with size (C, W + H). Subsequently, four fully connected branches were used to reduce their dimensions to one-fourth of the input channel, and four feature maps with the size of (C/2, W) were obtained. After four fully connected branches were formed, the resulting four feature maps were spliced to obtain a feature map with the size of (C, W + H). The idea of channel attention weighting is followed in the process of using four fully connected branches to reduce dimensionality and splicing. Drawing inspiration from SENetV2 [22], the use of multiple branches enhances the perception of feature information across various dimensions. The feature map was then split into two parts along the channel dimension and reshaped to produce two feature maps with sizes of (C, 1, W) and (C, 1, H), respectively. These dimensions are equivalent to those obtained from the mean pooling operation. Finally, the feature maps with sizes of (C, 1, W) and (C, 1, H) were multiplied by matrix multiplication to obtain an attention map ${\mathrm{M}}_{\mathrm{a}\mathrm{t}\mathrm{t}}$ with dimensions (C, H, W). Once the attention map is acquired, the output $\mathrm{Y}$ of the LFPM can be expressed as follows:

$$\mathrm{Y}={\mathrm{M}}_{\mathrm{a}\mathrm{t}\mathrm{t}}\times \mathrm{X}.$$

(3)

By leveraging the attention map ${\mathrm{M}}_{\mathrm{a}\mathrm{t}\mathrm{t}}$, the method enhances the rich local information of small organs within the input feature map, with important channels being weighted during the training process to highlight their significance.

2.2.2. Cross-Shaped Transformer Block

In OARs images, there is strong interdependence among organs, necessitating enhanced long-distance dependency perception by the network. Traditional U-shaped networks struggle to capture long-distance dependencies using only convolutional structures. To address this challenge, we introduced a module that integrates the Local Encoding Module (LEM) with the Cross-Shaped Transformer Block (CSTB) in the final two stages of the OAR-UNet architecture to enhance long-distance dependency perception. Inspired by Vision Transformer (ViT), our design adapts the concept of tokenization for image representation. In ViT, the image is segmented into small patches, which are then vectorized into tokens and processed by the Transformer encoder to lower computational complexity. However, the patch size, a hyperparameter, can limit the model’s long-range dependency perception. To address this, we propose a cross-shaped Transformer module that employs cross-shaped self-attention. Unlike ViT’s patch-based tokens, our method uses horizontal and vertical stripes as tokens. This allows us to compute interdependencies more effectively, enhancing the model’s ability to perceive long-range dependencies. The use of stripes provides greater modeling capabilities, strengthening the associations between distant pixels. Consequently, our approach more effectively captures long-range perceptual dependencies compared to the patch-based token method.

To address the inherent convergence issues of Transformers, we introduce LEM, a local enhancement module, before the cross-shaped Transformer. The LEM module is designed using depthwise separable convolution, which maintains a lightweight structure while offering a larger receptive field compared to traditional convolution, allowing it to capture more prior features. The combination of the LEM module with the cross-shaped Transformer can not only accelerate convergence but also strengthen the model’s ability to capture long-distance dependencies. The LEM module integrated with the Cross-Shaped Transformer Block is shown in Figure 3.

As shown in Figure 3, given the input feature map $\mathrm{X}\in {\mathbb{R}}^{\left(\mathrm{H}\times \mathrm{W}\right)\times \mathrm{C}}$, the LEM operation can be defined as follows:

$$\mathrm{L}\mathrm{E}\mathrm{M}\left(\mathrm{X}\right)={\mathrm{D}\mathrm{C}\mathrm{o}\mathrm{n}\mathrm{v}}_{3\times 3}\left(\mathrm{B}\mathrm{N}\left({\mathrm{C}\mathrm{o}\mathrm{n}\mathrm{v}}_{1\times 1}\left(\mathrm{X}\right)\right)\right)+\mathrm{X}.$$

(4)

In Equation (4), the LEM module consists of a 1 × 1 convolution, Batch Norm, and a 3 × 3 depth separable convolution [23]. The features processed by the LEM module are fed into the Cross-Shaped Transformer, which consists of the Layer Norm, cross-shaped self-attention, and MLP modules. Each Layer Norm module is followed by using a residual connection. With two layers, MLP uses the GELU activation function. The computation for the cross-shaped Transformer is expressed as follows:

$${\widehat{\mathrm{X}}}_{\mathrm{l}}=\mathrm{C}\mathrm{S}\mathrm{S}\mathrm{A}\left(\mathrm{L}\mathrm{N}\left({\mathrm{X}}_{\mathrm{l}-1}\right)\right)+{\mathrm{X}}_{\mathrm{l}-1,}$$

(5)

$${\mathrm{X}}_{\mathrm{l}}=\mathrm{M}\mathrm{L}\mathrm{P}\left(\mathrm{L}\mathrm{N}\left({\widehat{\mathrm{X}}}_{\mathrm{l}}\right)\right)+{\widehat{\mathrm{X}}}_{\mathrm{l}},$$

(6)

where ${\mathrm{X}}_{\mathrm{l}-1}$ represents the output of the previous layer, ${\mathrm{X}}_{\mathrm{l}}$ represents the output of the current layer, and the feature map inputs the long-range-dependent perceptual capability module. In the cross-shaped Transformer, cross-shaped self-attention (CSSA) is achieved by creating a cross-shaped window and performing self-attention computations in parallel within the horizontal and vertical strips that form this window. The structure of the cross-shaped self-attention is illustrated in Figure 4.

According to the self-attention mechanism, the input feature $\mathrm{X}\in {\mathbb{R}}^{\left(\mathrm{H}\times \mathrm{W}\right)\times \mathrm{C}}$ is first linearly projected onto the K heads. Furthermore, each head performs a localized self-attention computation within a horizontal or vertical strip. For the horizontal strip self-attention computation, X is uniformly partitioned into equal-width nonoverlapping horizontal strips $[{\mathrm{X}}^1,{\mathrm{X}}^2,\dots ,{\mathrm{X}}^{\mathrm{M}}]$, each containing sw × W tokens, where sw represents the strip width, which can be adjusted to balance the learning capacity and the computational complexity. Formally, it is assumed that the dimensions of the projected query, key, and value of the k-th head are ${\mathrm{d}}_{\mathrm{k}}$. The output of the horizontal strip self-attention computation of the k-th head (denoted as ${\mathrm{H}\mathrm{A}\mathrm{t}\mathrm{t}}_{\mathrm{k}}\left(\mathrm{X}\right)$) can be expressed as follows:

$$\mathrm{X}=[{\mathrm{X}}^1,{\mathrm{X}}^2,\dots ,{\mathrm{X}}^{\mathrm{M}}],$$

(7)

$${\mathrm{Y}}_{\mathrm{k}}^{\mathrm{i}}=\mathrm{A}\mathrm{t}\mathrm{t}({\mathrm{X}}^1{\mathrm{W}}_{\mathrm{k}}^{\mathrm{Q}},{\mathrm{X}}^2{\mathrm{W}}_{\mathrm{k}}^{\mathrm{K}},{\mathrm{X}}^{\mathrm{M}}{\mathrm{W}}_{\mathrm{k}}^{\mathrm{V}}),$$

(8)

$${\mathrm{H}\mathrm{A}\mathrm{t}\mathrm{t}}_{\mathrm{k}}\left(\mathrm{X}\right)=[{\mathrm{Y}}_{\mathrm{k}}^1,{\mathrm{Y}}_{\mathrm{k}}^2,\dots ,{\mathrm{Y}}_{\mathrm{k}}^{\mathrm{M}}],$$

(9)

where ${\mathrm{X}}^{\mathrm{i}}\in {\mathbb{R}}^{(\mathrm{s}\mathrm{w}\times \mathrm{W})\times \mathrm{C}}$, $\mathrm{M}\in \mathrm{H}/\mathrm{s}\mathrm{w}$,$\mathrm{i}=1,\dots ,$ M. ${\mathrm{W}}_{\mathrm{k}}^{\mathrm{Q}}\in {\mathbb{R}}^{\mathrm{C}\times {\mathrm{d}}_{\mathrm{k}}}$, ${\mathrm{W}}_{\mathrm{k}}^{\mathrm{K}}\in {\mathbb{R}}^{\mathrm{C}\times {\mathrm{d}}_{\mathrm{k}}}$, ${\mathrm{W}}_{\mathrm{k}}^{\mathrm{V}}\in {\mathbb{R}}^{\mathrm{C}\times {\mathrm{d}}_{\mathrm{k}}}$ denote the projection matrices of the queries, keys, and values for the k-th head, respectively, and ${\mathrm{d}}_{\mathrm{k}}$ is C/K. The self-attention of the vertical strip can be derived similarly, and the output of its k-th head is denoted as $\mathrm{V}{\mathrm{A}\mathrm{t}\mathrm{t}}_{\mathrm{k}}\left(\mathrm{X}\right)$. It is assumed that there is no directional bias in the head and neck medical images. This study divides K heads into two groups (each with K/2 heads, where K is usually an even number). The first group of heads performs horizontal strip self-attention, whereas the second group performs vertical strip self-attention. Finally, the outputs of the two groups were concatenated. The process can be expressed as follows:

$$\mathrm{C}\mathrm{S}\mathrm{S}\mathrm{A}\left(\mathrm{X}\right)=\mathrm{C}\mathrm{o}\mathrm{n}\mathrm{c}\mathrm{a}\mathrm{t}\left({\mathrm{h}\mathrm{e}\mathrm{a}\mathrm{d}}_1,{\mathrm{h}\mathrm{e}\mathrm{a}\mathrm{d}}_2,\dots ,{\mathrm{h}\mathrm{e}\mathrm{a}\mathrm{d}}_{\mathrm{k}}\right){\mathrm{W}}^{\mathrm{O}},$$

(10)

$$where,{\mathrm{h}\mathrm{e}\mathrm{a}\mathrm{d}}_{\mathrm{k}}=\left\{\begin{array}{c}{\mathrm{H}\mathrm{A}\mathrm{t}\mathrm{t}}_{\mathrm{k}}\left(X\right), k=1,\dots ,{\displaystyle \frac{\mathrm{K}}{2}}\\ {\mathrm{V}\mathrm{A}\mathrm{t}\mathrm{t}}_{\mathrm{k}}\left(X\right), k={\displaystyle \frac{\mathrm{K}}{2}},\dots ,K\end{array}\right.,$$

(11)

where CSSA represents cross-shaped self-attention and ${\mathrm{W}}^{\mathrm{O}}$ represents the commonly used projection matrix for projecting the self-attention results to the target output dimension (set to C by default). As mentioned above, a key insight in the design of our self-attention mechanism is to categorize the multiheads into different groups and apply different self-attention operations accordingly. In other words, the attention region of each token in each Transformer block is extended by grouping the multiheads. In contrast, the existing self-attention mechanism applies the same self-attention operation to different multiheads. The experimental section reveals that this design achieves better performance.

3. Data Preparation and Implementation Details

3.1. Data Preparation

In this section, we aim to validate the accuracy and robustness of OAR-UNet by utilizing two publicly available datasets for high-risk organ segmentation in the head and neck region, SegRap2023 and PDDCA. These datasets are used to assess the performance of the proposed method.

SegRap2023 datasets [24]: The SegRap2023 dataset provides CT images of the head and neck organs at risk, as well as accurate segmentation annotations. SegRap2023 contains 200 datasets, including 120 cases of training data, 20 cases of validation data, and 60 cases of test data. This study used nine anatomical categories in the dataset (Brain, BrainStem, TemporalLobe, OralCavity, Mandible, Parotid, Mastoid, SpinalCord, and Larynx) and pre-processed the raw data by removing noise from binarized images, extracting the maximally connected region, and generating a non-zero mask to identify the structural regions of interest. Subsequently, the images were computed and cropped based on non-zero masks to focus on the regions of interest and reduce the data size. To ensure that the images and labels have a uniform spatial resolution, this study performed a resampling process and normalized the images based on their median, mean, and standard deviation before slicing, resulting in 15,312 training images, 2520 validation images, and 7656 test images.

PDDCA datasets [25]: PDDCA, as a CT dataset for head and neck endangerment and organ segmentation tasks, contains CT images from 48 patients, and we used five anatomical categories in the dataset (BrainStem, Chiasm, Mandible, OpticNerve, Parotid), together with the manual identification of the bone markers. The dataset was divided into 38 training and 10 test samples. During pre-processing, we first sliced them in two dimensions (2D). After slicing, the images were binarized to obtain a common bounding box based on the labels, and the region of interest was obtained based on the common bounding box. In addition, morphological manipulation is used to optimize the shape of the bounding box to ensure that the final extracted region contains the structure of interest in its entirety. Finally, the original and labeled images are cropped according to the expanded bounding box coordinates, and 1250 training images and 313 test images are obtained.

In this study, the input images were scaled to a size of 512 × 512 pixels before inputting them into OAR-UNet for training. To broaden the data distribution for training purposes, we employed data augmentation techniques on the dataset throughout the training process, which involved random cropping and random warping of the images. Augmentation during the training process only expands the data distribution but does not change the amount of input data, which improves the generalization ability of the model without increasing the amount of data.

3.2. Loss Function

During training, this study used a hybrid loss function that combined Cross-Entropy (CE) loss with Dice loss. The effectiveness of this combined loss function was further validated through ablation experiments described in the experimental section.

Cross-Entropy Loss: The Cross-Entropy loss function was used to compare the predicted segmentation mask with the true mask in semantic segmentation tasks. The Cross-Entropy loss is defined as follows:

$${\mathrm{C}\mathrm{E}}_{\mathrm{l}\mathrm{o}\mathrm{s}\mathrm{s}}=-{\displaystyle \frac{1}{\mathrm{N}}}\sum _{\mathrm{i}=1}^{\mathrm{N}}\sum _{\mathrm{c}=1}^{\mathrm{C}}{\mathrm{y}}_{\mathrm{i},\mathrm{c}}\mathrm{l}\mathrm{o}\mathrm{g}\left({\widehat{\mathrm{y}}}_{\mathrm{i},\mathrm{c}}\right)$$

(12)

where $\mathrm{C}$ is the number of classes, and ${\mathrm{y}}_{\mathrm{i},\mathrm{c}}$ indicates whether the $\mathrm{i}$-th sample belongs to class $\mathrm{c}$, ${\widehat{\mathrm{y}}}_{\mathrm{i},\mathrm{c}}$ represents the predicted probability that the $\mathrm{i}$-th sample belongs to class $\mathrm{c}$. The loss function calculates the sum of the logarithms of the predicted probabilities weighted by the true labels, providing a measure of how well the model predictions align with the actual classes. In the segmentation tasks, the Cross-Entropy loss function was adapted to measure the difference between the predicted and actual pixel-wise classifications. The goal of segmentation is to classify each pixel of an image into one of several classes; thus, the Cross-Entropy loss is computed for each pixel across the entire image.

Dice Loss: The Dice loss function aims to quantify the overlap or similarity between the predicted segmentation and ground truth label. Dice loss is defined as follows:

$${\mathrm{D}\mathrm{i}\mathrm{c}\mathrm{e}}_{\mathrm{L}\mathrm{o}\mathrm{s}\mathrm{s}}=1-{\displaystyle \frac{2\left|{\mathrm{y}}_{\mathrm{i}}\cap {\widehat{\mathrm{y}}}_{\mathrm{i}}\right|}{\left|{\mathrm{y}}_{\mathrm{i}}\right|+\left|{\widehat{\mathrm{y}}}_{\mathrm{i}}\right|}}$$

(13)

where $\left|{\mathrm{y}}_{\mathrm{i}}\cap {\widehat{\mathrm{y}}}_{\mathrm{i}}\right|$ denotes the number of pixels that are correctly predicted as positive samples (part of the object of interest) and are positive samples in the ground truth label. $\left|{\mathrm{y}}_{\mathrm{i}}\right|$ represents the number of positive sample pixels in the true label and $\left|{\widehat{\mathrm{y}}}_{\mathrm{i}}\right|$ denotes the total number of pixels predicted as positive samples.

Hybrid Loss: The hybrid loss is formed by combining two loss functions, Cross-Entropy loss and Dice loss, which can be expressed as follows:

$${\mathrm{H}\mathrm{i}\mathrm{b}\mathrm{i}\mathrm{r}\mathrm{d}}_{\mathrm{L}\mathrm{o}\mathrm{s}\mathrm{s}}={\mathsf{\alpha }\mathrm{C}\mathrm{E}}_{\mathrm{L}\mathrm{o}\mathrm{s}\mathrm{s}}+{\mathsf{\beta }\mathrm{D}\mathrm{i}\mathrm{c}\mathrm{e}}_{\mathrm{L}\mathrm{o}\mathrm{s}\mathrm{s}},$$

(14)

where $\mathsf{\alpha }$ and $\mathsf{\beta }$ represent the weights of the Cross-Entropy loss and Dice loss, respectively, and the specific values are determined in the experiments.

3.3. Evaluation Criteria

Dice similarity coefficient (DSC) and intersection over ratio (IoU) metrics were used to evaluate the segmentation results by measuring the degree of similarity between the predicted segmentation masks and the true labeling masks. They are formulated as follows:

$$\mathrm{D}\mathrm{S}\mathrm{C}={\displaystyle \frac{2\left|{\mathrm{y}}_{\mathrm{i}}\cap {\widehat{\mathrm{y}}}_{\mathrm{i}}\right|}{\left|{\mathrm{y}}_{\mathrm{i}}\right|+\left|{\widehat{\mathrm{y}}}_{\mathrm{i}}\right|}}={\displaystyle \frac{2\mathrm{T}\mathrm{P}}{2\mathrm{T}\mathrm{P}+\mathrm{F}\mathrm{P}+\mathrm{F}\mathrm{N}}},$$

(15)

$$\mathrm{I}\mathrm{o}\mathrm{U}={\displaystyle \frac{\left|{\mathrm{y}}_{\mathrm{i}}\cap {\widehat{\mathrm{y}}}_{\mathrm{i}}\right|}{\left|{\mathrm{y}}_{\mathrm{i}}\cup {\widehat{\mathrm{y}}}_{\mathrm{i}}\right|}}={\displaystyle \frac{\mathrm{T}\mathrm{P}}{\mathrm{T}\mathrm{P}+\mathrm{F}\mathrm{P}+\mathrm{F}\mathrm{N}}},$$

(16)

where ${\mathrm{y}}_{\mathrm{i}}$ denotes the predicted segmentation mask and ${\widehat{\mathrm{y}}}_{\mathrm{i}}$ denotes the true labeling mask. True Positive (TP) refers to correctly identified positive cases, False Positive (FP) refers to incorrectly identified positive cases, and False Negative (FN) refers to missed positive cases.

3.4. Implementation Details

OAR-UNet was implemented on a server with Intel(R) Xeon(R) Silver 4216 CPU and two Quadro GV100 GPUs, using PyTorch version 1.10 and MMSegmentation version 0.29. All methods were implemented in MMSegmentation with the same training configurations and parameters for each method. The batch size was set to 8, and a Poly learning rate adjustment strategy was applied, with an initial learning rate of 10⁻⁴ and a final learning rate of 10⁻⁶. The Adam optimizer was used, and training was conducted for 20,000 iterations.

4. Results

4.1. Ablation Experiments

In this section, an ablation study of OAR-UNet is conducted, which includes ablation experiments for the LFPM, CSTB, and LEM modules, as well as the hybrid loss function. The weight loss of the hybrid was experimentally determined.

Table 1. Ablation experiments with LFPM, CSTB, and LEM. The best results are shown in bold.

			SegRap2023		PDDCA
LFPM	CSTB	LEM	mIoU	mDice	mIoU	mDice
×	×	×	0.6751	0.7014	0.7894	0.8145
√	×	×	0.7043	0.7355	0.8164	0.8451
×	√	×	0.6932	0.7251	0.8041	0.8355
×	×	√	0.6843	0.7151	0.7918	0.8189
×	√	√	0.7241	0.7598	0.8411	0.8741
√	√	√	0.7478	0.7822	0.8545	0.8942

presents the results of ablation studies for the LFPM, CSTB, and LEM modules. As shown, using the LFPM improves the mIoU and mDice metrics in three ways. The metrics improved by only two points with the CSTB module alone and less than 1 point with the LEM module alone. However, when both the LEM and CSTB modules are used together, the metrics improve by nearly five points compared to the baseline and by four points relative to the use of only the CSTB module. This suggests that the LEM module effectively reduces the training complexity of the CSTB module. The highest metric values are achieved when all three modules are used simultaneously.

In

Table 2. Ablation experiments with Loss Function. Where CE represents cross-entropy loss, Dice represents Dice loss, and ratio represents the ratio of $\mathsf{\alpha }$ to $\mathsf{\beta }$ in Equation (14). The best results are shown in bold.

			SegRap2023		PDDCA
CE	Dice	Ratio	mIoU	mDice	mIoU	mDice
√	×	—	0.7255	0.7613	0.8378	0.8798
×	√	—	0.7135	0.7589	0.8276	0.8681
√	√	1:1	0.7372	0.7764	0.8436	0.8842
√	√	2:1	0.7478	0.7822	0.8545	0.8942
√	√	3:1	0.7391	0.7781	0.8492	0.8897
√	√	1:2	0.7295	0.7710	0.8398	0.8791
√	√	1:3	0.7218	0.7643	0.8341	0.8746

, we conducted ablation experiments on the loss function, comparing the accuracy of using both the CE loss function and the Dice loss function with the accuracy of using only one loss function. The results show that the accuracy is higher when both the CE loss function and the Dice loss function are used together compared to when only one of them is used. We designed experiments to determine the optimal loss-function ratio. Then, the highest metrics were obtained when the loss weight of $\mathsf{\alpha }$ to $\mathsf{\beta }$ reached 2:1. Thus, the weight loss ratio of $\mathsf{\alpha }$ to $\mathsf{\beta }$ was determined to be 2:1.

In this section, we aim to validate the accuracy and robustness of OAR-UNet using two publicly available datasets for high-risk organ segmentation.

4.2. Comparative Experiments

Table 3. Comparative results of OAR-UNet and other seven methods: classic UNet, UNet++, ResUNet, TransUNet, SwinUNet, and state-of-the-art methods SegReg and FocusNetv2. The best results are shown in bold.

	SegRap2023		PDDCA
Method	mIoU	mDice	mIoU	mDice
UNet [8]	0.6542	0.6851	0.7715	0.8074
UNet++ [9]	0.6689	0.7007	0.7811	0.8123
ResUNet [10]	0.6734	0.7015	0.7826	0.8193
TransUNet [12]	0.6918	0.7202	0.7945	0.8310
SwinUNet [14]	0.7101	0.7387	0.8068	0.8432
SegReg [18]	0.7213	0.7674	0.8352	0.8752
FocusNetv2 [17]	0.7362	0.7695	0.8223	0.8785
OAR-UNet	0.7478	0.7822	0.8545	0.8942

presents a comparison of segmentation results among different methods. The results show that the proposed OAR-UNet outperforms the seven comparative methods in terms of mIoU and mDice metrics across two different datasets. Specifically, this study conducted a comparison with classic convolution-based methods such as UNet [8], UNet++ [9], and ResUNet [10], alongside Transformer-based methods like TransUNet [12] and SwinUNet [14]. Additionally, the state-of-the-art methods SegReg [18] and FocusNetv2 [17] were also evaluated for their performance in the task of segmenting head and neck OARs. The experimental results demonstrate that our method achieved the highest mDice scores of 0.7822 and 0.8942 on the SegRap2023 and PDDCA, respectively. OAR-UNet achieves higher accuracy than the latest SegReg and FocusNetv2 methods and does not require additional information or a “human-in-the-loop” strategy, making it suitable for clinical applications.

Figure 5 and Figure 6 illustrate the segmentation accuracy of the different methods across various anatomical structures using box plots. These box plots effectively capture the accuracy characteristics of OAR-UNet for each anatomical category and facilitate comparisons of the distribution characteristics across multiple experimental groups. The box plots show that across all anatomical categories and datasets, the distribution of IoU and Dice metrics for OAR-UNet is superior to that of other methods. Specifically, the OAR-UNet model consistently achieves higher overall performance metrics compared to all other evaluated methods, particularly exhibiting significant advantages when segmenting smaller organs, such as the parotid gland. These results substantiate the superior efficacy of the OAR-UNet in these specific contexts.

4.3. Visualization

This study visually compares OAR-UNet with seven other methods to illustrate the subjective advantages of the OAR-UNet method. Visual comparisons are displayed in Figure 7 and Figure 8, showing the performance differences between OAR-UNet and the other methods. We used two publicly available datasets, SegRap2023 and PDDCA, for visualization. By selecting several example cases and extracting relevant slices, we can generate visual representations of segmentation outcomes. The visual results indicate that our proposed OAR-UNet model yields segmentation results that are visually more aligned with the ground truth than the benchmark methods. The comparative methods exhibited notable instances of missed or erroneous identifications, particularly in the segmentation of smaller anatomical structures. For example, both the UNet and UNet++ models struggle to identify smaller organs. In the context of segmenting small targets, the OAR-UNet model demonstrates distinct superiority, with no instances of missed or incorrect identifications observed.

5. Discussion

5.1. Findings

The OAR-UNet model demonstrates significant improvements in segmentation accuracy for head and neck organs at risk (OARs), driven by the integration of the Local Feature Perception Module (LFPM) and the Cross-Shaped Transformer Block (CSTB). The introduction of LFPM enhances the network’s capability to perceive detailed local features, particularly for smaller organs within at-risk regions. By incorporating local perceptual attention mechanisms early in the network, LFPM effectively captures and preserves crucial spatial information that is otherwise prone to loss during the downsampling process in traditional U-shaped networks. The CSTB, on the other hand, addresses the challenge of capturing long-distance dependencies that are essential for accurate OAR segmentation. By employing cross-shaped self-attention mechanisms and LEM, CSTB improves the model’s ability to perceive relationships between distant pixels, surpassing the traditional patch-based methods in efficiency and accuracy. The combination of LFPM and CSTB not only boosts the model’s performance in interpreting complex OAR images but also significantly enhances segmentation maps, as evidenced by improved metrics across the datasets.

Our ablation studies reveal that integrating both LFPM and CSTB modules results in the highest performance metrics, underscoring the importance of leveraging both local and long-range feature extraction strategies. Specifically, the use of LFPM alone improves the mean Intersection over Union (mIoU) and mean Dice coefficient (mDice) by approximately 2 points, while the CSTB module alone contributes less than 1 point to these metrics. However, when both modules are combined, the improvements reach nearly 5 points compared to the baseline and 4 points relative to CSTB alone. This highlights the synergistic effect of combining local feature perception with long-range dependency modeling. In comparative experiments, OAR-UNet outperforms several state-of-the-art methods, including UNet, UNet++, ResUNet, and advanced Transformer-based methods like TransUNet and SwinUNet. The proposed model achieves the highest mDice scores of 0.7822 and 0.8942 on the SegRap2023 and PDDCA datasets, respectively.

These results validate the effectiveness of the OAR-UNet architecture and demonstrate its strong performance in accurately segmenting small and challenging anatomical structures, such as the parotid gland. Overall, the design and feature extraction capabilities of the OAR-UNet model set a new benchmark for head and neck OAR segmentation and provide a valuable tool for clinical diagnosis and treatment planning.

5.2. Limitation and Future Works

We recognize that transitioning research findings from the laboratory setting to actual patient care in a clinical setting requires a thorough understanding and consideration of the real-world clinical environment. This means that although the segmentation results achieved in the laboratory setting may be impressive, further validation and adjustment may be necessary in actual clinical practice to ensure that the algorithm can adapt to the complex realities of clinical settings.

For example, the model’s increased computational complexity due to the integration of the Local Feature Perception Module (LFPM) and Cross-Shaped Transformer Block (CSTB) may limit its applicability in real-time scenarios with constrained resources. Future work should focus on optimizing the model to reduce computational requirements, making it more practical for real-time use. Furthermore, broadening the assessment of the model to encompass a variety of datasets from multiple institutions and imaging modalities will be crucial to examine its robustness and versatility. Finally, integrating the OAR-UNet model into clinical systems and workflows could enhance its utility in clinical decision-making and treatment planning.

6. Conclusions

Current manual segmentation techniques for head and neck organs-at-risk (OARs) are known for their inefficiency and high cost. Moreover, existing automatic segmentation methods struggle with effectively segmenting small organs and managing the complex interdependencies among anatomical structures. To address these challenges, we propose the OAR-UNet segmentation method, which is based on a U-shaped network architecture. OAR-UNet incorporates three key design components: the Local Feature Perception Module (LFPM), the Local Encoding Module (LEM), and the Cross-shaped Transformer Block (CSTB). These innovations significantly enhance the model’s efficacy in segmenting head and neck OARs, ensuring precise segmentation of smaller organs and maintaining stability and accuracy in the face of complex organ interactions. Evaluations of the SegRap2023 and PDDCA datasets demonstrate the superior performance of the OAR-UNet method. Specifically, OAR-UNet achieved a mIoU of 0.7478 and an mDice of 0.7822 on the SegRap2023 dataset, and a mIoU of 0.8545 and an mDice of 0.8942 on the PDDCA dataset. These results reveal that OAR-UNet outperforms both classic medical image segmentation methods and the latest state-of-the-art methods. It is anticipated that this advancement will provide robust technical support for the accurate segmentation of organs at risk in head and neck tumor radiotherapy planning. In future research, the method will be further optimized, and its potential applications in additional clinical scenarios will be explored to contribute more significantly to personalized and precise tumor radiotherapy.