Transformer-Based Integrated Framework for Joint Reconstruction and Segmentation in Accelerated Knee MRI

Abstract

Magnetic Resonance Imaging (MRI) reconstruction and segmentation are crucial for medical diagnostics and treatment planning. Despite advances, achieving high performance in both tasks remains challenging, especially in the context of accelerated MRI acquisition. Motivated by this challenge, the objective of this study is to develop an integrated approach for MRI image reconstruction and segmentation specifically tailored for accelerated acquisition scenarios. The proposed method unifies these tasks by incorporating segmentation feedback into an iterative reconstruction algorithm and using a transformer-based encoder–decoder architecture. This architecture consists of a shared encoder and task-specific decoders, and employs a feature distillation process between the decoders. The proposed model is evaluated on the Stanford Knee MRI with Multi-Task Evaluation (SKM-TEA) dataset against established methods such as SegNetMRI and IDSLR-Seg. The results show improvements in the PSNR, SSIM, Dice, and Hausdorff distance metrics. An ablation study confirms the contribution of feature distillation and segmentation feedback to the performance gains. The advancements demonstrated in this study have the potential to impact clinical practice by facilitating more accurate diagnosis and better-informed treatment plans.

1. Introduction

The superior soft tissue contrast of magnetic resonance imaging (MRI) makes it an invaluable diagnostic instrument across a wide variety of diseases [1]. However, prolonged acquisition times can lead to patient discomfort, reduced throughput, and introduction of motion artifacts. Therefore, reducing the duration of MRI scans has become a pressing area of research [2].

One promising approach to achieve rapid MRI acquisition is the use of compressed sensing (CS) techniques [3]. These techniques violate the Nyquist–Shannon sampling theorem by undersampling and collecting fewer measurements than is conventionally required to reconstruct diagnostic-quality images. Despite their efficiency, they contradict the standard sampling theory, leading to aliasing artifacts. To mitigate this, researchers have incorporated additional a priori knowledge [4]. Recently, machine learning (ML) techniques have been integrated into the image reconstruction process [2]. The application of ML has led to the development of algorithms that can reconstruct high-quality images from sparsely sampled MRI data, significantly accelerating MRI scans and reducing acquisition time [5,6,7]. Reconstructed images often serve as a means to derive clinically relevant parameters through postprocessing steps such as segmentation and tissue characterization [8]. ML tools excel in this task, even automating dense image labeling tasks to match expert variability [6,9,10].

While many machine learning-based segmentation algorithms operate under the assumption of receiving a “clean” image, they do not necessarily account for the challenges posed by undersampled MRI scenarios [11]. Though traditional approaches to medical imaging often treat acquisition, reconstruction, and segmentation as distinct phases, it is widely understood within the research community that these stages are interrelated and can significantly impact one another. Nonetheless, the emphasis has often been on each individual stage, sometimes overlooking the cumulative effects [8]. Specifically, when prioritizing faster imaging speeds without proper image restoration, there is a risk of introducing residual aliasing and blurring, which can subsequently lead to errors in segmentation [12]. Figure 1 illustrates how a segmentation network trained on fully sampled (clean) data can yield misleading segmentation masks for undersampled (noisy) data. Thus, there is a need for efficient automated approaches that can simultaneously reconstruct MRI data and accurately segment the region of interest (ROI) [13]. However, the lack of segmentation datasets with k-space data required for MRI reconstruction poses a significant challenge in the development of joint reconstruction–segmentation algorithms [12].

Despite limited public datasets, several studies have addressed the joint MRI reconstruction and segmentation problem. SegMRI [8] uses CS and Gaussian mixture model segmentation on patch-based dictionaries for sparse image representation, thereby enhancing image reconstruction and segmentation. SegNetMRI [11] uses an iterative framework involving a data fidelity unit and UNet [14]-based denoising and segmentation networks. Each iteration shares an encoder between segmentation and denoising with unique decoders. The segmentation decoder is reused across iterations, while denoisers use different encoders–decoders across iterations. Multiple segmentation results are merged using a $1\times 1$ convolution at the end of iterations. FR-Net [13] presents a deep learning approach that includes a reconstruction network derived by unrolling the Fast Iterative Shrinkage–Thresholding Algorithm (FISTA) [15]. This is followed by a segmentation network that operates independently without sharing parameters with the reconstruction component. Lastly, IDSLR-SEG [12] introduces a framework for joint calibrationless Parallel MRI (PMRI) reconstruction and segmentation based on unrolling an iterative re-weighted least squares algorithm to minimize a CLEAR cost function [16] for calibrationless PMRI reconstruction [17]. The denoising network shared across all iterations and the segmentation network both use a shared encoder and are trained end-to-end using a few-shot learning strategy.

