Biologically Inspired Selective Attention for Prostate MRI: A Segmentation-Guided Vision Transformer Framework for Computer-Aided Diagnosis

1. Introduction

Prostate cancer (PCa) is one of the most common cancers in men and remains a major global health concern. Early and accurate detection is important for clinical monitoring, treatment planning, and improved patient outcomes [1]. Over the past decade, many studies have explored PCa identification using statistical and machine learning (ML) methods. Although these methods have shown promise, their performance is often limited by simple model structures, small training datasets, and strong dependence on hand crafted features. As a result, their generalization ability remains limited [2]. More recent research has investigated advanced learning paradigms, including transfer learning and Vision Transformers (ViTs). ViTs analyze images through a self attention mechanism, which allows the model to focus on the most informative regions of the input and learn richer feature representations [2]. This is particularly useful in medical imaging, where annotated datasets are often limited. Transfer learning further improves this setting by allowing models pre trained on large datasets to transfer useful visual knowledge to the target task. This can improve diagnostic performance when task specific medical data are scarce [3].

Traditional ML methods, such as support vector machines (SVMs) and linear discriminant analysis (LDA), have also been applied to multiparametric magnetic resonance imaging (mpMRI) for prostate cancer detection. However, these approaches face two main limitations. First, they rely on manually designed features that may not capture the complex and heterogeneous patterns present in prostate MRI. Second, their classification performance is often insufficient for reliable clinical use [3]. Convolutional neural networks (CNNs) have improved many medical image analysis tasks, including segmentation and classification. Even so, they still have important limitations. CNNs mainly learn local spatial patterns and may not model long range anatomical dependencies effectively. This is a major issue in prostate MRI, where the shape of the gland, zonal structure, and broader anatomical context are important for diagnosis. Several extensions, such as larger kernels, deeper architectures, global pooling, dilated convolutions, and non local attention modules, have been proposed to address this problem. However, these solutions often increase model complexity.

ViTs provide a strong alternative. By replacing local convolution with global self attention, ViTs can capture relationships between distant image regions and model broader contextual information more effectively [4]. This property makes them attractive for prostate MRI analysis, where global anatomical context is important. Recent studies have reported encouraging results with ViT based models in several medical imaging tasks, including lung cancer detection from CT scans [5], multi label lung cancer classification from histopathological images [6] and prostate localization [7]. Despite these advances, important gaps remain. Most current studies treat prostate gland segmentation and cancer detection as separate tasks. As a result, the anatomical information learned during segmentation is not fully used to support the diagnosis. In addition, ViT based methods for prostate MRI are still limited, and their use in segmentation guided detection frameworks remains underexplored.

To address these limitations, this paper proposes a segmentation guided Vision Transformer framework for automated prostate cancer detection from mpMRI. The framework first uses a modality aware 3D ViT based model to segment the prostate gland and produce anatomically consistent prostate regions. These segmented regions are then used to guide downstream cancer detection.

The main contributions of this work are as follows:

1.

We propose a segmentation-guided framework for automated prostate cancer detection from mpMRI.

2.

We develop a modality aware 3D Vision Transformer model for prostate gland segmentation from volumetric mpMRI that captures long range anatomical context.

3.

We use the segmented prostate region as an anatomical prior for downstream cancer detection.

4.

We integrate segmentation-guided analysis with both conventional machine learning classifiers and a ViT-based detection model.

5.

We demonstrate strong segmentation and detection performance on the publicly available dataset.

2. Related Work

Prostate cancer is one of the most frequently diagnosed cancers in men and remains a major cause of cancer related mortality worldwide [8]. For this reason, improving prostate cancer diagnosis has become an important research direction. In recent years, artificial intelligence and deep learning have attracted growing interest for prostate MRI analysis [9]. Several AI based strategies have been reported for prostate cancer detection. Yan et al. proposed a classifier based on pixel energy relationships [10]. Midiri et al. developed a computer aided detection system for MRI that identifies local maxima, extracts voxel level features, and segments the prostate using a multi atlas strategy [11]. They also used apparent diffusion coefficient (ADC) maps together with contrast and GLCM homogeneity features derived from T2-weighted images to estimate Gleason scores [12]. In another line of work, radiomic features have been studied using probabilistic support vector machines [13], while T2-weighted MRI has also been used in automatic pipelines for monitoring disease progression [14].

