A Novel Attention-based Network for Geometry Reconstruction with Error Estimation from Medical Images

Abstract

Instance segmentation of anatomical structures from medical images can help enhance clinical outcomes such as disease diagnosis, surgical planning accuracy, and treatment efficacy. However, since segmentation masks often lack point-to-point correspondence between patients, instance segmentation masks may not be directly applicable for clinical studies that require measurement of medical parameters defined on an anatomical shape atlas. For such applications, meshes with correspondence between patients are preferred representations of object geometries. The conversion from segmentation masks to meshes can be error-prone due to segmentation artifacts, and therefore, it is desirable to directly obtain mesh representations of object geometries from medical image data, bypassing segmentation masks. In this work, we propose novel attention-based neural networks for geometry reconstruction and error estimation, which offers a direct pathway from medical images to high-quality mesh representations. We introduce an innovative attention-based feature extraction network and incorporate image self-attention and shape self-attention with cross-attention between them to capture consistent features. Based on the extracted features, we develop a geometry reconstruction network that deforms a mesh template to reconstruct the geometry of the object, which automatically ensures mesh correspondence. In addition, we design a shape error estimation network to evaluate the reliability of the output from geometry reconstruction, that is, to estimate the point-to-point error of a reconstructed geometry. We demonstrate our approach for the application of lumbar spine geometry reconstruction and compare our geometry reconstruction network with UNet++, UTNet, Swin UnetTR, SLT-Net, and nnUNet, using the Dice metric. In this application, our geometry reconstruction network has much higher accuracy and artifact-free segmentation results, and our shape error estimation network facilitates quality control for clinical use. The source code is available at https://github.com/linchenq/SPIE2025-GoemReconstruction-with-ShapeErrorNet.

1. INTRODUCTION

Instance segmentation of anatomical structures from medical images can help improve clinical outcomes such as disease diagnosis, surgical planning precision, and treatment efficacy.¹ For example, accurate segmentation of the lumbar spine can help assess the state of spinal degeneration,² and automated liver segmentation provides surgeons with an overview of the spatial relationships of the liver organ and its internal structures and reduces the risks of surgical resection,³ and segmentation of cardiac structures can help detect earlier cardiac abnormalities.⁴ Because of the lack of point-to-point correspondence between pixel-level segmentation masks of different patients, instance segmentation masks may not be applicable for clinical studies that require measurement of medical parameters defined on an anatomical shape atlas. For clinical applications that require stress or hemodynamic analysis,^5–7 segmentation masks need to be converted to meshes. Furthermore, if the mesh correspondence is established (i.e., the same number of nodes/points and elements with the same connectivity), then detailed geometry/stress analysis and comparison^6,8 can be directly performed at the population level, and virtual clinical trials on virtual patients can be facilitated.⁹

Currently, deep learning approaches, particularly convolutional neural networks (CNNs) and transformers, have demonstrated superior performance across various medical tasks, including medical image segmentation. Networks such as UNet++,¹⁰ UTNet,¹¹ Swin UnetTR,¹² SLT-Net,¹³ nnUNet¹⁴ have been proposed for image segmentation on different medical imaging modalities and organs. As illustrated in Fig.1, despite improved accuracy, these segmentation models often generate artifacts such as fragmented, overlapping, and inconsistent regions.¹⁵ For applications that need mesh representations of object geometries, the conversion from segmentation masks to meshes can be error-prone because of the common segmentation artifacts. Thus, it is desirable to directly obtain the mesh representation of an object’s geometry from medical image data, bypassing the segmentation masks.

Figure 1.

The comparison between segmentation and geometry reconstruction methods.

In this work, our main contributions are as follows. (1) We propose an innovative attention-based neural network as the backbone for feature extraction, which utilizes image self-attention and shape self-attention as well as cross-attention between image and shape to capture consistent features. (2) Based on the feature extraction network, we design a geometry reconstruction network that deforms a mesh template in the input image, and the deformed mesh is the reconstructed geometry of the object, which automatically ensures the mesh correspondence between different patients. (3) Also based on the feature extraction network, we design a shape error estimation network to estimate the error of a reconstructed geometry on the point level, i.e., evaluating the reliability of the output from the geometry reconstruction network. We demonstrate our approach on the application of lumbar spine geometry reconstruction from mid-sagittal MR images, and compare our geometry reconstruction network with UNet++, UTNet, Swin UnetTR, SLT-Net, and nnUNet, using the Dice metric. In this application, our geometry reconstruction network has much higher accuracy and artifact-free segmentation results, and our shape error estimation network facilitates quality control for clinical use.