In summary, while existing models offer a foundational approach through end-to-end joint training and a shared encoder for dual tasks, they do not completely harness the latest progress in multitask learning. The method introduced in this study aims to enhance this by creating a more synergistic framework for simultaneous knee MRI reconstruction and segmentation. In this paper, we introduce an innovative approach that takes undersampled k-space data as input and outputs both reconstructed knee MRI images and ROI segmentation.

Leveraging recent multitask learning developments, a framework is proposed for joint segmentation and reconstruction. The primary contributions of this paper are shown below:

• Segmentation-Integrated Unrolled Reconstruction: we proposes a unique cost function for unrolling the reconstruction algorithm that integrates segmentation results into the reconstruction process.
• Enhanced Encoder–Decoder Architecture: this paper employs a Swin Transformer [18,19]-based encoder–decoder architecture for multitask denoising and segmentation. A shared attention mechanism is implemented wherein Query and Key vectors in the self-attention module of the task-specific decoder are computed using the shared encoder’s output.
• Feature Distillation in Multitask Decoders: the proposed model introduces integration of features between decoders through a distillation process by applying spatial attention features from each task that are then incorporated into the other task’s decoder.

2. Materials and Methods

2.1. Background

2.1.1. Compressed Sensing in MRI

Reconstructing MRI images can be formulated as an optimization problem [20,21], as presented in (1):

$$\begin{array}{cc}\hfill \widehat{x}& ={arg\; min}_{x}f\left(x\right)+\beta \mathrm{R}\left(x\right),\hfill \end{array}$$

(1)

where x represents the reconstructed images, $f\left(x\right)=\frac{1}{2}{\parallel Ax-y\parallel }_2^2$ is the data fidelity term, y is the data acquired in the k-space domain, and A stands for the imaging model, which includes coil sensitivity profile maps, a Fourier transformation operation, and data subsampling. The regularization function $\mathrm{R}\left(x\right)$ and its corresponding parameter $\beta$ constrain the problem when datasets are highly subsampled [21], preventing ill-posed conditions where multiple solutions might fulfill (1).

One way to solve (1) is to use a two-step iterative process alternating between gradient descent and a proximal operation, as provided in Equations (2) and (3) [22]:

$$\begin{array}{cc}\hfill x^{\left(n^{\prime }\right)}& =x^{\left(n\right)}-t\nabla f\left(x^{\left(n\right)}\right)\hfill \end{array}$$

(2)

$$\begin{array}{cc}\hfill x^{(n+1)}& ={\mathbf{prox}}_{\beta \mathrm{R}}\left(x^{\left(n^{\prime }\right)}\right).\hfill \end{array}$$

(3)

In these equations, n represents the n-th iteration, t is a scalar that indicates the gradient’s step size, and ${\mathbf{prox}}_{\beta \mathrm{R}}(·)$ denotes the proximal operator of $\mathrm{R}$.

A promising approach to developing efficient reconstruction algorithms is to use a data-driven method to learn optimal trainable parameters in regularization functions. Here, the proximal step is often replaced with a deep neural network, which directly learns a parameterized form of the regularization function [23,24]. Consequently, the proximal update in (3) is redefined [21,25] as follows:

$$\begin{array}{c}\hfill x^{(n+1)}=N_{\theta }\left(x^{\left(n^{\prime }\right)}\right),\end{array}$$

(4)

where $N_{\theta }$ is the neural network and $\theta$ denotes learnable parameters, which can either be shared or distinct across iterations. The iterative process provided by (2) and (4) is “unrolled” into a model $U_{\theta }$, which is then trained by minimizing the loss function ${min}_{\theta }{\sum }_{i}L(U_{\theta }(y_i,A_i),x_i)$ [21,25], where $L(·)$ measures the distance between its inputs and $x_i$ represents the ith ground-truth example.