Other studies have explored texture analysis and conventional machine learning methods for prostate cancer assessment. Liu et al. extracted texture features such as energy, entropy, correlation, homogeneity, and inertia from MR images and analyzed them using generalized estimating equations [15]. Kong et al. compared traditional approaches with deep learning methods for prostate cancer detection and reported that GLCM with KNN Cosine was the strongest conventional method in their study [16]. These findings show that hand crafted texture descriptors remain useful, especially when combined with well established classifiers. However, their effectiveness still depends heavily on feature design and may be limited in more complex imaging settings. Deep learning has further improved automated prostate MRI analysis. Abbasi et al. proposed a three stage framework that included lesion localization, ADC normalization, and a stacked auto encoder with a non negativity constraint for prostate cancer detection [17]. Santhirasekaram et al. developed a computer aided diagnosis system using a deep CNN and biopsy defined points to identify prostate cancer from mpMRI [18]. Lopez et al. introduced a deep learning detection system that used filtered block samples for prostate localization [19]. Zavala Romero et al. combined instance and semantic segmentation to integrate local and global image features [20]. In addition, XmasNet used 3D rotation and slicing to better exploit volumetric mpMRI data for lesion analysis [21].

Several related studies have also focused on classification and multimodal analysis. Vos et al. proposed a two stage approach for benign and malignant classification [22]. Chung et al. used voxel level radiological features in a radiation driven framework for cancer detection [23]. Khalvati et al. performed feature selection across multiple MRI modalities and included texture based descriptors in their analysis [24]. Wang et al. later introduced a deformation network combined with multimodal registration for prostate cancer detection [25]. Although these studies have contributed valuable ideas, many of them still treat localization, segmentation, and diagnosis as separate components rather than as a unified framework. Beyond imaging methodology, epidemiological studies have shown that prostate cancer incidence increases markedly with age and rises sharply after the age of 40 [26]. Risk also varies across population groups. African American men face a substantially higher risk than Caucasian men [27], while Asian populations generally show lower incidence rates, although this pattern changes after migration to Western countries [28]. These findings highlight the broader clinical importance of accurate and robust diagnostic tools. Despite substantial progress, important limitations remain. Most existing methods treat prostate gland segmentation and cancer detection as separate tasks, which limits the use of anatomical information during diagnosis. In addition, CNN based models mainly learn local patterns and may not capture long range anatomical dependencies effectively. This is a major limitation in prostate MRI, where gland shape, zonal structure, and broader anatomical context are important. Vision Transformers have recently shown strong performance in medical imaging because of their global self attention mechanism. However, their application to prostate MRI remains limited, especially in segmentation-guided detection frameworks. Moreover, most transformer based studies focus on segmentation accuracy alone and do not examine how segmentation quality affects downstream cancer detection. These limitations motivate the present study, which develops a unified and anatomically informed framework for segmentation-guided prostate cancer detection from mpMRI.

Figure 1 illustrates the proposed framework. Raw mpMRI volumes are first pre-processed and then passed to the segmentation module. The resulting prostate masks are subsequently used to guide the final cancer detection stage.

3. Materials and Methods

3.1. Dataset Description

The PROSTATEx dataset, released by The Cancer Imaging Archive (TCIA), is a publicly available collection of prostate mpMRI studies developed for the SPIE-AAPM-NCI PROSTATEx grand challenge. It includes 346 patients and 349 prostate MRI scans, with imaging sequences such as T2-weighted (T2w), diffusion-weighted imaging (DWI) with apparent diffusion coefficient (ADC) maps, dynamic contrast-enhanced (DCE) imaging, and proton density-weighted images. The challenge framework defined 204 subjects for training and 140 for testing. Clinically significant prostate cancer was defined as biopsy-proven Gleason Grade Group $\ge 2$. Owing to its rich multimodal imaging content, PROSTATEx has been widely used for both prostate cancer detection and segmentation research.

All images were acquired on Siemens 3T MRI scanners, including MAGNETOM Trio and Skyra systems. T2-weighted images were acquired using a turbo spin-echo sequence, with a slice thickness of 3.6 mm and an in-plane resolution of approximately 0.5 mm. DCE time series were acquired using a 3D turbo flash gradient-echo sequence, with a temporal resolution of 3.5 s, a slice thickness of 4 mm, and an in-plane resolution of approximately 1.5 mm. In this study, 200 patients were randomly selected from the PROSTATEx cohort for prostate gland segmentation [29]. Their mpMRI scans were acquired on Siemens 3T MRI scanners. The T2-weighted images had a slice thickness of 3.0 to 4.0 mm and an in-plane resolution of 0.3 to 0.7 mm. The most common resolution, seen in 74% of cases, was $0.5\times 0.5\times 3.0$ mm. The selected subset comprised 134 patients with clinically significant prostate cancer and 66 without clinically significant disease.