2. METHODS

2.1. Architecture

As illustrated in Fig.1, the goal of our application is to reconstruct the lumbar spine geometries from 2D mid-sagittal MR images and then measure the medical parameters and estimate the error of any geometry. For this purpose, we designed a novel attention-based network for geometry reconstruction with error estimation. We developed the shape feature extraction network to capture consistent features from both image and shape input. Then, based on the extracted features, we designed the geometry reconstruction network to deform the input template and get the reconstructed geometry.

The architecture of the shape feature extraction network is depicted in Fig. 2. Our network applies two groups of CNN layers as pixel/region-level feature extractors at two spatial resolutions (512×512 and 128×128). The network utilizes self-attention on non-overlapping patches from the low-resolution CNN feature map, producing a sequence of image features with aggregated global context. Our approach uses a normalized mesh grid for image position embedding, on which each patch is centered on one of the mesh grid points. Then, shape self-attention and cross-attention between images and shapes are performed to cross-examine and fuse information from individual patches of the feature maps and points of the shape.

Figure 2.

The architecture of out feature extraction network.

The input shape (i.e., a template shape) is used to sample the low-resolution CNN feature map, where each point on the shape is assigned with specific image features by bilinear interpolation. This ensures that the shape is represented not only by a set of spatial points but also by their embeddings/features. The cross-attention acts as a bridge between image and shape. Whereas the sampling method can filter irrelevant information during the initialization, this method is confined to its perception field, the area of neighbor pixels. Our network utilizes cross-attention where each shape token (i.e., point with features) serves as a query vector to align itself with the entire image tokens (patches), thereby fostering the long-range dependencies. This approach allows our network to gather proper information from relevant image patches, even if the shape is initially far from the true shape. The final output of the feature extraction network is a shape feature matrix that contains information about the shape and image.

The structure of our geometry reconstruction network is shown in Fig.3. Given the input image and template shape, the output of the feature extraction network is used to predict the displacement vectors at the points of the template shape, and an intermediate shape is obtained after the first deformation. Then, this intermediate shape is further deformed by using the higher-resolution CNN feature. The instance segmentation masks are directly obtained from the reconstructed shape. We note that shape, mesh, and geometry are used interchangeably in this paper.

Figure 3.

The architecture of our geometry reconstruction network that deforms a mesh/shape template.

2.2. Attention with Relative Position Embedding

The attention mechanisms in ViT¹⁶ and its successors assume that tokens are located on a regular grid and therefore could not be used in our networks, in which points of a shape may be located at non-integer locations. To enable the new attention modules in our networks, we developed new equations for attention with relative position embedding. Let $X$ and $Y$ represent two matrices containing all input tokens. In our application, $X$ and $Y$ can be the same (i.e., $X=Y$), representing image or shape tokens, which are used for image or shape self-attention. Also, $X$ and $Y$ can be different: $X$ representing image tokens and $Y$ representing shape tokens, which is used for cross-attention in our application. The attention score matrix $A$ can be redefined as:

$$A=\text{softmax}\left(\frac{{\varphi }_Q{\left(X{W}_Q,R_1\right)}^T{\varphi }_K\left(Y{W}_K,R_2\right)}{\sqrt{d}}\right)$$ (1)

Here $W_Q,W_K$ are individual linear projections to generate query and key embeddings that can be replaced by MLPs in the implementation. $R_1$ and $R_2$ are two matrices encoding relative positions between tokens. The function ${\varphi }_Q,{\varphi }_K$ combines the original query and key matrix with the corresponding position matrix into a new query and key matrix using Fourier characteristics.¹⁷

To better explain the above Equation (1), we redefine query and key by incorporating sinusoidal functions to effectively extract relative position information. After the calculation, we derive that each entry, $a_{\mathit{\text{ij}}}$, of the attention score matrix, $A$, is proportional to the scalar dot product of query and key vectors $q_i,k_j$:

$$a_{\mathit{\text{ij}}}\propto \mathit{\text{dot}}\left(q_i,k_j\right)=\text{sum}\left(x_i{W}_Q\odot y_j{W}_K\odot \text{cos}\left(\left(p_i-p_j\right)W_1+b_1\right)+x_i{W}_Q\odot \text{cos}\left(\left(p_i-p_j\right)W_2+b_2\right)\right)$$ (2)

Here, each entry $a_{\mathit{\text{ij}}}$ of the attention score matrix is a function of the relative position vector $p_i-pj$. More details are shown in our previous work.¹⁵ We propose this novel relative position embedding to improve the performance of our geometry reconstruction method compared to absolute position embedding. We will show an ablation study in our future work.

2.3. Shape Error Estimation Network

The structure of the shape error estimation network is shown in Fig.4. It takes an image and an arbitrary shape as the inputs, and estimates the point-level difference between the input shape and the ground truth shape of the object in the image. This predicted error depicts the reliability of each point in the input shape that may be the output of a network for geometry reconstruction. During model training, the input shape is randomly selected from the training dataset and then added with noise, and the MSE loss is used.

Figure 4.

The architecture of our proposed network for shape error estimation.

3. EXPERIMENTS

3.1. Evaluation of the Geometry Reconstruction Network

In our experiment, we used the public dataset SSMSpine,¹⁵ which has 7000 images for training, 250 images for validation, and 2500 images for testing. To facilitate the comparison with image segmentation models, we converted each output mesh of our geometry reconstruction network into a segmentation mask and calculated the Dice Score. We used Pytorch 2.0.0¹⁸ for model implementation. The training process for each model was completed within 24 hours on a single NVIDIA RTX A6000.

We conduct qualitative and quantitative evaluations on the accuracy and quality of the predicted segmentation of six vertebrae and five intervertebral discs. We utilize Dice Score (a.k.a. Dice Similarity Coefficient, DSC, ranging from 0 to 1, the higher the better) to measure the spatial overlap of predicted and ground-truth segmentation masks numerically.

Tab.1 and Fig.5 show the performance of our network and the other five well-known models. Our proposed model demonstrates superior performance in all the listed spinal areas.

Table 1.

Dice Score Coefficient percentage score (%) comparison between our model and baselines. ”L1” ”L5” represent vertebrae bodies, ”D1” ”D5” represent disks, and “S1” denotes the first sacral vertebra. Red color highlights the best score among the methods.

Network	Ours	nnUNet	SLT-Net	Swin UnetTR	UTNet	UNet++
L1	97.8±1.0	95.1±4.0	91.9±11.2	93.1±8.4	82.7±28.3	65.1±40.8
L2	97.9±0.9	97.1±1.7	96.7±2.4	96.3±3.3	80.8±31.0	69.5±36.4
L3	96.9±2.4	94.7±10.3	93.9±12.7	93.3±14.5	84.1±27.0	70.0±31.7
L4	97.2±1.1	95.2±6.4	92.5±18.4	95.3±5.1	85.7±27.9	66.9±37.2
L5	97.5±1.0	94.9±6.0	94.6±8.9	94.7±4.9	91.1±16.2	81.4±28.0
S1	96.8±1.7	93.8±5.9	92.4±8.0	93.1±5.5	91.3±9.0	90.5±13.2
D1	96.3±1.7	94.0±3.2	92.7±7.2	93.1±5.4	76.8±31.0	64.6±40.3
D2	96.2±2.4	95.0±3.5	93.8±4.9	94.2±4.2	80.6±28.2	72.0±31.8
D3	95.1±3.2	92.9±7.8	89.9±18.7	92.0±9.2	83.9±27.0	63.1±38.2
D4	95.0±2.4	91.2±9.5	90.4±11.8	90.4±6.8	85.4±20.6	72.9±31.2
D5	96.0±1.9	92.4±6.4	91.1±11.4	91.3±6.2	90.7±11.8	87.2±20.3
Average	96.6±1.2	94.2±6.6	92.7±9.1	93.3±5.0	84.8±19.8	73.0±23.3

Figure 5.