2.1.2. Transformers in Medical Imaging

Integration of attention mechanisms [26] into architectures influenced by Convolutional Neural Networks (CNNs) [27] has been a significant focus within the computer vision community. Consequently, it has led to the development of “Vision Transformer” (ViT) models [28]. Their popularity has grown due to their capacity to encode long-range dependencies and generate effective feature representations [29].

ViT models have demonstrated considerable potential in MRI restoration and analysis tasks. In the domain of MRI restoration, notable work includes that by Feng et al. [30], who have developed a cross-attention module capable of extracting and merging complementary features from auxiliary imaging modalities [29]. In the field of MRI analysis, a standout example is Swin UNETR [31], which has achieved leading performance in the Brain Tumor Segmentation (BraTS) 2021 challenge [32]. This model combines a Swin Transformer encoder with a CNN-based decoder. The Swin Transformer using a patch partition layer can create non-overlapping patches from the input data and construct windows for self-attention computations. These processed feature representations are then forwarded to a CNN decoder via skip connections at multiple resolutions [29].

2.1.3. Multi-Task Learning for Dense Predictions

Multi-Task Learning (MTL) [33] seeks to build generalized ML models able to generate all pertinent task outputs from a given input [34]. MTL can improve generalization capability with shared representation learning from multiple task-specific training signals [35]. MTL offers advantages over single-task learning such as increased inference speeds by avoiding repetitive feature calculations in shared layers [36], with potential performance improvements when tasks share information or are able to regularize each other [37].

Significant MTL work in pixel-level prediction tasks has led to innovations in network architecture and optimization techniques [34]. Optimization methods can maintain balance among tasks during training to prevent any single task’s dominance. For instance, Kendall et al. [38] have quantified homoskedastic uncertainty to balance single-task losses, and GradNorm [39] is able to equalize task-specific gradients. On the architectural front, methods are able to leverage shared information among tasks. Approaches include hard parameter sharing [40,41] (a shared encoder branching into task-specific heads), soft parameter sharing [42,43] (individual task parameters with cross-task feature sharing), and designs that first predict tasks and then leverage these predictions to enhance task outputs [44]. For example, Xu et al. [45] have developed a multimodal distillation module that can distill information from initial predictions of other tasks using spatial attention and then incorporate it into the task of interest in order to effectively utilize intermediate predictions’ complementary information.

2.2. Proposed Method

2.2.1. Incorporating Segmentation Feedback into the Reconstruction Cost Function

Improvement of reconstruction can inherently lead to enhanced segmentation; thus, in most joint MRI reconstruction and segmentation research, segmentation is typically executed subsequent to the reconstruction process. This approach has been employed in previous studies such as the study of Huang et al. [13], where the segmentation network was only employed after the reconstruction process. However, segmentation outcomes can reciprocally refine the reconstruction process. As such, a novel cost function is proposed in this study to facilitate an optimization algorithm that enables integration of segmentation results into the reconstruction process. This is achieved by appending a term $\mathrm{S}(x;\left\{w_k\right\})$ to (1):

$$\begin{array}{cc}\hfill \widehat{x}& ={arg\; min}_{x}f\left(x\right)+\beta \mathrm{R}\left(x\right)+\mu \mathrm{S}(x;\left\{w_k\right\}).\hfill \end{array}$$

(5)

Here, $\mu$ adjusts the impact of the added term, the set $\left\{w_k\right\}$ constitutes a transformed segmentation mask, and $\mathrm{S}(x;\left\{w_k\right\})$ is a term indicating the relationship between the reconstructed image and the transformed segmentation mask.

To comprehend the role of $\mathrm{S}\left(x\right)$, consider $\mathrm{S}(x;\left\{{\check{w}}_k\right\})=\frac{1}{2}{\sum }_{k=1}^K{\parallel C_{k}x\parallel }_{{\check{W}}_k}^2=\frac{1}{2}{\sum }_{k=1}^K{\left(C_{k}x\right)}^T$${\check{W}}_k\left(C_{k}x\right)$, where ${\check{W}}_k=diag\left\{{\check{w}}_k\right\}$, ${\check{w}}_k$ represents an indicator image for the ROI boundary extracted from the segmentation mask, and $C_k$ denotes a finite differencing matrix in the $x,y,$ or z directions; thus, in this case, $K=3$. Each image update involves the gradient of $\mathrm{S}\left(x\right)$, defined as $\nabla \mathrm{S}\left(x\right)={\sum }_{k=1}^K{C}_k^T{\check{W}}_k{C}_{k}x$, and is zeroed where ${\check{w}}_k$ is zero, indicating the boundary region. Consequently, the update considering $\nabla \mathrm{S}\left(x\right)$ encourages spatial smoothness outside the boundary region while limiting smoothing across boundaries.