3.2. Preprocessing

All preprocessing steps were applied only after patient wise splitting into training, validation, and test subsets. This design prevented information leakage between subsets. During model development, prostate gland masks were used to localize the prostate region, thereby reducing computational cost and focusing the analysis on anatomically relevant tissue. Lesion level annotations were not used during preprocessing. Data augmentation was applied only to the training set in order to improve robustness and reduce overfitting. The validation and test sets were left unchanged to ensure an unbiased assessment of model performance.

3.2.1. Data Augmentation

Data augmentation was used to reduce overfitting and improve model robustness and generalization. In this study, a range of geometric and intensity-based transformations was applied to increase the diversity of training samples while preserving anatomical information. Flipping was used to reflect images along the horizontal or vertical axis [30]. Rotation introduced directional variation by transforming images through selected angles [31]. Shearing modified object shape by preserving one coordinate and shifting the other, while maintaining essential structural information [30]. Cropping was applied to simulate variation in the field of view while retaining relevant image content [31]. Translation repositioned objects without deformation, and the empty regions created during this process were filled with padding to avoid information loss [30].

Several contrast based augmentations were also employed. Gamma contrast adjusted brightness and contrast by varying the gamma value while preserving structural content. The gamma range was set from 0.5 to 2, with a probability of 0.2. Sigmoid contrast enhanced the dark and bright regions separately using a sigmoid intensity transformation, with the threshold coefficient set in the range 5 to 10 and a probability of 0.2. Linear contrast improved image uniformity by scaling intensities in a linear manner. Elastic transformation was used to simulate smooth deformations while preserving the identity and overall structure of the object.

Figure 2 shows representative examples of the augmentation methods applied during training.

3.2.2. Image Enhancement

Image enhancement was applied before feature extraction to improve image quality and increase the visibility of diagnostically relevant structures [32]. The following enhancement methods were used.

Brightness adjustment. Pixel intensities were adjusted using

$$k={\left({\displaystyle \frac{l-n}{h-n}}\right)}^{\gamma }·(d-o)+c,$$

(1)

where l is the lower pixel boundary, h is the upper boundary, n is the input value, o is the output value, c is the minimum output intensity, and d is the maximum output intensity. When the gamma value is greater than one, the image becomes darker, which can enhance contrast and highlight fine anatomical details [33]. However, excessive adjustment may introduce artifacts.

Gamma correction. A nonlinear intensity transformation was used to better match image brightness to the human visual response [34]:

$$T\left(l\right)=l_{\max }{\left({\displaystyle \frac{l}{l_{\max }}}\right)}^{\gamma }.$$

(2)

Figure 3. Effect of image enhancement methods on prostate MRI, including corresponding activation heatmaps. Brightness adjustment, gamma correction, and contrast stretching improve the visibility of anatomical structures compared with the original image.

Figure 3. Effect of image enhancement methods on prostate MRI, including corresponding activation heatmaps. Brightness adjustment, gamma correction, and contrast stretching improve the visibility of anatomical structures compared with the original image.

Contrast stretching. Intensity values were normalized to the range $[a,b]$ using

$$P_{\mathrm{out}}=(P_{\mathrm{in}}-c)\left({\displaystyle \frac{b-a}{d-c}}\right)+a,$$

(3)

where $[a,b]$ denotes the desired output range and $[c,d]$ corresponds to the 5th and 95th percentile input intensities. Using percentile based limits reduces the influence of extreme outliers and improves contrast normalization [35].

3.3. Prostate Gland Segmentation

Prostate gland segmentation is a key component of the proposed framework. In this study, it serves as an anatomically informed intermediate step for downstream prostate cancer detection. Rather than analysing the full MRI volume directly, the framework first identifies the prostate gland and then uses this localized region to guide diagnosis. This design reduces the influence of irrelevant background structures and improves the anatomical focus of the detection stage. In medical imaging, segmentation tasks are commonly divided into semantic segmentation and instance segmentation. Semantic segmentation assigns a class label to each pixel or voxel, whereas instance segmentation additionally separates individual objects of the same class [36]. In the present study, prostate gland delineation is treated as a semantic segmentation problem. MRI provides the anatomical detail needed for this task and supports accurate gland level annotation [37]. Early deep learning-based segmentation methods relied on convolutional encoders to extract latent image representations [38]. More recent methods have extended this idea by incorporating stronger contextual modeling.