Examples of semantic segmentation using different methods, Objects (vertebra or disc) are shown in different colors as indicated by the right-side colormap: (a) an input image (b) ground truth (c) our network (d) UNet++ (e) UTNet (f) Swin UnetTR (g) SLT-Net (h) nnUNet

3.2. Evaluation of the Shape Error Estimation Network

We used the shape error estimation network to estimate the output errors of the geometry reconstruction network on the test set of the SSMSpine dataset. We computed different correlation coefficients between the estimated errors and the true errors. The Pearson correlation coefficient is 0.9254, and the Spearman’s rank correlation coefficient is 0.8989, which shows that our shape error estimation network has high performance on error estimation.

3.3. Medical Parameter Measurement and Quality Control

All of the current machine learning (ML) models for segmentation and geometry reconstruction, including our models, are purely data-driven and therefore have no guarantee in output accuracy. Thus, the output shape from an ML model has to be checked by a human operator for clinical use. To reduce human negligence, those shapes with larger (estimated) errors will be checked first when the operator is vigilant, and those with smaller (estimated) errors will be checked later when the operator possibly becomes fatigued. Motivated by this, we conducted the experiment on the test set of the SSMSpine dataset to show how the shape error estimation network can facilitate quality control for medical parameter measurement.

In our experiment, we focused on the diameters and heights of the vertebrae and the disks of the lumbar spine. As illustrated in Fig.6, the vertebrae are measured by diameters and heights, including UVD, LVD, MVD, AVH, PVH, and MVH, while disks are measured by height only, including ADH, MDH, and PDH.

Figure 6.

Definitions of the medical parameters on an anatomical shape atlas. Corner vertices ($d_1,d_2,d_3$, and $d_4$), laying at the intersections between a vertebra and disc, are identified as landmarks. Midpoints ($d_5,d_6,d_7$, and $d_8$) on edges are also identified. For vertebrae, we introduce three main parameters, including upper, lower and middle vertebra diameters (UVD, LVD, MVD). Additionally, we measure the height of anterior, posterior and middle vertebra (AVH, PVH, MVH) to comprehensively depict the vertebra in cause of its abrasion. Similar to vertebrae, we introduce three main parameters including anterior, posterior and middle disc heights (ADH, PHD, MDH) to measure disc dimensions.

Tab. 2 shows the relative error of various medical parameters of patients in the test set.

Table 2.

The relative error’s comparison of different metrics at various $\alpha \%$ levels. UVD, LVD, MVD, AVH, PVH, MVH, ADH, PDH, and MDH are represented across different $\alpha \%$ levels.

$\mathit{\alpha }\mathbf{\%}$	UVD	LVD	MVD	AVH	PVH	MVH	ADH	PDH	MDH
100	1.75	2.23	2.44	2.10	2.24	2.06	4.48	7.23	4.73
80	1.54	2.04	2.20	1.90	2.06	1.73	3.91	6.57	4.21
60	1.49	1.95	2.02	1.80	1.98	1.59	3.68	6.42	4.05
40	1.47	1.87	1.89	1.70	1.95	1.46	3.41	6.45	3.94
20	1.45	1.77	1.76	1.56	1.89	1.33	3.10	6.43	3.67

The reconstructed shapes in the test set are ranked by the errors predicted by the shape error estimation network. These medical parameters can be evaluated using $\alpha \%$ of the shapes with the lowest (estimated) errors. Here we take $\alpha$ as 20, 40, 60, 80, and 100. To enhance clarity, we have converted the relative error to the percentage error (%) in Table 2. It is evident that the relative error of the medical parameters decreases when samples with higher (estimated) errors are filtered out.

4. CONCLUSION

In this work, we propose an innovative neural network architecture for geometry reconstruction from medical images, which utilizes image self-attention and shape self-attention as well as cross-attention between image and shape, to capture consistent features. In addition, we propose a shape error estimation network, sharing the same feature extraction network with the geometry reconstruction network. The experiment on the SSMSpine dataset of lumbar spine MR images demonstrates that our geometry reconstruction network has higher accuracy and is free of artifacts, and our shape error estimation network facilitates quality control for clinical use. Our approach could be extended to the geometry reconstruction of other human organs, which is our ongoing work.