Building on this understanding, a modification to (2) with $\mathrm{S}(x^{\left(n\right)};\left\{w_k^{\left(n\right)}\right\})=\frac{1}{2}{\sum }_{k=1}^K{\parallel c_k\ast x^{\left(n\right)}\parallel }_{W_k^{\left(n\right)}}^2$ can be proposed, as shown below:

$$\begin{array}{cc}\hfill x^{\left(n^{\prime }\right)}& =x^{\left(n\right)}-t\nabla f\left(x^{\left(n\right)}\right)-t^{\prime }\nabla \mathrm{S}(x^{\left(n\right)};\left\{w_k^{\left(n\right)}\right\}),\hfill \end{array}$$

(6)

where $t^{\prime }$ sets the step size for the added term and

$$\begin{array}{cc}\hfill \nabla \mathrm{S}(x^{\left(n\right)};\left\{w_k^{\left(n\right)}\right\})& =\sum _{k=1}^K{\tilde{c}}_k\ast \left(W_k^{\left(n\right)}\left(c_k\ast x^{\left(n\right)}\right)\right).\hfill \end{array}$$

(7)

Here, $W_k^{\left(n\right)}=diag\left\{w_k^{\left(n\right)}\right\}$, $w_k^{\left(n\right)}={\sum }_{k^{\prime }=1}^{K^{\prime }}g(c_k^{\prime }\ast m^{\left(n\right)})$, $g(·)$ is an activation function for nonlinearity, ${\tilde{c}}_k$ is a flipped convolution kernel of $c_k$, and $m^{\left(n\right)}$ designates the ROI mask obtained from the segmentation network at the $n\mathrm{th}$ iteration. Both $\left\{c_k\right\}$ and $\left\{c_k^{\prime }\right\}$ denote K and $K^{\prime }$ sets of convolutional filters trained end-to-end alongside the denoising and segmentation network, thereby allowing the data to guide how the model utilizes the segmentation result for both tasks. Although $\mathrm{S}(x^{\left(n\right)};\left\{w_k^{\left(n\right)}\right\})$ is not considered as a part of the data consistency term, it is differentiable its gradient can be easily found. Therefore, the update for the $\mathrm{S}(x^{\left(n\right)};\left\{w_k^{\left(n\right)}\right\})$ term is included in (2) rather than in the proximal step. This update is denoted as the modified data consistency step in Figure 2a.

2.2.2. Swin Transformer-Based Encoder–Decoder Approach

The proposed method leverages a Swin Transformer-based encoder–decoder architecture [46] for multitask [47] denoising and segmentation, as depicted in Figure 2b. The encoder, inspired by ResNet [48], includes a patch embedding process that adjusts the spatial resolution and channel dimension, followed by several Transformer blocks. This configuration produces a feature pyramid compatible with many vision task architectures. The Swin Transformer’s design [18,19] of alternating window partitioning and shifted window partitioning is then applied.

The decoder module [46,47], influenced by CNN-based decoders, replaces convolutional layers with Swin Transformer blocks. The four-stage decoder comprises two Swin Transformer blocks per stage, with an upsampling layer in between to double the spatial resolution and halve the channel dimension. Here, a “stage” refers to a set of multiple Transformer blocks that process data at the same spatial resolution. In contrast to [47], the encoder features are fed to the decoder via a skip connection at the same resolution.

