SPGNet: A Shape-prior Guided Network for Medical Image Segmentation

1 Introduction

Medical image segmentation has always been critical to medical image processing. Currently, most mainstream methods focus primarily on high-precision pixel-level supervision. Despite significant achievements in a variety of segmentation tasks in differing domains, limitations persist when dealing with medical images. Primarily, anatomical structures in medical images often exhibit shape patterns and geometric information that generic pixel-level supervision could fails to leverage fully, especially when shape regularities are prominent and perceptibly advantageous.

Past studies indicate integrating shape prior knowledge could benefit traditional segmentation algorithms [Nosrati et al., 2016]. For instance, introducing shape priors to level set contour evolution techniques proved to enhance the accuracy of segmentation [Chen et al., 2002]. Integrating elastic shape priors into frameworks was also shown to be able to align the outcomes closer to the actual shape variations [Schoenemann et al., 2007]. However, these conventional segmentation methods often struggle with complex medical image scenarios due to constraints in data distribution assumptions and susceptibility to noise and contextual information.

∗Corresponding author

Code will be available at https://github.com/Chris97s/spgnet

Deep learning methods, in comparison, exhibit greater flexibility in adapting to diverse shapes and backgrounds [Bohlender et al., 2021]. It has been shown that this paradigm of segmentation methods could also benefit from shape priors to better cater for specific task requirements. For example, a model based on sparse representation and local repulsive deformation was proposed for normalizing the former deep convolutional neural network segmentation and constraining the segmentation results within an effective shape domain [Xing et al., 2015]. Although post-processing has often demonstrated effective for further improving segmentation [Li et al., 2017; Medley et al., 2019], its usefulness still hinges on the dependability of the segmentation model. Another model [Lee et al., 2019] based on template deformation implemented segmentation by deforming the shape prior template. However, this approach may be constrained by the expressiveness of the priors and would likely be undermined by structural variations in medical images, such as when handling images of various organs with distinct shapes and morphological characteristics.

While these methods aim to constrain segmentation using shape priors, they have not entirely endowed deep neural networks with the capacity to perceive shape patterns in the data. We propose to enable deep learning methods to fully incorporate learned priors from training data as a form of regularization. This enables segmentation networks to be influenced by prior shape knowledge when classifying pixels. As shown in Figure 1, with shape prior guidance, attention in shallow layers concentrates completely on the target area. In deep layers, the excellent alignment between the attention map and the target area is maintained, while the most activated attention areas shift towards the edges which serves to improve the accuracy and smoothness of segmentation. Conversely, a network without guidance exhibits fragmented and misaligned attention highly susceptible to noise and variability.

As shown in Figure 2, our end-to-end dual-path collaborative network integrates a multi-class statistical shape model for incorporating a wealth of shape priors. The segmentation path focuses on dense pixel-level classification, while the shape-prior path regresses prior flows. We introduce a collaboration module establishing interactions between encoders to enhance robustness by exchanging features. Simultaneously, by collaboratively learning multi-class statistical shape model deformation, the segmentation encoder is enabled to perceive shape distribution and utilise this explicit shape prior knowledge to guide segmentation. We also use cropped local features from the segmentation path to refine shapes across different scales, guiding attention to focus on boundaries.

Our contributions are summarized as follows:

• We propose a novel dual-path collaborative segmentation network, SPGNet, which embeds explicit and diverse shape priors. The dual-path structure enhances the representation capability of segmentation path encoders. The segmentation path, guided by the explicit shape priors, reinforces shape understanding and enhances attention at target edges, addressing single-path deficiencies in capturing shape features. We also design a cluster strategy to learn shape regularities from the training set.
• We explore the collaborative effects of SPGNet and demonstrate the effectiveness of the shape-prior path in improving segmentation accuracy. Specifically, we validate the efficacy of each component within the shapeprior path.
• We evaluate SPGNet on a chest X-ray dataset with prominent shape regularities and a breast ultrasound dataset with potential regularities. Results demonstrate superior accuracy over baselines and existing state-of-the-art methods, particularly in edge smoothness. We also validate its adaptability to medical images with low signal-to-noise ratios, blurry boundaries, and significant shape and positional variations of lesions. The superiority of our method has been demonstrated.

