Weakly supervised segmentation of intracranial aneurysms using a novel 3D focal modulation UNet

While 3D UNet was a common choice for UIA detection and segmentation [3], the newer Vision Transformer (ViT) and its variants (e.g., Swin Transformer) that leverage self-attention mechanisms [11, 12] have become increasingly popular in medical imaging applications, providing superior performance. More recently, to better encode contextual information with a lighter model, Yang et al. [9] proposed focal modulation for DL in vision tasks, and demonstrated that it performed better than the Swin Transformer. In focal modulation networks, self-attention is replaced by a stack of depth-wise convolutional layers that focally encode visual contexts and selectively gather them into a modulator using a gated aggregation. Next, the modulator is injected into the query and passed to the next block. Both self-attention and focal modulation techniques involve linearly determining the key, query, and value based on input tokens. Nevertheless, the primary distinction lies in their operational sequence: in self-attention, a query-key interaction occurs before aggregation with the value, whereas in focal modulation, the value is first aggregated with contexts around each key, and subsequently, the result is adaptively modulated with the query. Despite excellent performance in classification and segmentation of natural images, focal modulation’s performance in 3D medical images is yet to be verified. Recently, Naderi et al. [13] proposed a 2D UNet-like architecture with focal modulation blocks as the encoder and decoder for abdominal CT segmentation. It offered a higher mean Dice score (not confirmed by statistical tests) than the Swin-Unet [14], another UNet-like model with Swin Transformers as the encoder and decoder. Different from [13], we followed the approach of the recent 3D Swin-UNETR [15], a popular Swin-CNN hybrid model to build a new architecture by replacing the Swin Transformer with 3D focal modulation for the encoder. At the encoder branch, each layer’s output, consisting of image tokens and their embeddings, is reshaped to a volumetric patch and passed through a residual block to work as skip connections. Starting from the bottleneck, each feature map is expanded using a transposed convolutional block and then concatenated with the corresponding skip connection before passing through another residual block to the next layer. The detailed architecture is depicted in Fig. 1.

As the proposed deep learning (DL) models (i.e., FocalSegNet) were trained on rough segmentations that overestimate the true aneurysm shapes, we used a fully connected CRF [16] to further refine our model’s initial prediction ($P(x_{i})$). If we define $X_{i}$ as a random variable representing the assigned label to pixel $i$ (foreground/background), and $I_{i}$ as a global observation characterizing the pixel information in the input image, the CRF is modeled as the pair of $(I,X)$ which follows a Gibbs distribution of the form $P(X=x|I)=\frac{1}{Z(I)}exp(-E(x|I))$ [17]. Here, $Z(I)$ is a normalization factor, and $E(x|I)$ is the energy function of assigning labels $x$ to the image. Through minimizing this energy function, the new labels will have a greater likelihood of being assigned to the image. Inspired by [18], we define the energy function as:

Here, $\psi_{u}(x_{i})$ is the unary term measuring the cost of assigning label $x_{i}$ to pixel $i$, and $\psi_{p}(x_{i},x_{j})$ is the pairwise term that shows the cost of assigning labels $x_{i}$ and $x_{j}$ to pixels $i$ and $j$ simultaneously. The unary term is derived from our FocalSegNet output logits, and the pairwise term is computed by applying a bilateral kernel and a Gaussian kernel on the original MRA patch image. The first kernel forces the pixels with similar positions ($p$) and intensities ($I$) to have similar labels, and the second enforces the smoothness of prediction by only considering the spatial proximity. $\omega_{1}$ and $\omega_{2}$, $\sigma_{\alpha}$, $\sigma_{\beta}$, and $\sigma_{\gamma}$ are hyper-parameters that control the weight and scale of these kernels and are chosen empirically. We used connected-component analysis [19] to identify different clusters when the algorithms predict multiple UIAs, and each connected component was treated as an individual detected aneurysm. To further remove noise, any cluster that had lower than 10 voxels was discarded.

Unfortunately, we were not given access to the ADAM challenge dataset [3]. Instead, we relied on a publicly accessible TOF-MRA dataset of UIAs described in [20] to develop and validate our algorithms. It contains 284 subjects (170 females, 127 healthy controls/157 patients with 198 aneurysms). For 246 subjects, coarse spheres that enclose a whole aneurysm were manually labeled, while 38 subjects have voxel-wise segmentation. All scans were pre-processed with skull-stripping and bias field correction [21], and finally resampled to a median resolution of $0.39\times 0.39\times 0.55$ $mm^{3}$. To allow efficient computation, we performed segmentation based on 3D image patches instead of the entire brain volume. Furthermore, we followed the “anatomically informed” approach in [20] to extract image patches of $64\times 64\times 64$ voxels guided by a probabilistic cerebrovascular atlas with anatomical landmarks on the most probable locations of having an aneurysm, resulting in $\sim$50 patches per subject. In summary, several positive patches with different offsets around every aneurysm were extracted, and for negative patches, a variety of vessel-like, landmark-centered, and random patches were obtained. For more details and the script for image patch extraction, we refer the readers to the original data paper [20]. It is worth mentioning that we used the same patch-extracting strategy for train and test sets. Finally, for DL model training and inferencing, all image patches were normalized with z-transformation. Note that the cases with coarse labels were divided subject-wise for training (95%) and validation (5%), and those with refined segmentation masks were saved for model testing only.