A shared attention mechanism [47] is employed as well. In standard Vision Transformer self-attention, each multi-head self-attention layer independently creates its Query, Key, and Value vectors using only its own input. These vectors are used to calculate an attention score, with the Query and Key vectors determining the score and the Value vector generating a weighted sum to form the self-attention output. In contrast, the shared attention mechanism modifies this process within the last Transformer block of the decoder at each stage. Specifically, it computes Query and Key vectors from the shared encoder’s output that corresponds to the same spatial resolution, while the Value vector is derived from the preceding decoder stage for the specific task, ensuring task-specific outputs. This configuration mirrors the decoder in the original Transformer design [26]. By incorporating shared encoded features into the computation of the Query and Key vectors, this shared attention mechanism enhances each task’s ability to utilize cross-task relationships and dependencies.

2.2.3. Feature Sharing and Distillation across Multitask Decoders

The proposed model, as outlined in Section 2.2.2, utilizes an encoder–decoder framework with a shared encoder and two distinct decoders for MRI reconstruction and segmentation tasks. This design with separate paths for each task can facilitate task-specific feature processing.

A key aspect of the proposed architecture involves sharing and integration of features between decoders. This design draws inspiration from a previous study [45]. However, unlike the previous study, the proposed approach applies a distillation process to intermediate features of decoders rather than applying it to initial predictions. Before expanding the spatial dimension, features from each task are subjected to a spatial attention process and then incorporated into the other task’s decoder:

$$\begin{array}{cc}\hfill I_r^{i+1}& =O_r^i+\sigma \left(W_{s,r}{O}_s^i\right)\odot O_s^i\hfill \end{array}$$

(8)

$$\begin{array}{cc}\hfill I_s^{i+1}& =O_s^i+\sigma \left(W_{r,s}{O}_r^i\right)\odot O_r^i.\hfill \end{array}$$

(9)

Here, $I_r^{i+1}$ and $I_s^{i+1}$ denote inputs to the $(i+1)$-th stage of the reconstruction and segmentation decoders. Each stage contains multiple transformation blocks processing vectors of identical dimensions, $O_r^i$ and $O_s^i$ symbolize outputs from the i-th stage of respective decoders, $\sigma (·)$ denotes the sigmoid function, and $W_{s,r}$ and $W_{r,s}$ represent tunable parameters.

The term $\sigma \left(W_{s,r}{O}_s^i\right)$ presents a spatial mask applied to the segmentation decoder feature for the reconstruction decoder. This spatial attention mechanism enables the model to highlight critical spatial locations within feature maps through gating, thereby controlling the information flow between decoders. Following this refinement, these features are blended into the other task’s decoder before expanding the spatial dimensions. This approach takes advantage of mutual benefits between reconstruction and segmentation tasks, potentially augmenting the overall effectiveness of the network.

3. Results

3.1. Dataset Details

3.1.1. SKM-TEA Dataset

The Stanford Knee MRI with Multi-Task Evaluation (SKM-TEA) dataset [6] offers a substantial pool of quantitative knee MRI (qMRI) scans, enabling evaluation of MRI reconstruction and analysis methods. The dataset comprises around 25,000 slices from 155 patients, including raw-data measurements, scanner-generated DICOM images, manual segmentation of four tissues (Patellar Cartilage, Femoral Cartilage, Tibial Cartilage, and Meniscus), and annotations for sixteen pathologies. We employed this dataset to benchmark the comparing methods and the proposed method.

In the provided dataset, an inverse Fourier transform is applied to the fully-sampled k-space in the readout direction to generate a hybridized k-space ($x\times k_y\times k_z$). Sensitivity maps for each 2D axial slice were estimated using JSENSE [49] and the fully-sampled k-space was reconstructed using SENSE [50], which then served as the target image for the reconstruction task. Dataset acquisition used double-echo steady-state (qDESS) MRI method, which provides two sets of 3D images (termed echoes—E1 and E2). In this study, we utilized only E1 data from the two available echoes. The majority of the data consisted of 8-channel coil. Data with 16-channel coil (8 out of 155 samples) were omitted due to GPU memory constraints. This study adhered to the training, validation, and test splits provided within the dataset.

3.1.2. Data Preprocessing

