Hierarchically enhanced feature fusion and loss prevention for prostate segmentation on micro-ultrasound images

Deep learning methods for prostate segmentation

In recent years, deep learning technology (Long, Shelhamer & Darrell, 2015; Ronneberger, Fischer & Brox, 2015; Chen et al., 2022; Lei et al., 2019; Zhu, Du & Yan, 2019; Wang et al., 2019) has achieved great success in various computer vision-related tasks such as object detection, image classification, semantic segmentation, and instance segmentation. Traditional methods usually rely on manually designed features (Jiang et al., 2023), while convolutional neural networks (CNNs) have the ability to automatically learn features without manual design. Many deep learning-based prostate segmentation methods for various imaging modalities have been proposed and have shown impressive performance, including TRUS (Lei et al., 2019; Wang et al., 2019; Yang et al., 2017; Ghavami et al., 2018; Orlando et al., 2020; Pellicer-Valero et al., 2021; Peng et al., 2022, 2023) and MRI (Chen et al., 2021b; Tian et al., 2018; Li et al., 2023). Yang et al. (2017) used a recurrent neural network (RNN) to learn the shape prior information of the prostate and implanted its core into a multi-scale automatic context scheme to achieve a gradual refinement of the prediction map details. Wang et al. (2019) used an attention mechanism to select multi-layer features integrated from different layers for utilization, and the prostate details were enhanced by suppressing non-prostate noise. Ghavami et al. (2018) replaced the convolutional layers in U-Net with residual network units to enhance the feature representation ability and optimize the prostate segmentation performance. Karimi et al. (2019) used the difference between models trained on simple and complex ultrasound scans to improve segmentation. Lei et al. (2019) combined the 3D supervision mechanism with the backbone network V-Net, and implemented the deep supervision mechanism by combining the binary cross entropy (BCE) and Dice loss functions in the training stage. Pellicer-Valero et al. (2021) combined checkpoint integration, deep supervision and neural resolution enhancement with the convolutional neural network (CNN) based on DenseNet-ResNet, and achieved excellent performance in prostate segmentation in multiple imaging modes. Peng et al. (2022) combined the improved principal curve method with the evolutionary neural network, and used redefined closed K-segment principal curve (RCKPC) and the modified cuckoo search and differential evolution algorithm (MCSDE) techniques at different stages to optimize the segmentation performance.

Previous prostate segmentation methods for 3D ultrasound images have addressed several challenges, including significant intensity variations, blurred prostate boundaries caused by speckle noise, low contrast between the prostate and surrounding tissues due to a low signal-to-noise ratio, and shadow artifacts. Micro-ultrasound imaging also faces difficulties with boundary clarity. Traditional approaches often lack real-time processing capability, as they typically operate on 3D volumes after the probe sweep. While micro-ultrasound offers much higher resolution and better image contrast than conventional ultrasound, it is particularly susceptible to artifacts from prostate calcifications, which complicate segmentation. Jiang et al. (2024) introduced MicroSegNet, the first deep learning-based method for prostate segmentation in micro-ultrasound images. It enables real-time processing of pseudo-sagittal views. However, MicroSegNet struggles to effectively capture and fuse highly differentiated features at multiple scales, primarily due to limited feature integration and inadequate modelling of dependencies among multi-scale features.

Transformers for medical image segmentation