2 RelateWorks

Shape Clustering and Statistical Modeling. Shape clustering groups shapes by extracting descriptors and clustering based on similarity distances. For instance, a skeleton-based approach captures intrinsic structural information for sameclass shapes, clustering using a node-matching matrix [Shen et al., 2013]. However, skeleton-based methods often express relatively coarse shape features. In contrast, a hierarchical clustering method is used in a different study for contour-based shapes to learn probabilistic models from shape clusters [Srivastava et al., 2005]. In another approach, to address unlabeled longitudinal shape data, a flexible nonlinear mixture model is established by learning average shape trajectories and variances for each cluster [Debavelaere et al., 2020]. For instance segmentation, k-means clustering of training masks are used to obtain centres of shape clusters, establishing a linear prior model e.g., [Kuo et al., 2019]. A pipeline combining segmentation, clustering, and modeling has also been proposed [Bruse et al., 2017]. With the rise of deep neural networks, exploring clustered shape information with statistical models still has significant potential.

Shape-prior Guided Segmentation. Numerous methods have attempted to leverage shape priors for segmentation. For example, to address blurred overlapping regions in the cell cytoplasm, a generator utilizing a prior template to generate masked was proposed [Song et al., 2020]. Similarly, a cyclic registration network was also designed to integrate anatomical context specificity with priors [Jiang and Veeraraghavan, 2022]. For a different medical application, a deep neural network was designed to predict PCA layers for improving the segmentation of the left ventricle in ultrasound images [Milletari et al., 2017]. Generating threshold-based priors and optimizing outputs via a spatial transform network was also proposed as another strategy for harnessing prior shape knowledge [Zhao et al., 2021]. More recently, it was also found that introducing an additional branch transforms visible regions into complete areas through supervision, and therefore facilitates holistic shape understanding [Gao et al., 2023]. Similarly, another study proposed a generative invariant shape prior network that introduced a branch to learn invariant priors, mimicking human perceptual learning of basic shapes [Li et al., 2023]. Other than considering two-dimensional shapes, a different method used three-dimensional reconstructed shapes as priors and reconstructed occluded objects before projecting them to predict complete mask [Li et al., 2022]. However, these methods may struggle to handle complex and heterogeneous targets due to insufficient shape diversity [Zhao et al., 2021; Jiang and Veeraraghavan, 2022]. Implicit priors may fail to generalize [Gao et al., 2023; Li et al., 2023]. In contrast, our approach establishes an explicit multi-class shape statistical model to guide segmentation.

3 Methodology

3.1 Overview

In Figure 2, we introduce SPGNet, a novel image segmentation algorithm with embedded shape priors. Section 3.2 discusses offline multi-class shape statistical modeling, Section 3.3 details SPGNet's internal modules and Section 3.4 introduces the hybrid loss function for network training.

3.2 Offline Multi-class Shape Statistical Modeling

Shape Preparation. As shown in Figure 3, we derive shape data from mask annotations, focusing on regions with heightened curvature to build a statistical model with prominent shape features. The shape generation process involves: (i) Employing B´ezier curves to smoothly fit mask contours, yielding a set of smoothed contour points. (ii) Determining the number of sampling points as p and calculating the absolute curvature for each point, followed by normalization. A base value is introduced to prevent neglecting points with extremely small curvature. The ratio of the sum of normalized curvatures to the number of sampling points serves as the sampling distance. (iii) Iterating through contour points, saving the current sum of curvatures and the number of sampled points, denoted as k. When the sum exceeds k times the sampling distance, the point is saved as a sample point and k increments by 1. The traversal ends when k reaches p. The sampled p points constitute the shape contour, represented as Si = ((xi1 , yi1 ), ..., (xip , yip ))T ∈ Rp×2 .