In our proposed FocalSegNet, we utilized a 3D version of the original Focal modulation network [9] as the encoder. Each layer of the encoder consists of two Focal Modulator Blocks, comprising hierarchical depth-wise convolutional layers with kernel sizes of 3, 5, and 7, alongside GELU activation functions. For the baseline models to compare against the FocalSegNet, we also implemented a 3D Residual-UNet that unlike the simple UNet used in the original data paper [20], takes advantage of residual connections for better gradient and information flow, and a Swin-UNETR with default parameters [15] to compare two similar mechanisms, self-attention and focal modulation. The only difference between Swin-UNETR and FocalSegNet is their encoders, which makes our comparison valid. All networks had four layers/hierarchies, equal embedding dimensions, and were trained with a batch size of 12 and an AdamW optimizer (weight decay for L2 regularization=1e-6 and initial learning rate=1e-3). We used a step learning scheduler that reduces the learning rate by 2% every 100 iterations. Furthermore, we used sampling, data augmentation, and a combination of different loss functions to tackle the class imbalance problem. Early stopping is also used to avoid model overfitting.

For the outcomes of the proposed algorithm and its counterparts, as well as the associated ablation studies, we included five different metrics to evaluate their performance in detecting and segmenting the aneurysm based on the test dataset on a per-image-patch basis. Specifically, for UIA detection quality, we measured the sensitivity and false positive (FP) rate. Note that within the same image patch, it is feasible to contain multiple UIAs. Thus, we separated distinct connected components in the prediction map and ground truth, and evaluated them per aneurysm. Here, a correct aneurysm detection is defined as when the centroid of the predicted segmentation component lies within the boundary of an aneurysm and is considered a false positive otherwise. Next, for all the correctly identified aneurysms, we further evaluated the segmentation accuracy using the Dice coefficient, Intersection over Union (IoU), and 95-$\%$ Hausdorff distance (95-HD) measured in voxels to verify the quality of region and boundary overlaps. Because an aneurysm may be detected by one model and missed by the other, we performed an unpaired two-sided t-test to compare the performance between different experimental setups. A $\text{p-value}<0.05$ was used to indicate a significant difference.

The UIA detection and segmentation performance for all three DL models (i.e., Residual-UNet, Swin-UNETR, and FocalSegNet) with CRF finetuning is listed in Table 1, with segmentation results for two subjects illustrated in Fig. 2. In terms of aneurysm detection, for FP rate, our FocalSegNet well outperformed the other two networks ($\text{p}<0.05$), while for sensitivity, all methods perform similarly ($\text{p}>0.05$). This implies that FocalSegNet is better at distinguishing true/false UIAs, which could be hard even for human raters at times. Although FocalSegNet has the best mean segmentation metrics, the difference with Swin-UNETR was not significant ($\text{p}>0.05$). However, both outperform the Residual-UNet ($\text{p}<0.05$).

Besides comparing the performance of the full setup, where fully-connected CRF was used to refine the initial results of the DL models, including FocalSegNet, Swin-UNETR, and Residual-UNet, we also conducted ablation studies to test the impact of CRF post-processing, as well as each component of the full loss function for the proposed FocalSegNet. We evaluated the metrics for UIA detection and segmentation for all the experiments to gain the required insights.

From Table 1, we observe that CRF filtering has a positive impact in boosting the accuracy of the initial segmentation of all models ($\text{p}<0.05$) while FocalSegNet offers the best performance without CRF ($\text{p}<0.05$). Table 2 demonstrates the impact of design choices in terms of the loss function on FocalSegNet. The network trained solely with the Cross-Entropy loss function performs poorly due to the aneurysm size being negligible compared to background voxels, leading to a bias toward negative predictions. Incorporating Generalized Dice Loss enhances performance, but the false positive (FP) rate remains high. To address this, boundary loss is added, enforcing closer proximity between predictions and ground truths. Ultimately, post-processing with a CRF further enhances segmentation performance.