Vaswani (2017) proposed the Transformer architecture for natural language processing tasks. Its self-attention mechanism allows capturing long-range dependencies and global context information, making it a remarkable success in text classification, machine translation, and language modeling, with subsequent extensions to medical image processing tasks. Parmar et al. (2018) made some modifications to the application of self-attention. The self-attention mechanism is not applied globally, but applied to the local neighborhoods of each query pixel. Wang et al. (2022) used the Transformer to optimize the UNet model, replacing the original skip connection with a connection consisting of Channel Cross fusion with Transformer (CCT) for multi-scale channel cross fusion and Channel-wise Cross-Attention (CCA) for guided fusion. Chen et al. (2021a) proposed TransUNet, which strong global context and models long-range dependencies by treating image features as sequences, while low-level CNN features can be well utilized through UNet. Zhang, Liu & Hu (2021) proposed Transfuse, which operates in parallel, combining the strengths of CNNs and Transformers. Cao et al. (2022) proposed Swin-UNet, which replaced the convolution operation in the U-Net encoder with a hierarchical Swin Transformer with shifted windows, aimed at enhancing global semantic information interaction. Valanarasu & Patel (2022) proposed UNext for medical image segmentation, which reduced the number of convolutional blocks in the shallow encoder and decoder, and the traditional UNet bottleneck layer was also replaced by tokenized MLP blocks to achieve a more lightweight structure. Ding et al. (2023) embedded multi-scale channel self-attention (CSA) in the skip connection of UNet. Jiang et al. (2024) combined multi-scale deep supervision in the TransUNet architecture (Chen et al., 2021a) to overcome the challenges of imaging artifacts caused by prostate calcification in micro-ultrasound images and unclear prostate boundaries. They also proposed an AG-BCE loss function, which assigned different weights to the prediction errors of identifying hard areas and easy areas, and more weights were assigned to the prediction errors of identifying more challenging areas. Our study leverages MicroSegNet, aiming two better handle the multi-scale features captured by its multi-scale deep supervision, optimize feature fusion, consider the dependencies between interrelated features extracted from different scales, and better handle multi-scale features in skip connections to reduce information loss and the imbalance between local details and global contextual information.

Overall structure

Figure 1 illustrates the proposed HEFFLPNet, which is based on the MicroSegNet model. MicroSegNet is based on the TransUNet architecture, a Transformer-based UNet model for medical image segmentation. MicroSegNet integrates a multi-scale deep supervision module with TransUNet, helping the model retain local information while capturing global contextual dependencies. It also introduces an AG-BCE loss function, which is incorporated across all intermediate layers and the final output layer to emphasize accurate segmentation of challenging regions during training.

Our proposed method is based on MicroSegNet, combining TransUNet with a multi-scale deep supervision module and AG-BCE loss. We propose two innovative modules: the TriCAFE module, which replaces the original skip connection, and the MSPMA module, positioned before the segmentation head of the final prediction output to process multi-scale features.

The TriCAFE Module enhances the function of skip connections, preserving essential features from earlier layers while enabling the network to adjust these features through learned weight. It mainly consists of three channels. The features of each channel are initially created through operations such as addition and concatenation, and then pass through channel attention (Woo et al., 2018) (CA) and spatial attention (Hu, Shen & Sun, 2018) (SA). The learned spatial weights modulate the features across the three channels, enabling adaptive feature fusion. We also designed a feature enhancement ratio (FER) mechanism to amplify the fused features to increase the weight of strong features in subsequent outputs, and finally concatenated with the decoder features to be fused to maximize the retention of original information, prevent information loss, and narrow the semantic gap between the encoder and decoder.

The MSPMA module allows the network to combine features across multiple scales effectively, thereby achieving multi-scale prediction that captures both high-level semantics and low-level spatial details. This module primarily consists of the PKI module (Cai et al., 2024), UpFA module, and Channel Aggregation (CA) block (Li et al., 2022). The PKI module captures multi-scale texture features and enhances high-level semantic features. The UpFA module upsamples and fuses low- and mid-level features through the attention mechanism, producing a final fused feature that matches the dimensionality of the lowest-level input feature. By extracting more representative information from multi-scale features, the UpFA module ensures that the network can capture high-level semantic information while selectively emphasizing to maintain low-level spatial details through the attention mechanism, ultimately improving the overall expressiveness of the model. The CA block aggregates channel information from low-level, spatially detailed features, adaptively forcing the network to encode expressive interactions that might otherwise be overlooked (Li et al., 2022). The feature enhancement ratio (FER) mechanism is also used here to achieve feature enhancement. Together, these mechanisms ensure that the network can produce accurate and well-calibrated predictions at multiple scales while leveraging semantic richness and spatial details. Overall, the MSPMA module processes and fuses features at each scale in a manner that maximally preserves and utilizes critical information, supporting robust and precise predictions across different scales.

MicroSegNet architecture

MicroSegNet (Jiang et al., 2024) is the first deep learning-based prostate segmentation method for micro-ultrasound images. It leverages the TransUNet architecture (Chen et al., 2021a) and integrates multi-scale deep supervision (Xie & Tu, 2015) to capture local details and broader contextual relationships at different scales.

Tri-cross attention with feature enhancement (TriCAFE) module