The reconstruction baseline method provided by authors of the SKM-TEA dataset used 2D k-space data in the axial direction ($k_y\times k_z$). However, through empirical findings, it was observed that training a segmentation network with axial slices posed a more significant challenge compared to the sagittal direction due to the comparatively sparse distribution of various tissue classes (such as tibial cartilage and meniscus). In light of these findings, 2D k-space data in the sagittal direction were used, which is consistent with the baseline segmentation methods of the SKM-TEA dataset. For this, ($k_x\times k_y\times k_z$) k-space data were produced by applying the Fourier transform in the readout direction, then an inverse Fourier transform along the z axis, resulting in a hybridized k-space ($k_x\times k_y\times z$). For undersampling, a 2D Poisson disc at an acceleration factors of 8 was utilized with the code provided by the authors of the SKM-TEA dataset. The undersampling mask was generated for the true acquisition region ($512\times 416$), then zero-padded to match with the kspace data size ($512\times 512$). During training, 10,000 precomputed undersampling masks were cached to ensure consistency across different training sessions. A fixed undersampling mask was generated for each scan in the test dataset for evaluation.

3.2. Baseline and Comparative Methods

As baseline, a 2D UNet trained for joint reconstruction and segmentation was employed utilizing an image-to-image approach. UNet features a shared encoder and two distinct decoders, with each decoder being dedicated to either reconstruction or segmentation. Moreover, the proposed method was compared with previously suggested methods for joint MRI reconstruction and segmentation, including SegNetMRI [11] and IDSLR-SEG [12]. These methods alternate between data consistency and denoising via neural networks, bearing close resemblance to the unrolled compressed sensing in (2)–(4). Considering the dataset’s distinct challenge of multicoil Knee MRI reconstruction and segmentation, these methods were adapted and reimplemented, in the course of which reimplementations were based on unrolled compressed sensing, utilizing U-Net as the denoiser and the segmentation network as detailed in previous studies [11,12]. For IDSLR-SEG, the method was adapted into a calibrated approach considering the provision of sensitivity maps with the dataset. The main divergences between implementations of SegNetMRI and IDSLR-SEG, lie in whether to use a shared or distinct denoising encoder across iterations and whether segmentation occurs multiple times during the iterative process. The distinctions between the models are summarized in

Table 1. Comparative analysis of different models (DSNA: Denoising and Segmentation Network Architecture; SEDSN: Shared Encoder between Denoising and Segmentation Networks; SDI: Shared Denoiser across Iterations; MSPI: Multiple Segmentation Predictions across Iterations).

Model	DSNA	SEDSN	SDI	MSPI
SegNetMRI	U-Net	Yes	No	Yes
IDSLR-Seg	U-Net	Yes	Yes	No
Proposed	Transformer	Yes	Yes	Yes

. A visual comparison between SegNetMRI and IDSLR-SEG has been provided previously in [12].

3.3. Details on Implementation and Training

The following training specifications were applied to all methods mentioned in this section. The Pytorch [51] deep learning library was leveraged for training. The training objective combined a complex ℓ1 loss for the reconstruction task and a soft Dice loss for the segmentation task. Input k-space data were normalized using the same standard deviation value of the target 3D volume as in the reconstruction baseline method of SKM-TEA dataset. Training was conducted over 200 epochs using the AdamW optimizer with a weight decay of 0.05. The learning rate was adjusted using a custom CyclicLR scheduler, with the parameters set as follows: maximum learning rate of 0.0005, gamma value of 0.5, and step size up of 15. Gradient accumulation steps were used to achieve an effective batch size of 16 across all methods. For a fair comparison, the number of trainable parameters for all methods was set at approximately 40 million. All unrolling-based methods employed a total of four unrolling iterations. The best epoch was determined using the sum of PSNR and the mean Dice score multiplied by 40, and was evaluated using the validation split of the SKM-TEA dataset for all methods presented in this section.

In the proposed method, both the encoder and each decoder featured two Transformer blocks at each stage. The window size was set as 8 and the number of heads in each stage was set as 3, 6, 12, and 24, respectively. K and $K^{\prime }$ in (7) were set as 16. To enhance the training stability, two approaches were employed in the implementation of the method described in Section 2.2.1: First, the terms $c_k\ast x^{\left(n\right)}$ and $c_k^{\prime }\ast m^{\left(n\right)}$ in (7) were calculated using a sequence of layers composed of a convolutional layer, followed by an instance normalization layer, and a PReLU activation function. Second, the multiplication involving $W_k^{\left(n\right)}$ in (7) was parameterized by concatenating $w_k^{\left(n\right)}$ and $c_k\ast x^{\left(n\right)}$, which was then followed by a similar sequence of layers as in the previous approach. For increased model flexibility, separate convolutional filters were used instead of employing the flipped version of $c_k$, denoted as ${\tilde{c}}_k$ in (7).