Procrustes Shapes Agglomerative Clustering. We employed the agglomerative clustering method, utilizing the Procrustes shape distance as the shape similarity metric, which requires aligning shapes before computing distances. The Procrustes shape distance calculation between two aligned shapes S1 and S2 is calculated by:

In Figure 3, the agglomerative clustering of shapes involves these steps: (i) Standardizing and aligning all shapes using the Procrustes analysis [Cootes et al., 1995] in the training set. (ii) Calculating the pairwise distance matrix using Procrustes shape distance. (iii) Determining the number of clusters k. Starting with k = 2 and iteratively incrementing k, we select the k value that maximizes the variance of distances between k cluster centers as the initialization. We limit k to a maximum of 100 to prevent excessive redundancy in categories. (iv) Initiating agglomerative clustering with complete linkage for calculating inter-cluster distances, minimizing dissimilar shape aggregation between clusters. After all iterations, we obtain k shape clusters.

Multi-class Shape Priors andModeling. We compute the mean shape within each cluster and conduct principal component analysis (PCA) on cluster shapes to derive eigenvectors representing major variations. The number of eigenvectors per cluster aligns with the maximum across clusters, denoted as t. Concatenating mean shapes and eigenvectors yields shape priors, denoted as Priors = {(meani, ρi)ki=1 } ∈ Rk×(1+t)×p×2 . Simultaneously, shape category labels C are obtained from the shape clusters. The Active Shape Model (ASM) [Cootes et al., 1995], a statistical model for shapes, necessitates a set of shape training samples Set = {(Si )ni=1 } for model construction. Initially, we compute the average shape vector for all shapes:

(2)

We employ PCA for dimensionality reduction on the n shape training samples, creating an approximate model that encapsulates the entire training dataset:

(3)

where Φ = (ρ1, ρ2, ..., ρj) ∈ Rj×p×2 represents the first j major eigenvectors corresponding to the eigenvalues of the covariance matrix. b ∈ R1×j denotes the shape deformation parameters. Subsequently, we can utilize a neural network to predict the deformation parameters b, facilitating the fitting of the shape statistical model to the shapes in the training set. Additionally, we introduce an affine transformation function a with parameters γ = (sγ, θγ, (t1)γ, (t2)γ) ∈ R1×4 . The four parameters include the scale parameter sγ, rotation parameter θγ, and translation parameters (t1)γ and (t2)γ. For a set of points S = (Sx, Sy) ∈ Rp×2 , affine transformations can be described as:

In summary, an ASM Transformation model is denoted as:

The ASM Transformation function is embedded in the network computational process. Priors are precomputed offline before training and inference, entering the network as constant parameters.

3.3 SPGNet

Dual-path CollaborationModule(DCM). In the encoder sections of both paths, we aim for mutual attention during training, allowing the segmentation path to attend to learned features from the shape-prior path and vice versa. To achieve this, we introduce the dual-path collaboration module (DCM), incorporating spatial and channel attention for enhanced feature interaction. As shown in Figure 4, the DCM takes features Fi ∈ RCi×HFi ×WFi and Si ∈ RCi×HSi ×WSi from the two paths as input, with Si resized to match Fi dimensions. Spatial attention is computed by averaging along the channel dimension and applying the sigmoid activation function σ(·). For channel attention, global average pooling (GAP) is applied along the spatial dimension, and linear layers, along with the sigmoid activation function σ(·), calculate channel attention. The concatenated features, after passing through consecutive convolutional layers, are multiplied separately by spatial attention and channel attention, before the results are summed. This process introduces spatial and channel attention, fostering interactive features. A residual structure combines these interactive attention features with features from each path. The resulting features Fi+1 and Si+1 are obtained after passing through the encoder blocks of each path.