Skip connections enable the direct transfer of low-level features to high-level layers, thereby enhancing information flow, preserving features, and preventing gradient vanishing. However, the simple skip connection scheme in the TransUNet architecture, as applied in MicroSegNet, faces the challenge of modeling global multi-scale contexts. In this scheme, skip connections pass low-level detail features in the encoder layer with a smaller receptive field directly to the decoder, while the deeper layers in the encoder extract features with a larger receptive field, leading to a mismatch when directly fused with shallow features via simple skip connections, which brings about the problem of receptive field mismatch, and simple operations like addition, concatenation, or hybrid operations fail to reconcile this discrepancy. In addition, inadequate fusion processing of these two features may lead to information loss and an imbalance between local details and global contextual information, potentially diminishing segmentation accuracy by omitting crucial local or structural details at broader levels.

To address these challenges, we propose the TriCAFE module as a replacement for the simple concatenation operation in the original skip connections. The structure of the TriCAFE module is shown in Fig. 2, which visualizes the internal flow and interactions among components.

The input to the TriCAFE module consists of two features: the decoder-side feature map (denoted as $X_d$) and the corresponding encoder-side skip feature ( $X_e$). These two inputs are simultaneously processed through three parallel attention branches. In each branch, the inputs are combined via both element-wise addition and channel-wise concatenation, producing two intermediate representations $f_1=X_d\oplus X_e$ and $f_2=\mathrm{C}\mathrm{o}\mathrm{n}\mathrm{c}\mathrm{a}\mathrm{t}\left(X_d,X_e\right)$, where “ $\oplus$” denotes element-wise addition, and $\mathrm{C}\mathrm{o}\mathrm{n}\mathrm{c}\mathrm{a}\mathrm{t}(\cdot )$ represents concatenation along the channel dimension. This dual representation captures both aligned spatial structure and enhanced feature diversity.

In the first branch, attention weights are generated by applying channel attention (CA) and spatial attention (SA) to $f_1$, forming $W_1=\mathrm{C}\mathrm{A}\left(f_1\right)\oplus \mathrm{S}\mathrm{A}\left(f_1\right)$. The second and third branches generate $W_2=CA\left(f_1\right)\oplus SA\left(f_1\right)\oplus SA\left(f_2\right)$ and $W_3=SA\left(f_1\right)\oplus SA\left(f_2\right)$. These attention maps dynamically guide the fusion of encoder and decoder features, enabling adaptive weighting of each source through learned masks.

Each branch constructs a channel-level fusion output by blending the decoder and encoder features using the learned attention weights, which can be formally written as Eq. (1):

(1) $$F_{\mathrm{c}\mathrm{h}\mathrm{a}\mathrm{n}\mathrm{n}\mathrm{e}\mathrm{l}i}=f_1\oplus \left[W_i\otimes X_d\oplus \left(1-W_i\right)\otimes X_e\right],i=1,2,3$$where “ $\otimes$” denotes element-wise multiplication, and $\left(1-W_i\right)$ acts as the complementary gate to balance decoder and encoder feature contributions.

These three fusion results are then aggregated through element-wise addition and scaled by a FER (set to 1.5), which increases the discriminative power of the fused features. The enhanced result is then concatenated with the original decoder feature $X_d$. This process can be expressed as Eq. (2):

(2) $$\mathrm{T}\mathrm{r}\mathrm{i}\mathrm{C}\mathrm{A}\mathrm{F}\mathrm{E}\left(X_d,X_e\right)=\mathrm{C}\mathrm{o}\mathrm{n}\mathrm{c}\mathrm{a}\mathrm{t}\left(\left[F_{\mathrm{c}\mathrm{h}\mathrm{a}\mathrm{n}\mathrm{n}\mathrm{e}\mathrm{l}1}\oplus F_{\mathrm{c}\mathrm{h}\mathrm{a}\mathrm{n}\mathrm{n}\mathrm{e}\mathrm{l}2}\oplus F_{\mathrm{c}\mathrm{h}\mathrm{a}\mathrm{n}\mathrm{n}\mathrm{e}\mathrm{l}3}\right]\times \mathrm{F}\mathrm{E}\mathrm{R},X_d\right).$$

Channel attention (CA): Channel attention weights the features of each channel to highlight useful channel features and suppress redundant channels. The channel attention $CA(\cdot )$ is given in Eq. (3):