3.4. Results: Quantitative and Qualitative Evaluation

The test split of the SKM-TEA dataset (34 samples, excluding two 16-channel coil samples) was used for performance assessments. The reconstruction quality was quantitatively evaluated using the peak signal-to-noise ratio (PSNR) and structural similarity index (SSIM [52]) metrics. The segmentation performance was measured using the Dice similarity coefficients and 95% Hausdorff distance for each tissue class. The PyTorch-based MEDDLR framework [53] was employed for evaluations. These metrics were computed based on 3D volumes, with the mean and standard deviation values reported from 34 volumes.

In the reconstruction task, the proposed model generally outperformed other methods in both PSNR and SSIM metrics, as highlighted in

Table 2. Comparison of different joint MRI reconstruction and segmentation methods using the peak signal-to-noise ratio (PSNR), structural similarity index measure (SSIM), Dice score, and 95% Hausdorff distance (mean (standard deviation)).

Method	Reconstruction		Segmentation
	PSNR (dB)	SSIM	DICE	Hausdorff (mm)
UNet	33.153 (0.977)	0.765 (0.027)	0.824 (0.053)	5.25 (3.27)
SegNetMRI	35.322 (1.025)	0.834 (0.022)	0.821 (0.051)	13.62 (17.37)
IDSLR-Seg	35.139 (1.002)	0.828 (0.022)	0.826 (0.056)	6.68 (8.42)
Proposed	35.550 (1.012)	0.834 (0.021)	0.825 (0.055)	4.63 (3.14)

. This numerical advantage was corroborated by the empirical observations. The images generated by the proposed model, as illustrated in Figure 3, displayed fewer or similar levels of errors when compared to those produced by alternative methods.

In the segmentation task, the proposed method yielded competitive Dice scores and demonstrated improvements in Hausdorff distances compared to other methods, as indicated by the data in Table 2. This performance is visually corroborated in Figure 3.

To sum up, the proposed model offers more accurate results in the reconstruction task and is on par with other models in terms of Dice scores in the segmentation task while improving the Hausdorff distances. The proposed architecture, which combines transformer-based encoders and decoders for an integrated approach to reconstruction and segmentation, contributed to these results.

4. Discussion

In this study, a novel approach has been introduced for integrated MRI image reconstruction and segmentation, challenging the traditional view of these tasks as separate entities. The proposed method synergistically combines these tasks, utilizing a combination of distillation applied to features between decoders and a reconstruction cost function guided by segmentation feedback. The proposed approach demonstrates advantages over existing techniques in terms of achieving improved image reconstruction and segmentation quality.

4.1. Discussion of Reconstruction Results

SegNetMRI and IDSLR-SEG produced reconstructed images of comparable quality, whereas U-Net removed many image details and introduced a higher level of blur, consistent with the numerical results shown in Table 2. These results suggest that design differences between SegNetMRI and IDSLR-SEG, as summarized in Table 1, do not significantly affect the final reconstruction outcome. However, the utilization of iterative methods was found to be a crucial factor in enhancing reconstruction performance. While IDSLR-Seg and SegNetMRI performed well, they were slightly outperformed by the proposed method. The proposed method showed the highest performance in terms of PSNR, with a value of 35.550, and was tied with SegNetMRI in SSIM with a score of 0.834.

4.2. Discussion of Segmentation Results

IDSLR-Seg led in terms of the DICE coefficient with a score of 0.826, closely followed by the proposed method at 0.825 (0.055) and UNet at 0.824. The proposed method excelled at minimizing the Hausdorff distance, achieving the lowest score of 4.63. The segmentation of meniscus pixels proved more challenging than that of other tissues. The SegNetMRI and IDSLR-SEG methods struggled to accurately segment the meniscus in the specific slice depicted in Figure 3. Contrary to the reconstruction results, design variations between SegNetMRI and IDSLR-SEG influenced their segmentation performance. Specifically, IDSLR-SEG outperformed SegNetMRI in terms of the Hausdorff distance, suggesting that that the use of distinct denoisers across iterations or the merging of multiple segmentation results using a $1\times 1$ convolutional operation can adversely affect segmentation performance.