Segmentation Path. In Figure 2, the segmentation path employs an encoder-decoder architecture with skip connections. All subsequent encoder blocks, except the initial one, receive interactive attention features from the DCM. Each encoder block is composed of two sets of sub-blocks and a max-pooling layer. The sub-blocks include a convolutional layer, ReLU activation and batch normalization. In the decoder section, features from the preceding layer are initially upsampled using bilinear interpolation. Following concatenation with the skip connection input from the encoder, the combined features enter the decoder. The ultimate layer of the decoder employs a convolutional layer to derive the probability distribution map for dense pixel-wise classification.

ASMTransformation. We acquired prior shape knowledge, denoted as Priors, which is loaded into the network before training. We utilised the ResNet18 [He et al., 2016] as the shape classifier backbone in the classification network. During training, the classifier learns the shape category weights for input images, and the appropriate shape prior is obtained by multiplying the classifier output Classes ∈ Rk with Priors. Specifically, the classifier output is used to calculate the shape mean mean ∈ R1×p×2 and t eigenvectors ρ ∈ Rt×p×2 . In the shape-prior path, the shape encoder block employs a combination of spatially separable convolution and convolutional layers with a stride of 2, focusing on capturing edge and shape features. The output of the last encoder block is fed through continuous convolutional layers to obtain deformable parameters δ ∈ R1×t and affine parameters α ∈ R1×4 . The shape L0 is obtained through ASM Transformation:

Shape-guided Feature Cropping. As illustrated in Figure 5, given the input shape points L(t-1) ∈ Rp×2 and the features F(4-t) from the segmentation path, where t = 1, 2, 3, and setting the relative length of the clipping patch l, we first convert F(4-t) into f(4-t) ∈ Rc×wf(4-t) ×hf(4-t) through convolution. For a shape point in the spatial direction of f(4-t), denoted as (Lt-1 x ,Lt-1 y ), we clip out c patches, each containing n×n sampled feature points. The feature at each sampled point is calculated on f(4-t) using bilinear interpolation, and we retain the bottom-left relative position coordinates of the patch, denoted as (Lt-1 x ,Lt-1 y )lt. By using p shape points, we obtain the clipped feature sequence Ωt ∈ R(p×c)×n×n and the bottom-left position sequence Lt-1 lt ∈ Rp×2 .

Long-range Refinement Module(LRM). The global correlation of shape points is crucial for accuracy. In Figure 6, we introduce the Long-range Refinement Module (LRM), which takes Ωt , Lt-1 lt , and the patch side length l as inputs. This facilitates the adjustment of shape points within the patch. LRM consists of a parallel structure employing convolution and a multilayer perceptron. Additionally, it utilizes a transformer encoder [Zheng et al., 2021] to handle the feature sequence. In the first direction, a n × n convolutional layer extracts global features from Ωt , followed by flattening and passing through a multilayer perceptron with two linear layers and ReLU activation, resulting in offset1 ∈ Rp×2 . In the second direction, another n × n convolutional layer extracts Ωt and projects it onto token sets with positional embedding. Afterwards, an 8-layer transformer encoder with multi-head self-attention is employed to establish global feature correlations. The output, reshaped into linear features, passes through a multilayer perceptron, yielding offset2 ∈ Rp×2 . The refined shape points are computed as follows:

In our model, the coarse output of the ASM Transformation is denoted as L0 , and the outputs of the three stages of LRM are (L1 , L2 , L3 ).

3.4 Hybrid Loss Function

We formulated a hybrid loss function for training SPGNet. The total loss L is a weighted sum of the segmentation loss Lseg derived from the segmentation path and the shape loss Lshape, along with the classification loss Lc. The segmentation loss is calculated as follows:

where ˆ Y is the output probability segmentation map from our segmentation path, and Y represents the ground truth for the probability segmentation mask. The losses Lce and Ldice correspond to the cross-entropy and dice loss [Li et al., 2019], respectively. The classification loss is calculated as:

where ˆ C represents the predicted shape category probabilities from classifier, and C refers to the ground truth for shape categories. Lbce corresponds to the binary cross-entropy loss [Chen et al., 2023]. The shape loss calculation is as follows:

where L represents the ground truth for the shape, and Ll1 refers to the L1 loss [Feng et al., 2018]. The total loss L is defined as:

4 Experiment

4.1 Dataset and Pre-Processing

In our comprehensive evaluation, we utilized two publicly available datasets: chest X-rays and breast ultrasound image data. Below, we outline the databases employed for these experiments.

JSRT Dataset. The JSRT database [Shiraishi et al., 2000] consists of 247 high-resolution X-ray images, comprising 154 conventional chest X-rays with lung nodules selected from 14 medical centers and 93 chest X-rays without lung nodules. The dataset provides manually annotated masks for three anatomical structures of interest: the left lung, right lung, and heart. We separated the masks for these three anatomical structures for individual evaluations.

Breast Ultrasound Dataset (BUS). The Breast Ultrasound Dataset (BUS) [Al-Dhabyani et al., 2020] includes 780 breast ultrasound images categorized into three classes: 133 normal, 437 benign, and 210 malignant images. For our experiments, we combined samples from both benign and malignant classes. The dataset offers manually annotated masks for tumor regions.

4.2 Implementation Details andMetrics

We implemented all evaluation methods on a server equipped with an NVIDIA GeForce RTX 4090 GPU. For a fair quantitative comparison, all methods for dense pixel-level classification underwent evaluation using the same 5-fold crossvalidation scheme, maintaining a standardized input resolution of 256 × 256. To prevent overfitting, consistent data augmentation was applied in experiments, including random rotation, random vertical flipping, and random changes in brightness and contrast. We performed 150 epochs of training on the JSRT dataset and 300 epochs on the Breast Ultrasound Dataset (BUS) while keeping the remaining training hyperparameters consistent. The batch size was set to 16, utilizing the Adam optimizer. The initial learning rate was 0.0001, with weight decay at 0.0005, and a learning rate decay of 90% every 15 epochs. During the validation phase, we assessed segmentation performance using the Dice coefficient(%) (Dice), Jaccard index(%) (Jaccard), and 95% Hausdorff Distance(mm) (HD95).

4.3 Comparison with State-of-the-art

Results on JSRT. Table 1 summarizes the experimental outcomes for the left lung, right lung, and heart components in JSRT. Our method, employing a shape point sampling of 128, surpasses other dense pixel-level classification approaches in the average Dice score. Although the improvement in average Dice may not be as pronounced compared to other methods, as depicted in Figure 7, our approach distinctly excels in edge smoothness and accuracy. Moreover, in methods based on points regression, our approach (with a shape point sampling of 64) achieves significantly higher average Dice scores in the shape path than other state-of-the-art points regression methods. Training was conducted in both scenarios, using our parameters and the optimal parameters specified in their respective papers. Compared to the state-of-the-art method [Bransby et al., 2023], our method demonstrates improvements in Dice scores of 1.56%, 2.44%, and 1.83% for the left lung, right lung, and heart, respectively. These results highlight the superior performance of SPGNet on datasets exhibiting prominent shape patterns.

Results on BUS. Table 2 presents the experimental results on the BUS dataset. The segmentation path of our method(with a shape point sampling of 128) achieved an average Dice score superior to other state-of-the-art methods based on dense pixel-level classification. For SPGNet, the Dice and Jaccard scores reached 78.4% and 69.7%, respectively, This represents an improvement of 0.72% and 0.76% compared to the previous state-of-the-art method [Chen et al., 2023]. As illustrated in Figure 7, our method excels in accurately identifying the location of breast tumors and maintaining overall shape and edge smoothness. The experimental results demonstrate that SPGNet can leverage shape priors to enhance performance on datasets with underlying shape patterns.

4.4 Ablation Study