(3) $$CA\left(x\right)=\sigma \left({\mathbb{C}}_1\left(P_m\left(x\right)\right)+{\mathbb{C}}_2\left(P_a\left(x\right)\right)\right)\odot x$$where $\sigma (\cdot )$ denotes the Sigmoid activation, ${\mathbb{C}}_1(\cdot )$ is a convolutional layer with a kernel size of 1 to reduce the channel dimension 16 times, ${\mathbb{C}}_2(\cdot )$ is another convolutional layer to recover the original channel dimension, $P_m(\cdot )$ denotes the adaptive maximum pooling, $P_a(\cdot )$ denotes the adaptive average pooling, and $\odot$ refers to the Hadamard product.

Spatial attention (SA): Spatial attention generates an attention map based on the spatial dimension to enhance the features of key spatial positions in the image. The spatial attention $SA(\cdot )$ is given in Eq. (4):

(4) $$SA\left(x\right)=\sigma \left(\mathcal{H}\left(Concat\left(R_m\left(x\right),R_a\left(x\right)\right)\right)\right)\odot x$$where $\mathcal{H}(\cdot )$ refers to a convolutional layer with a kernel size of 7 and padding set to 3, $R_m(\cdot )$ and $R_a(\cdot )$ denote the maximum and average values obtained along the channel dimension.

Multi-scale prediction map attention (MSPMA) module

Following MicroSegNet, we employ multi-scale deep supervision (Xie & Tu, 2015) to integrate information across multiple scales during segmentation, enhancing robustness against variations in the shape, size, and appearance of input images. However, MicroSegNet certain limitations in handling highly differentiated features at different scales. Its inadequate modeling of interrelationships during feature fusion before the segmentation head output may lead to decreased performance, especially in capturing critical features such as shapes and boundaries. For instance, when small, intricate local features are combined with coarse global ones, important details are often masked or overlooked, which ultimately affects model accuracy. Moreover, MicroSegNet lacks in-depth modelling of multi-scale features, failing to fully utilize these features to represent complex structures accurately. This limitation can result in suboptimal performance when handling objects with complex geometric structures or edges, thereby reducing segmentation accuracy.

To address these challenges, we propose the MSPMA module, which enables specific and sufficient modeling of features at each scale, allowing the network to effectively integrate multi-scale features for predictions that capture both high-level semantics and fine-grained spatial details. As illustrated in Fig. 3, the MSPMA module mainly consists of four components: PKI block, UpFA, CA block, and the FER mechanism. The PKI block is applied to the shallowest feature map $X_1$, where it extracts global semantic cues and injects position-aware knowledge to enhance low-level representations. The output is further refined via a residual connection to preserve fine texture details, forming prediction map $P_1$. The UpFA module is responsible for fusing the mid- and low-level features $X_2$, $X_3$, and $X_4$, where $X_2$ and $X_4$ is passed through a 1 × 1 convolution and then combined with the upsampled $X_3$ by UpFA module. This fusion undergoes attention-guided refinement, allowing the model to adaptively focus on semantically meaningful regions while retaining spatial precision. To further enhance the fused features, we employ the FER mechanism, which introduces channel-wise modulation through element-wise multiplication to produce the output $P_2$. The $X_3$ is directly passed to generate $P_3$. Meanwhile, the deepest encoder feature $X_4$ is fed into the CA block, where global channel dependencies are modeled to recalibrate channel responses. This helps the network emphasize discriminative feature channels and suppress irrelevant ones. A residual connection is then applied to obtain the final output $P_4$. By designing parallel attention-enhanced paths for multi-scale prediction, the MSPMA module effectively bridges the semantic gap across features at different depths.

These mechanisms collectively ensure that the network can produce accurate and well-calibrated predictions at multiple scales while exploiting semantic richness and spatial details. The MSPMA module ensures that features at each scale are processed and fused in a way that maximizes the retention and exploitation of important information, thereby ensuring robust and accurate predictions across scales. Four prediction outputs

$P_i,i\in \left\{1,2,3,4\right\}$, can be formally written as Eq. (5):