Figure 4 summarizes the feature categories relevant to anatomical segmentation. These include motion related information, appearance features, shape based constraints, and contextual relationships. In the present work, appearance is strengthened through image enhancement, while shape and context are modeled through the proposed segmentation framework. The aim is to produce anatomically consistent prostate gland masks that can serve as reliable inputs to the downstream cancer detection stage.

Figure 4. Feature categories used for anatomical segmentation.

Figure 4. Feature categories used for anatomical segmentation.

Figure 5. Prostate MRI segmentation using adaptive thresholding. Original and thresholded images are shown for both classes.

Figure 5. Prostate MRI segmentation using adaptive thresholding. Original and thresholded images are shown for both classes.

3.4. Adaptive Thresholding

Global thresholding performs poorly when image contrast is non-uniform or the background lacks uniformity. Adaptive thresholding addresses this by computing local thresholds. A binary segmentation mask is defined as:

$$I(a,b)=\left\{\begin{array}{cc}1\hfill & \mathrm{if}I(a,b)>\tau ,\hfill \\ 0\hfill & \mathrm{otherwise},\hfill \end{array}\right.$$

(4)

where $\tau$ is the locally computed threshold.

Figure 6. Prostate MRI segmentation using morphological operations. Original and post-morphology images are shown for both classes.

Figure 6. Prostate MRI segmentation using morphological operations. Original and post-morphology images are shown for both classes.

3.5. Optimized U-Net Architecture

U-Net remains one of the most widely used architectures for medical image segmentation because it combines multi scale feature extraction with effective dense prediction. Its encoder captures local image patterns, while its decoder reconstructs spatial detail at the output resolution. However, convolutional operations mainly model local neighborhoods and may not capture long range anatomical dependencies effectively. This limitation is important in prostate MRI, where gland shape and broader anatomical context influence segmentation quality. For this reason, the optimized U-Net is used in this work as a strong baseline for comparison with the proposed transformer based segmentation model.

Figure 7. Optimized U-Net architecture used for prostate gland segmentation.

Figure 7. Optimized U-Net architecture used for prostate gland segmentation.

3.6. Proposed Segmentation Model

The proposed segmentation model is designed to generate anatomically consistent prostate gland masks that can guide downstream cancer detection. Unlike conventional convolution based architectures, the proposed model uses modality aware inputs together with global self attention to capture long range anatomical dependencies in prostate MRI. This enables the network to model both local gland boundaries and broader contextual structure within the volume.

3.6.1. Problem Formulation

Given a multi-modal 3D prostate MRI volume:

$$\mathbf{X}\in {\mathbb{R}}^{C\times D\times H\times W},$$

(5)

where $C=3$ corresponds to the T2w, ADC, and DWI modalities, the objective is to predict a voxel wise binary segmentation mask: $\widehat{\mathbf{Y}}\in {[0,1]}^{1\times D\times H\times W}.$ The segmentation model learns the mapping: $f_{\theta }:\mathbf{X}\to \widehat{\mathbf{Y}}$ where $\theta$ denotes trainable parameters. The mapping is optimized not only for voxel-level overlap but also to preserve the global anatomical coherence needed for reliable downstream cancer detection.

3.6.2. Input Normalization

Each modality is independently normalized using min max normalization:

$$X^{\prime }={\displaystyle \frac{X-\min \left(X\right)}{\max \left(X\right)-\min \left(X\right)+ϵ}}.$$

(6)

3.6.3. Patch Embedding

The input volume is partitioned into non-overlapping 3D patches of size $P=(p_D,p_H,p_W)$. The total number of patches is:

$$N={\displaystyle \frac{D}{p_D}}·{\displaystyle \frac{H}{p_H}}·{\displaystyle \frac{W}{p_W}}.$$

(7)

Each patch is projected into a latent embedding space of dimension E via a 3D convolution:

$${\mathbf{Z}}_0=\mathrm{Flatten}\left(\mathrm{Conv}3\mathrm{D}\left(\mathbf{X}\right)\right),{\mathbf{Z}}_0\in {\mathbb{R}}^{N\times E}.$$

(8)

Learnable positional embeddings are added to preserve spatial ordering: ${\mathbf{Z}}_0\leftarrow {\mathbf{Z}}_0+{\mathbf{E}}_{\mathrm{pos}}.$

3.6.4. Transformer Encoder

Each transformer layer consists of a multi-head self attention (MHSA) block and a feed forward network (FFN).

Multi-Head Self-Attention.

Query, key, and value matrices are computed as:

$$Q=\mathbf{Z}{W}_Q,K=\mathbf{Z}{W}_K,V=\mathbf{Z}{W}_V.$$

(9)

Scaled dot-product attention is:

$$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left({\displaystyle \frac{Q{K}^{\top }}{\sqrt{d_k}}}\right)V.$$

(10)

The MHSA output concatenates H attention heads:

$$\mathrm{MHSA}\left(\mathbf{Z}\right)={\displaystyle \underset{h=1}{\overset{H}{⨁}}}{\mathrm{Attention}}_h.$$

(11)

A residual connection followed by layer normalization is applied: ${\mathbf{Z}}^{\prime }=\mathbf{Z}+\mathrm{MHSA}\left(\mathrm{LN}\left(\mathbf{Z}\right)\right).$

Feed-Forward Network.

The FFN uses GELU activation: $\mathrm{FFN}\left(x\right)=W_2\mathrm{GELU}\left(W_{1}x\right).$ With a residual connection: ${\mathbf{Z}}_{l+1}={\mathbf{Z}}^{\prime }+\mathrm{FFN}\left(\mathrm{LN}\left({\mathbf{Z}}^{\prime }\right)\right).$

3.6.5. Decoder and Reconstruction

Transformer output tokens are reshaped into a 3D feature map: $\mathbf{F}\in {\mathbb{R}}^{E\times D_p\times H_p\times W_p}.$ The decoder progressively upsamples $\mathbf{F}$ to the original resolution: $\widehat{\mathbf{Y}}=\sigma \left(\mathcal{D}\left(\mathbf{F}\right)\right)$ where $\mathcal{D}(·)$ denotes convolutional and transposed convolutional layers, and $\sigma$ is the sigmoid activation function.

3.6.6. Loss Function

The total training loss is a sum of binary cross-entropy (BCE) loss and Dice loss:

$$\mathcal{L}={\mathcal{L}}_{\mathrm{BCE}}+{\mathcal{L}}_{\mathrm{Dice}}.$$

(12)

Binary Cross-Entropy Loss.

$${\mathcal{L}}_{\mathrm{BCE}}=-{\displaystyle \frac{1}{N}}{\displaystyle \sum _{i=1}^N}\left[y_i\log p_i+(1-y_i)\log (1-p_i)\right].$$

(13)

Dice Loss.

$$\mathrm{Dice}={\displaystyle \frac{2{\sum }_i{p}_i{y}_i+ϵ}{{\sum }_i{p}_i+{\sum }_i{y}_i+ϵ}},{\mathcal{L}}_{\mathrm{Dice}}=1-\mathrm{Dice}.$$

(14)

For empty ground truth masks (${\sum }_i{y}_i=0$), the Dice loss term is excluded from the update.

3.6.7. Optimization

Model parameters are optimized using AdamW:

$${\theta }_{t+1}={\theta }_t-\eta \left({\displaystyle \frac{{\widehat{m}}_t}{\sqrt{{\widehat{v}}_t}+ϵ}}+\lambda {\theta }_t\right).$$

(15)

A warm up cosine learning rate schedule is used:

$${\eta }_t={\eta }_{\min }+{\displaystyle \frac{1}{2}}({\eta }_{\max }-{\eta }_{\min })\left(1+\cos \left(\pi t\right)\right),$$

(16)

where ${\eta }_{\max }=3\times {10}^{-4}$ is the peak learning rate, ${\eta }_{\min }=1\times {10}^{-6}$ is the minimum learning rate, and $t\in [0,1]$ is the normalized training progress after the warm-up phase. A linear warm-up over the first 5 epochs increases the learning rate from 0 to ${\eta }_{\max }$. Cosine annealing is applied for all remaining epochs.

Figure 8. Vision Transformer (ViT) architecture used for prostate cancer detection.

Figure 8. Vision Transformer (ViT) architecture used for prostate cancer detection.

4. Prostate Cancer Detection Methods

After segmentation, the localized prostate region is used for downstream cancer detection. Two complementary detection strategies are considered in this study. The first is a conventional machine learning pipeline based on hand crafted texture features. The second is a ViT based classifier. This dual evaluation makes it possible to assess the value of segmentation guided analysis across both classical and deep learning settings.

4.1. Feature-Based Machine Learning Pipeline

Classical machine learning approaches rely on hand crafted features extracted before classification. In this work, texture descriptors are computed from the segmented prostate region and then used to train established classifiers for cancer detection [39].

4.1.1. Feature Extraction

The textural properties of cancerous prostate tissue differ from those of normal tissue. Measures such as mean, standard deviation, and entropy are commonly used to capture these textural fluctuations. GLCM-based features are extracted following the method of Iqbal et al. [40].

Gray-Level Co-occurrence Matrix (GLCM) Features. Table 1 represents GLCM features. The GLCM encodes spatial co-occurrence statistics between pixel pairs and captures both spatial dependencies and textural correlations.

where $\mu$ is the mean intensity and ${\sigma }_k$, ${\sigma }_l$ are standard deviations of the row and column marginal distributions.

4.1.2. Classifiers

Naive Bayes. The Naive Bayes classifier applies Bayes’ theorem under the assumption of feature independence. It estimates the posterior probability of each class given the observed features and selects the most probable class [41]. Despite its simplifying assumptions, Naive Bayes is computationally efficient and performs well across diverse medical datasets [42].

Decision Tree. Decision trees guide each sample through a series of attribute-based decisions toward a predicted class label [43]. They provide intuitive feature importance rankings, supporting dimensionality reduction as a preprocessing step.

RUSBoosted Tree. RUSBoost is an ensemble method that combines random under-sampling with gradient boosting to handle class imbalance [44,45]. Misclassified instances receive progressively higher weights across boosting iterations.

Support Vector Machine (SVM). Given a labeled training set $T={\left\{(a_i,b_j)\right\}}_{i=1}^n$ with $a_i\in {\mathbb{R}}^d$ and $b_j\in \{-1,+1\}$, an SVM seeks the hyperplane: $W·I+B=0$ that maximizes the inter-class margin:

$$\mathrm{Margin}={\displaystyle \frac{2}{\parallel W\parallel }}.$$

SVMs are well-suited to moderate-sized, high-dimensional datasets and offer interpretable decision boundaries [46,47].

4.2. Vision Transformer with Transfer Learning

The second detection branch uses a Vision Transformer based classifier. Transformers model long range dependencies through self attention and can therefore capture broader contextual relationships than purely local architectures. In this work, the ViT based classifier operates on the segmented prostate region and is used as the main deep learning detection model. Let $X\in {\mathbb{R}}^{T\times C}$ be the input, where T denotes the number of tokens and C the embedding dimension. Query, key, and value projections are computed as:

$$\begin{array}{cc}\hfill Q& =X\times W^Q,\hfill \end{array}$$

(17)

$$\begin{array}{cc}\hfill K& =X\times W^K,\hfill \end{array}$$

(18)

$$\begin{array}{cc}\hfill V& =X\times W^V,\hfill \end{array}$$

(19)

where $W^Q,W^K\in {\mathbb{R}}^{C\times d_k}$ and $W^V\in {\mathbb{R}}^{C\times d_v}$. The attention output is:

$$h=\mathrm{softmax}\left({\displaystyle \frac{Q{K}^{\top }}{\sqrt{d_k}}}\right)V.$$

(20)

The scaling factor $1/\sqrt{d_k}$ prevents excessively small gradients in the softmax operation.

Positional information is encoded via sinusoidal embeddings:

$$E(p,2i)=\sin \left({\displaystyle \frac{p}{{10000}^{2i/d_{\mathrm{model}}}}}\right),$$

(21)

where p is the token position and i the dimension index.

Multi-head attention extends single-head attention by projecting queries, keys, and values into multiple subspaces and concatenating their outputs. This allows the model to jointly attend to information from different representations sub-spaces at different positions.

5. Performance Measures

The following metrics were used to evaluate segmentation and classification performance.

5.1. Dice Similarity Coefficient (DSC)

The Dice Similarity Coefficient measures the spatial overlap between predicted and ground truth segmentation masks:

$$\mathrm{DSC}={\displaystyle \frac{2|\mathrm{Target}\cap \mathrm{Predicted}|}{\left|\mathrm{Target}\right|+\left|\mathrm{Predicted}\right|}}\times 100={\displaystyle \frac{2\mathrm{TP}}{\mathrm{TP}+\mathrm{FN}+\mathrm{FP}}}\times 100,$$

(22)

where TP, FP, and FN denote true positives, false positives, and false negatives, respectively.

5.2. Intersection over Union (IoU)

The IoU (Jaccard Index) measures the ratio of the intersection to the union of the predicted and ground truth regions []:

$$\mathrm{IoU}={\displaystyle \frac{\left|\mathrm{Area}\mathrm{of}\mathrm{Overlap}\right|}{\left|\mathrm{Area}\mathrm{of}\mathrm{Union}\right|}}\times 100.$$

(23)

5.3. Precision

Precision (positive predictive value) measures the fraction of predicted positives that are truly positive:

$$\mathrm{Precision}={\displaystyle \frac{\mathrm{TP}}{\mathrm{TP}+\mathrm{FP}}}.$$

(24)

5.4. F1-Score

The F1-score is the harmonic mean of precision and recall:

$$F_1={\displaystyle \frac{2\times \mathrm{Precision}\times \mathrm{Recall}}{\mathrm{Precision}+\mathrm{Recall}}}.$$

(25)

5.5. Receiver Operating Characteristic (ROC) Curve

The area under the ROC curve (AUC) summarizes classifier performance across all decision thresholds. An AUC of 1.0 indicates a perfect classifier; 0.5 indicates chance performance. AUC values in the range $[0.9,1.0]$ are generally considered excellent. The ROC metric is particularly informative for imbalanced datasets, as it evaluates both sensitivity (true positive rate) and specificity (1 − false positive rate) [].

5.6. Implementation Details

All experiments were conducted on a workstation equipped with an 11th generation Intel Corei9-11900KF processor (3.50GHz), 64GB of system RAM, and an NVIDIA GeForce RTX 3060 GPU with 12 GB of dedicated memory. The framework was implemented in PyTorch. Standard scientific computing libraries (NumPy, Pandas, Nibabel) were used for handling volumetric MRI data. The model uses a 3D ViT-based encoder–decoder architecture. Input volumes were resized to $96\times 128\times 128$ voxels. The patch embedding module divides the volume into non-overlapping $8\times 16\times 16$ patches, which are projected into a 192-dimensional embedding space. The transformer encoder comprises 12 layers with 6 self attention heads and a feed forward expansion ratio of 4.0. Positional embeddings are learned. The decoder reconstructs the segmentation map using convolutional and transposed convolutional layers. The model is trained with a combined Dice + BCE loss. A batch size of 2 and gradient accumulation of 2 steps yield an effective batch size of 4. Optimization uses AdamW with ${\eta }_{\max }=3\times {10}^{-4}$ and weight decay $\lambda =1\times {10}^{-4}$. The learning rate follows a warm-up cosine schedule over 100 epochs (5 warm-up epochs, minimum learning rate $1\times {10}^{-6}$).

Figure 9. Visual comparison of original image, ground truth mask, and the predicted mask produced by the proposed 3D ViT segmentation model. Row 1: The predicted mask captures the overall shape of the prostate gland but slightly overestimates its extent on the left. Row 2: The predicted mask aligns closely with the ground truth, with only minor discrepancies at fine boundary details.

Figure 9. Visual comparison of original image, ground truth mask, and the predicted mask produced by the proposed 3D ViT segmentation model. Row 1: The predicted mask captures the overall shape of the prostate gland but slightly overestimates its extent on the left. Row 2: The predicted mask aligns closely with the ground truth, with only minor discrepancies at fine boundary details.

Figure 10. Training curves of the optimized U-Net model: accuracy and Dice coefficient versus epoch. Both metrics increase steadily and stabilize after approximately epoch 100, indicating convergence.

Figure 10. Training curves of the optimized U-Net model: accuracy and Dice coefficient versus epoch. Both metrics increase steadily and stabilize after approximately epoch 100, indicating convergence.

Figure 11. Overlay visualization of the ground truth mask and predicted mask from the proposed 3D ViT segmentation model on a representative case.

Figure 11. Overlay visualization of the ground truth mask and predicted mask from the proposed 3D ViT segmentation model on a representative case.

6. Results and Discussion

Table 2 and Table 3 summarize the segmentation results. The proposed model achieves a Dice score of 0.9048 and an IoU of 0.728, indicating substantial overlap with the ground truth masks. Precision (0.889) and recall (0.803) are well balanced, indicating good control over both false positives and false negatives. Boundary fidelity is confirmed by a Hausdorff Distance of 2.11mm and an Average Surface Distance of 3.31mm. High ROC-AUC (0.987) and PR-AUC (0.928) further confirm robust pixel-level discriminative capacity. All metrics exceed those of the optimized U-Net baseline, demonstrating the benefit of incorporating global self-attention into the segmentation architecture.

6.1. Cancer Detection Performance

Both deep and classical classifiers were evaluated: SVM with radial basis function (RBF) kernel, Naive Bayes, RUSBoosted Tree, Decision Tree, and the proposed ViT with transfer learning. A patient-wise split was enforced throughout to prevent data leakage and ensure unbiased evaluation on entirely unseen patients.

Figure 12 summarizes the detection performance of the evaluated classifiers. The proposed ViT based classifier achieves the strongest overall balance across sensitivity, specificity, PPV, NPV, precision, and accuracy, indicating reliable discrimination between cancer and non cancer cases within the segmented prostate region. Some conventional classifiers achieve slightly higher AUC values, showing that they remain competitive in ranking performance. Overall, however, the proposed segmentation guided ViT pipeline provides the most consistent detection performance across the broader set of reported evaluation measures.

Table 4 reports per class classification results for the ViT model with adaptive thresholding and morphological preprocessing. The model achieves 99.8% overall accuracy and balanced performance across both classes. Precision of 0.998 indicates that nearly all predicted positive cases are true positives. A recall of 0.987 indicates that the model misses very few genuine PCa cases. The F1 score of 0.987 for the PCa class reflects a good balance between precision and recall. These characteristics are important for clinical deployment, where both false positives and false negatives carry significant consequences.

Figure 13. ROC curve for prostate cancer detection using the proposed 3D ViT segmentation-guided model. The AUC of 0.987 indicates excellent discriminative performance.

Figure 13. ROC curve for prostate cancer detection using the proposed 3D ViT segmentation-guided model. The AUC of 0.987 indicates excellent discriminative performance.

Figure 14. Training loss curve of the proposed segmentation model over 200 epochs. Loss decreases rapidly in the early epochs as the model learns coarse anatomical and prostate structures, followed by a gradual, stable decline. The logarithmic scale reveals continued optimization at later epochs. Convergence is confirmed by minor fluctuations in the final 50 epochs.

Figure 14. Training loss curve of the proposed segmentation model over 200 epochs. Loss decreases rapidly in the early epochs as the model learns coarse anatomical and prostate structures, followed by a gradual, stable decline. The logarithmic scale reveals continued optimization at later epochs. Convergence is confirmed by minor fluctuations in the final 50 epochs.

6.2. Ablation Study

Table 5 presents the ablation results. Each added component improves overall performance. Preprocessing and augmentation increase the AUC from 0.850 to 0.880. Adding segmentation guidance further improves the AUC to 0.900 and yields a Dice score of 0.900. The full model achieves the best overall performance, indicating that segmentation guidance, preprocessing, and augmentation each contribute to the final result.

6.3. Comparison with State of the Art

Table 6 compares the proposed method with recent published work. Hussain et al. achieved AUC = 0.99 using morphological and SIFT features [48]. Jung et al. reported AUC = 0.90 with a radiomics-only approach [49]. Ali et al. obtained DSC = 0.80 for volumetric segmentation on two datasets [50]. Giganti et al.reported AUC = 0.85 with a combined CAD and deep learning system [51]. The proposed framework achieves a Dice score of 0.905 and a detection AUC of 0.987, which are competitive with those reported in the literature while also providing explicit segmentation guidance. The proposed method offers the additional advantage of integrating prostate gland segmentation directly into the detection pipeline, thereby improving anatomical focus and strengthening its relevance to clinically meaningful mpMRI analysis. Regarding computational characteristics: ViT uses a self attention mechanism that scales linearly with the number of patches for large images, whereas ResNet101 scales quadratically with image size due to its convolutional operations [39,52]. ViT primarily stores the attention matrix, and optimizations such as relative attention and windowed attention significantly reduce its memory footprint [53]. This makes ViT is more scalable for high-resolution volumetric inputs such as prostate mpMRI.

7. Conclusion and Limitations

This study proposed a segmentation guided Vision Transformer framework for automated prostate cancer detection from multiparametric MRI. The framework first segments the prostate gland and then uses the segmented prostate region to guide downstream cancer detection. By combining anatomical localization with diagnostic prediction, the proposed method improves the anatomical focus of the detection stage and captures long range contextual information that is difficult to model with conventional convolution based approaches. Experimental results on the publicly available PROSTATEx dataset showed strong segmentation performance and high downstream detection performance. The findings suggest that anatomically guided analysis can improve the reliability of automated interpretation of prostate MRI by reducing the influence of irrelevant background structures. These results support the value of combining prostate segmentation with transformer based global contextual modeling for computer aided prostate cancer diagnosis. This study also has important limitations. The framework was evaluated on a single publicly available dataset, and broader validation is needed to assess generalizability across institutions, scanner settings, and imaging protocols. In addition, the current method focuses on prostate gland segmentation and does not include lesion level segmentation. Future work will therefore examine multi center validation, integration of clinical variables, lesion level analysis, and improvements in computational efficiency for clinical deployment.