4.3. Ablation Study

To evaluate the influence of specific design elements on the performance of the proposed method in MRI reconstruction and segmentation tasks, an ablation study was performed. The analysis was mainly focused on two components: (i) feature distillation between task-specific decoders, and (ii) implementation of a reconstruction cost function guided by segmentation feedback.

To ascertain the individual contributions of feature distillation (FD) and segmentation feedback (SF), the proposed model was evaluated under three different conditions: without feature distillation and segmentation feedback, with only segmentation feedback, and with both. The results were then used to determine the effects of these features on the model’s performance compared to a baseline model devoid of these components.

Table 3. Ablation study results comparing the effectiveness of Feature Distillation (FD) and Segmentation Feedback (SF) on the peak signal-to-noise ratio (PSNR), structural similarity index measure (SSIM), Dice score, and 95% Hausdorff distance (mean (standard deviation)).

Proposed Method		Reconstruction		Segmentation
FD	SF	PSNR (dB)	SSIM	DICE	Hausdorff (mm)
No	No	35.524 (1.024)	0.836 (0.021)	0.826 (0.054)	4.70 (2.55)
No	Yes	35.554 (1.016)	0.840 (0.020)	0.819 (0.057)	5.23 (3.42)
Yes	Yes	35.550 (1.012)	0.834 (0.021)	0.825 (0.055)	4.63 (3.14)

presents the results of the ablation study. The data suggest that both FD and SF contributed to the enhanced performance of the proposed model. Specifically, When neither FD nor SF was applied, the PSNR was 35.524 dB, the SSIM was 0.836, the DICE was 0.826, and the Hausdorff distance was 4.702 mm. Implementing SF alone improved both the PSNR (35.554 dB) and SSIM (0.840), indicating enhanced reconstruction performance; however, it led to a slight degradation in the DICE (0.819) and an increase in the Hausdorff distance (5.225 mm), suggesting that the segmentation performance suffered. Incorporating both FD and SF yielded a PSNR of 35.550 dB and an SSIM of 0.834, maintaining strong reconstruction performance. Notably, the DICE score was almost identical to the baseline (0.825) and the Hausdorff distance improved to 4.63 mm, suggesting a balanced improvement across both reconstruction and segmentation tasks. In summary, while the inclusion of segmentation feedback (SF) did improve reconstruction performance as measured by PSNR and SSIM, it slightly compromised the segmentation performance in terms of DICE and increased the Hausdorff distance. Incorporating feature distillation (FD) mitigated this sacrifice in segmentation performance and resulted in a balanced overall performance improvement.

4.4. Limitations

This study has certain limitations tied to the dataset used in the experiment. The method’s training and validation were limited to a specific set of MRI images sourced from a single hospital and obtained using scanners from the same manufacturer. This might restrict the generalizability of this study’s findings. Future studies should aim to validate the proposed model with more diverse datasets incorporating images from a variety of MRI machines and diverse patient populations.

Although our ablation study indicated improved performance due to the Transformer-based network architecture, feature distillation, and segmentation feedback, the individual contributions of these techniques deserve further exploration. Additionally, the computational cost of the proposed model might pose challenges in time-sensitive clinical settings.

5. Conclusions

The proposed model offers threefold benefits: enhanced accuracy in MRI reconstruction, precise tissue segmentation, and substantial time savings during both MRI acquisition and post-acquisition image analysis. Employing an acceleration factor of eight for undersampling theoretically reduces the MRI scan time by a factor of eight. Concurrent segmentation capabilities further streamline the diagnostic process, potentially saving additional time in clinical workflows. The model’s proficiency in detail-rich reconstruction and precise tissue segmentation holds promise for earlier diagnosis and more timely treatments. These compelling results warrant further research to fully realize the model’s potential in clinical settings.

Future work should concentrate on optimizing the algorithm to reduce computational time without sacrificing performance. Furthermore, future studies should investigate the scalability of the model [54] and assess the relationship between the number of parameters and performance. Insight into this relationship could inform the development of more efficient model architectures. Future investigations should explore pretraining strategies that make use of extensive datasets. Studying the effects of large-scale pretraining on the model’s performance could reveal new avenues for improving MRI reconstruction and segmentation.