(5) $$\begin{array}{c}\{\begin{array}{c}{P}_1=PKI\left(X_1\right)+X_1\hfill \\ P_2=UpFA\left(Conv\left(X_1\right),X_2,Conv\left(X_3\right)\right)\otimes FER\hfill \\ P_3=X_3\hfill \\ P_4=CA\left(X_4\right)+X_4\hfill \end{array}\end{array}$$where $PKI(\cdot )$ refers to PKI module, $UpFA(\cdot )$ refers to UpFA Module, $CA(\cdot )$ refers to CA Module, $Conv(\cdot )$ is a convolutional layer with a kernel size of 1, and $FER$ denotes the feature enhancement ratio, which is set to 2 here.

Upsample fusion attention (UpFA) module

The UpFA module in MSPMA is designed to fuse mid-level and low-level semantic prediction maps in multi-scale deep supervision. As illustrated in Fig. 4, the high-level features are fused with the mid-level features by addition after upsampling, and are fused with the low-level features by addition after upsampling through the PreActBottleneck block. The low-level features are output by addition after passing through the channel attention (CA) and the spatial attention (SA) respectively, and concatenated with the fused high-level and mid-level features, and finally fused with the high-level features after passing through the PreActBottleneck block by addition after $1\times 1$ convolution. The feature dimension of the final output aligns with the dimension of the lowest-level features in the input. This fusion strategy integrates multi-level feature representations, making the UpFA module particularly effective in enhancing the model’s segmentation performance, preserving detail, and strengthening global features. As a result, this module is highly suitable for multi-scale information processing in complex image segmentation tasks. The process can be formally represented by Eq. (6):

(6) $$UpFA\left(X_1,X_2,X_3\right)=U_{2x}(PA{B}_2\left(PA{B}_1\left(X_1\right)\right))+(Conv(Concat((U_{2x}\left(X_1\right)+X_2),\left(SA\left(X_3\right)+CA\left(X_3\right)\right)))+X_3)$$where $PA{B}_1(\cdot )$ refers to the PreActBottleneck block that keeps the number of intermediate channels the same, and then doubles the number of output channels, $PA{B}_2(\cdot )$ refers to the PreActBottleneck block that halves the number of intermediate channels and keeps the number of output channels the same.

PreActBottleneck: The PreActBottleneck (He et al., 2016a) block is an improved form of the bottleneck structure based on ResNet (He et al., 2016b), aimed at optimizing information flow in deep networks, thereby improving model performance without additional computational cost. This structure enhances gradient flow and feature representation. Its design is particularly well-suited for networks with multiple stacked layers, making it ideal for processing features rich in semantic information from high-level networks. The pre-activation structure minimizes information loss between layers, enabling deeper stacking to capture complex features efficiently.

Training loss

BCE loss is widely used as a loss function in binary image segmentation tasks. It quantifies the discrepancy between the predicted probability map $P$ and the ground truth segmentation $G$. The BCE loss is defined as shown in Eq. (7):

(7) $$BCE\left(P,G\right)=-\frac{1}{J}\sum _{i=1}^J\left(y_i\cdot \mathrm{l}\mathrm{o}\mathrm{g}{p}_i+\left(1-y_i\right)\cdot \mathrm{l}\mathrm{o}\mathrm{g}\left(1-p_i\right)\right)$$where $J$ refers to the number of pixels in the image, $y_i$ denotes the ground truth label for pixel $i$, and $p_i$ denotes the predicted probability for pixel $i$.

BCE loss assigns equal weights to all pixels in the image. However, this approach does not account for the increased difficulty of classifying pixels in areas affected by artifacts in micro-ultrasound images. This limitation makes BCE loss less suitable for prostate segmentation in these images. To address this, Jiang et al. (2024) introduced the AG-BCE function, which places greater emphasis on penalizing prediction errors in challenging areas. In our proposed HEFFLPNet, we followed MicroSegNet, adopted the AG-BCE function, combined it with multi-scale deep supervision, and applied the AG-BCE loss to the segmentation predictions generated by all layers. The AG-BCE loss is defined in Eq. (8):

(8) $$AG\text{-}BCE\left(P,G\right)=-{\displaystyle \frac{1}{J}\sum _{i=1}^J{w}_i\left(y_i\cdot \mathrm{l}\mathrm{o}\mathrm{g}{p}_i+\left(1-y_i\right)\cdot \mathrm{l}\mathrm{o}\mathrm{g}\left(1-p_i\right)\right)}$$where $w_i$ refers to the weight assigned to the pixel $i$. Consistent with the original definition in MicroSegNet (Jiang et al., 2024), we set $w_i=1$ for pixels in easy regions, and $w_i=4$ for pixels in hard regions, which has been empirically validated to enhance segmentation in challenging areas.

We apply the AG-BCE loss to segmentation predictions generated by all layers. The loss function for training HEFFLP can be written as Eq. (9):

(9) $$loss=AG\text{-}BCE\left(P_1,G_1\right)+AG\text{-}BCE\left(P_2,G_2\right)+AG\text{-}BCE\left(P_3,G_3\right)+AG\text{-}BCE\left(P_4,G_4\right)$$where $P_i$ prefers to the predicted segmentation at layer $i$, and $G_i$ prefers to the ground truth segmentation at layer $i$, $i=\left\{1,2,3,4\right\}$.

To provide a clear overview of the proposed architecture and its forward pass, we summarize the key steps of HEFFLPNet in Algorithm 1. This representation illustrates how the Transformer encoder, TriCAFE modules, and MSPMA module work together to produce multi-scale prostate segmentation predictions.

Dataset

In this study, we utilized the micro-US dataset published by Jiang et al. (2024) to evaluate the learning ability of HEFFLP, as well as the CCH-TRUSPS dataset (Feng et al., 2023) to assess the model’s generalization ability. The micro-US dataset comprises micro-ultrasound scans with human-prostate annotations from 75 patients who underwent micro-ultrasound-guided prostate biopsies at the University of Florida. The training set consists of 2,060 images from 55 patients, while the test set includes 768 images from 20 patients. Additionally, the CCH-TRUSPS dataset provides 1,820 2D ultrasound images of prostate cancer, collected from 364 male patients aged 40 to 90 at Chongqing University Cancer Hospital. For each patient, five representative images were taken from the apex, body, and base of the prostate. The preprocessing steps for the micro-US dataset, including downsampling images for efficiency, converts them to NumPy arrays, normalizes pixel intensities to 0–255, extracts 2D slices from 3D medical scans, resizes them to 224 × 224, and binarizes segmentation masks before saving them as PNG images. The DOI/URL of Micro-US dataset is provided by Jiang et al. (2024) at: https://zenodo.org/records/10475293, and the CCH-TRUSPS (Feng et al., 2023) dataset is introduced by Feng et al. (2023) together with their Multi-Stage FCN method in: https://doi.org/10.1016/j.bbe.2023.08.002.

Analysis of computational efficiency

As shown in Table 1, although HEFFLPNet introduces additional architectural modules compared to MicroSegNet, the overall computational overhead remains modest. Specifically, HEFFLPNet contains 94.234 million parameters and 29.738 GFLOPs, representing only a slight increase from MicroSegNet’s 93.235 million parameters and 24.683 GFLOPs. This marginal rise in complexity enables the integration of multi-scale attention mechanisms and enhanced feature fusion strategies, which we demonstrate in the experiment section to be beneficial for segmentation performance and robustness.

Evaluation metrics

We use the Dice coefficient to measure the relative overlap between the predicted prostate segmentation (P) and the true prostate segmentation (G), as defined in Eq. (10):

(10) $$DSC\left(P,G\right)={\displaystyle \frac{2\times \left|A\cap B\right|}{\left|A\right|+\left|B\right|}={\displaystyle \frac{2\times TP}{2\times TP+FP+FN}}}$$where $\left|A\right|$ denotes the number of positive pixels in the predicted segmentation images, $\left|B\right|$ denotes the number of positive pixels in the ground truth images, $TP$ is true positive instances, $FP$ is false positive instances, and $FN$ is false negative instances. In addition to using the Dice coefficient, we also apply the Hausdorff distance to evaluate the maximum separation between the two prostate boundaries, $P$ and $G$. To reduce the impact of minor outliers, we use the 95th percentile of boundary point distances (HD95) instead of the absolute maximum distance (HD). The definition is provided in Eq. (11):

(11) $$HD=max\left\{d_{PG},d_{GP}\right\},$$where $d_{GP}$ refers to the distance between $P$ and $G$, and $d_{PG}$ refers to the distance between $G$ and $P$.

Implementation details

During training, we retained the same hyperparameter settings as MicroSegNet to ensure a fair and consistent comparison under identical training conditions, using a patch size of 16, a batch size of 8, a learning rate of 0.01, a momentum of 0.9, a weight decay of $1e-4$, and a fixed random seed of 1,234 for all experiments. These hyperparameters were kept consistent for all models. We used the same dataset of 2,060 micro-ultrasound images from 55 patients as MicroSegNet for training, and the same 758 images from 20 patients as the test dataset to test the learning ability. All images were resized to 224 × 224 to match the standard input size of the Vision Transformer backbone. Intensity values were normalized to the [0, 1] range to stabilize training and maintain consistency with pre-trained encoder expectations and commonly adopted preprocessing pipelines. We further tested the generalization ability of the models on the CCH-TRUSPS dataset. To prevent overfitting, all models were trained for 30 epochs. The training and evaluation process was performed on an NVIDIA RTX 3070 laptop GPU with 8 GB of memory. We used Python 3.9.19 and PyTorch 2.3.0 in all experiments. In this article, we compare the proposed method with x state-of-the-art image segmentation models, including UNet (Ronneberger, Fischer & Brox, 2015), UNet++ (Zhou et al., 2018), TransUNet (Chen et al., 2021a), Swin-UNet (Cao et al., 2022), MicroSegNet (Jiang et al., 2024).

Quantitative analysis of learning ability

Table 2 presents the average Dice coefficient and Hausdorff distance for six segmentation models on 758 images of 20 test cases. Our HEFFLPNet outperforms the other five models in both metrics. By substituting the basic concatenation operation in the original skip connection with the TriCAFE module and incorporating an MSPMA module before the multi-scale deep supervised segmentation head for multi-scale feature processing, we enhance the Dice coefficient of MicroSegNet from 0.931 to 0.938 and reduce the Hausdorff distance from 2.17 to 2.12 mm. Compared with TransUNet, leveraged by MicroSegNet, our model improves the Dice coefficient by 0.015 and reduces the Hausdorff distance by 0.17 mm, demonstrating outstanding learning ability. Notably, the UNet model, which without Transformer integration, performs significantly worse, underscoring the advantage of Transformer-based architectures in capturing long-range dependencies essential for complex image segmentation tasks.

To further validate the robustness of the improvements, we conducted statistical significance tests between HEFFLPNet and the baseline MicroSegNet. The Dice improvement of HEFFLPNet (0.938 vs. 0.931) was statistically significant (paired t-test, p = 0.018; Wilcoxon, p = 0.026), with a moderate effect size (Cohen’s d = 0.564). For HD95, although HEFFLPNet achieved a slightly lower error (2.12 vs. 2.17 mm), the difference relative to the baseline was not statistically significant (p > 0.05).

Quantitative analysis of generalization ability

Table 3 presents a performance comparison of our HEFFLPNet model with five other methods on the CCH-TRUSPS dataset, which was previously unseen. HEFFLPNet achieves a Dice coefficient increase from 0.903 to 0.914 over MicroSegNet, while reducing the Hausdorff distance from 4.22 to 3.63 mm. In comparison to the TransUNet model leveraged by MicroSegNet, our model improves the Dice coefficient by 0.011 and reduces the Hausdorff distance by 0.56 mm, demonstrating its superior generalization capability. Notably, the UNet model, which lacks Transformer integration, continues to exhibit weak generalization, underscoring the advantages of a Transformer-based architecture in capturing the long-distance dependencies essential for complex image segmentation tasks.

Based on this analysis, our HEFFLPNet model demonstrates exceptional learning and generalization performance in the challenging task of prostate segmentation on micro-ultrasound images, surpassing other SOTA methods.

Qualitative analysis

To comprehensively evaluate the effectiveness of our proposed method, we selected two high-performing representatives from five SOTA methods and compared them with our HEFFLPNet based on visual results. As illustrated in Fig. 5, micro-US ultrasound images demonstrate superior quality, fewer artifacts, and more defined prostate boundaries compared to traditional ultrasound images. Among the SOTA methods, our HEFFLPNet provides more accurate and smoother boundary segmentation for prostates of various sizes and shapes in both micro-US and traditional ultrasound images, underscoring the model’s accuracy and its robustness to artifacts. This improvement is largely due to our replacement of the original skip connection in TransUNet with an innovative TriCAFE module and the integration of an MSPMA module ahead of the multi-scale deep supervised segmentation head, enhancing multi-scale feature processing. By hierarchically reinforcing feature fusion and preventing loss, we preserve semantic richness and spatial detail across scales, which minimizes information loss and facilitates robust, precise predictions.

Specifically, as shown in Fig. 5, in terms of boundary accuracy, HEFFLPNET demonstrates superior performance, particularly in the 6th and 8th rows (CCH-TRUSPS), where it closely adheres to the true anatomical boundaries with smooth and well-fitted contours. In contrast, MicroSegNet exhibits slight edge blurring in these rows, while TransUNet shows boundary over-expansion and reduced curvature conformity in the 8th row. For shape preservation, HEFFLPNET maintains high structural fidelity, as evidenced in the 1st and 5th rows (micro-US) and the 7th row (CCH-TRUSPS), producing segmentations that closely match the ground truth without noticeable geometric distortions. MicroSegNet and TransUNet, however, tend to over-segment or under-segment, especially around object boundaries, leading to shape inaccuracies. Overall, HEFFLPNET consistently produces results with higher conformity to the ground truth across both imaging modalities, indicating stronger generalization ability and robustness in multi-domain prostate segmentation tasks.

Ablation studies

We conducted ablation experiments to assess the impact of three core components in HEFFLPNet: the MSPMA module, the UpFA sub-module, and the TriCAFE module. Table 4 presents the results on both the micro-US and CCH-TRUSPS datasets using mDice and HD95 as evaluation metrics.

When the TriCAFE module is removed, the model’s performance on CCH-TRUSPS declines notably: mDice drops from 0.914 to 0.905 and HD95 increases sharply from 3.63 to 4.44. This suggests that TriCAFE plays a critical role in improving segmentation quality in more complex scenarios, particularly those involving anatomical ambiguity and fine-grained boundaries. In contrast, on the micro-US dataset, the changes are more moderate (mDice from 0.938 to 0.934, HD95 from 2.12 to 2.15), indicating that TriCAFE’s benefits are especially pronounced in datasets with more challenging spatial structures. This aligns with our design goal for TriCAFE: to strengthen skip connections by reducing the semantic gap between encoder and decoder features and better preserving spatial consistency.

Similarly, the removal of the full MSPMA module or its internal UpFA component also leads to consistent degradation in performance, particularly in HD95, which reflects boundary precision. For instance, removing MSPMA increases HD95 to 2.19 on micro-US and to 3.73 on CCH-TRUSPS, indicating that the module plays a key role in fusing features across scales and preserving structural information throughout the network. MSPMA is designed to model multi-scale feature representations explicitly, ensuring that both fine-grained local textures and broader contextual cues are retained and utilized during decoding. This helps the model better delineate ambiguous or low-contrast regions.

The UpFA component within MSPMA is particularly important for effectively integrating mid-level and low-level features. Its removal results in a decrease in mDice on the micro-US dataset (from 0.938 to 0.932) and a slight degradation in HD95. This reflects UpFA’s role in selectively enhancing useful spatial details while avoiding the introduction of noise. Through attention-based fusion, UpFA ensures that important structural features from different encoder stages are emphasized appropriately, improving segmentation coherence and boundary localization.

These results validate the complementary contributions of each module. While TriCAFE enhances skip connection fidelity, MSPMA and its internal UpFA sub-module improve hierarchical feature integration and detail preservation. Together, these modules allow HEFFLPNet to produce more accurate and robust segmentation results, particularly in anatomically variable or noisy clinical scenarios.

We present visualization results in Fig. 6 to further demonstrate the effectiveness of the proposed HEFFLPNet architecture and the impact of its core modules, including MSPMA, UpFA, and TriCAFE. As shown, the removal of any individual component leads to a noticeable decline in segmentation quality, particularly in boundary accuracy and structural completeness.

Specifically, when the MSPMA module is omitted, the resulting masks exhibit incomplete contours and inaccurate boundaries, especially in cases with low contrast or blurred edges. This can be attributed to the loss of effective multi-scale feature modeling and fusion, which are crucial for precise delineation in complex anatomical regions.

Removing the UpFA module leads to under-segmentation in several cases, indicating a weakened ability to preserve and integrate fine spatial details from earlier encoder stages. Without UpFA, the network’s capacity to selectively highlight relevant spatial features is significantly reduced.

The absence of the TriCAFE module results in a decline in boundary accuracy and structural completeness. This suggests that TriCAFE plays an essential role in strengthening skip connections and bridging the semantic gap between encoder and decoder features, thereby maintaining spatial coherence.

These visual results align with our quantitative findings and confirm that each module in HEFFLPNet contributes complementarily to its overall segmentation robustness and precision across different datasets.