Ablation of Each Component in Shape-prior Path. Table 3 illustrates the performance enhancement trends of the segmentation path with the introduction of each component in our dual-path collaborative network compared to the baseline. To validate each module, various coefficients in the mixed loss function were adjusted to reflect the combined effects. It was observed that three factors primarily contributed to the improvement: (i) DCM (Dual-path Collaboration Module): Coefficients λ1 and λ2 for Lseg were set during training to highlight the impact of DCM. Utilizing DCM, which enables collaborative learning of input image features by dualpath encoders, yielded enhancements of 1.65% for BUS and 0.03% for the right lung, emphasizing the effectiveness of our dual-path collaborative network structure. (ii) Supervision on Shape: Building upon the first factor (DCM), coefficients λ3 and λ4 for Lc and Lshape were introduced. This supervises the shape weights and guides the ASM Transformation output shape, leading to a significant 2.87% improvement for BUS and 0.12% for the right lung. This underscores the effectiveness of our offline-modeled, multi-class shape priors in guiding the segmentation network, demonstrating that deep learning segmentation networks, guided by diverse shape prior information, can significantly compensate for deficiencies in shape perception. (iii) Multi-stage Shape Refinement: Building upon the second factor (DCM+Tasm), coefficients λ5, λ6, and λ7 were sequentially introduced for each stage of shape refinement. With the introduction of each refinement stage, there has been a relative improvement of 0.3%, 0.81% and 1.07% for BUS, and 0.02%, 0.05% and 0.07% for the right BUS Heart Right Lung Number of shape points lung, respectively, compared to the result of (DCM+Tasm). This demonstrates that our combination of multi-stage shape refinement modules (Crop+LRM) is effective. The experimental results indicate that the shape-prior path could lead to improved performance in the segmentation path.

Figure 9: Ablation on the number of shape points in BUS and JSRT (Heart) was performed, with fold 1 of BUS, fold 3 of JSRT (Heart) and fold 4 of JSRT (Right Lung) as the validation sets.

Ablation on the Number of Clustering. We validated the impact of the number of clusters in Procrustes shapes agglomerative clustering on segmentation performance. Our criterion is to select the value that maximizes the variance of cluster distances as the offline clustering quantity for SPGNet. As shown in [Figure 8 (a)], when we chose the clustering quantity k value according to our rule (our rule selected 30), the variance of cluster centers reached 9.31%, and SPGNet achieved a dice value of 79.63%. The segmentation performance is higher compared to situations when we chose the values of the cluster of numbers (k) around the one selected according to our criterion. A similar trend can be observed in [Figure 8(b)]. This suggests that selecting the value that maximizes the variance of cluster distances can approximate an optimal quantity of shape types, avoiding situations with insufficient categories or redundancy. This serves as a suitable rule for selecting the clustering quantity to enhance segmentation performance.

Ablation on the Number of Shape Points. We assessed the impact of varying quantities of shape points on segmentation performance, as depicted in Figure 9. Seven scenarios, ranging from 32 to 256 shape points, were selected to observe their influence on segmentation performance. As the number of shape points increased, our segmentation performance gradually improved, reaching optimal performance at 128 points. It can be observed that the number of points (e.g., 96, 128, 192, 256) that effectively represent shapes leads to a significant improvement in model performance, with little variation among them. Therefore, selecting sufficient and reasonable numbers of points can better promote model performance improvement.

5 Conclusion

In this paper, we introduce a novel shape-prior guided segmentation network. The core concept is to enable the network to grasp the shape distribution within the data by learning from statistical shape model enriched with shape prior knowledge, thereby enhancing segmentation accuracy. We validate the effectiveness of the proposed network through extensive experiments on two public datasets.

Acknowledgments

This study was supported by grants from the National Natural Science Foundation of China (Grant No. 81873631, 81370866, 81070612), Guangzhou Science and Technology Planning Project (Grant No. 202002020047), and Guangdong Basic and Applied Basic Research Foundation (Grant No. 2019A1515011078), and the Royal Society (Grant Reference: IEC\NSFC\201041).