SwinVNETR: Swin V-net Transformer with non-local block for volumetric MRI Brain Tumor Segmentation

Figures & data

Figure 1. Overall architecture of Swin VNETR.

Figure 2. (a). Swin 3D blocks included W-MSA-3D and MLP, (b). 3D convolutional block.

Figure 3. Structure of non-local block.

Figure 4. Flow diagram for proposed model execution.

Figure 5. BraTs2021 native (T1) and post-contrast (T1Gd) T1-weighted (T1Gd) volumes, T2-weighted (T2) and T2-FLAIR volumes 3D MRI Images.

Table

Pseudo-code for Proposed Model
1  input_dimensions = (batch_size, channels, height, width, depth) # Input shape for 3D MRI scans
2  num_classes = 3 # Number of classes: ET, WT, TC
3  def encoder(x):
4    x = ConvBlock(x, filters = [64, 128, 256]) # filter sizes
5    # Add Swin Transformer layers here
6    swin_layers = []
7    for _ in range(num_swin_layers):
8        swin_layer = SwinTransformerLayer(embedding_dim, num_heads, window_size)
9      swin_layers.append(swin_layer)
10   for i, swin_layer in enumerate(swin_layers):
11      res_block = ResidualConnection(swin_layer(x))
12      x = DownSampling(res_block, stride = stride)
12   # Non-Local Block
13    x = NonLocalBlock(x)
14   return x
15  def decoder(features, skip_connections):
16      upsampled_features = UpSampling(features)
17      x = concatenate([upsampled_features, skip_connections])
18      x = ConvBlock(x, filters = [256, 128, 64])
19      x = ConcatenateFeatures() # Combine multiple levels of features
20      x = ConvBlock(x, filters = [64, num_classes])
21      x = SoftmaxActivation(axis = −1)
22      return x
23  def BrainTumorSegmentationNetwork():
24      input_layer = Input(shape = input_dimensions)
25      encoded_features = encoder(input_layer)
26      skip_connections = []
27      for _ in range(len(encoder.__dict__[‘__globals__’][‘swin_layers’])):
28        skip_connection = UpSampling(encoded_features)
29        skip_connections.append(skip_connection)
30        encoded_features = encoded_features + skip_connection
31    segmentation_output = decoder(encoded_features, skip_connections)
32    model = Model(inputs = input_layer, outputs = segmentation_output)
33    compile_model(model, optimizer = 'adamW’, loss = 'binary_crossentropy’)
34    return model

Figure 6. A) Brain image with tumour, B) Whole Tumor highlighted in green, C) Enhanced Tumor in yellow, D) Tumor Core in red, visualized using ITK-SNAP [25] software tool used to view the brain tumour segmentation.

Figure 7. A. an Flair image as an input B. the ground truth of the validation data C. the feature extraction process during the Non-Local block D. the final segmented WT, TC, and ET.

Figure 8. Curves illustrating the variation in Dice coefficients were used to assess the model's fitting capability on the validation dataset of Brats2021. These curves depict the Dice coefficients for the ET, TC, and WT channels.

Table 1. Comparison table of brats 2021 validation dataset.

		Mean Dice			Hausdorff distance
Name of the model	Params Size	ET	TC	WT	ET	TC	WT
3D U-Net	15.834	0.825	0.844	0.900	3.71	7.36	4.82
V-Net	182.432	0.815	0.840	0.751	2.88	6.17	6.43
UnetR	204.899	0.842	0.853	0.905	3.12	8.46	4.59
TransBTS	65.975	0.824	0.843	0.889	3.89	5.54	4.22
SwinBTS	71.394	0.828	0.843	0.896	4.02	5.16	4.68
AttentionUnet	47.257	0.841	0.851	0.870	3.57	6.83	4.06
Swin Pure Unet3D	67.163	0.817	0.822	0.885	4.32	7.56	3.81
Proposed model	67.487	0.838	0.871	0.905	3.51	5.21	2.89

Table 2. p-value comparison table of all other model dice scores from Swin VNETR.

	P – Value
Model Name	ET	TC	WT
3D U-Net	0.2154	0.0122	0.0002
V-Net	0	0	0
UnetR	0	0	0
TransBTS	0.0028	0.0001	0.0428
SwinBTS	0.2941	0.3174	0.5322
AttentionUnet	0	0	0
Swin Pure Unet3D	0	0	0

Figure 9. Visualization of Input Image, Ground Truth, Whole tumour, Tumor core, Enhanced Tumor, and Grad-CAM highlighted.

Yushkevich PA, Piven J, Cody Hazlett H, et al. User-guided 3D active contour segmentation of anatomical structures: significantly improved efficiency and reliability. Neuroimage. 2006;31(3):1116–1128.

 PubMed Web of Science ®Google Scholar