Improving Prostate Cancer Segmentation on T2-Weighted MRI Using Prostate Detection and Cascaded Networks

1. Introduction

Prostate cancer is one of the leading causes of cancer-related deaths in men [1]. One of the reasons for this high mortality is the difficulty in diagnosing the disease in its early stages. The classic method for detecting cancer is Prostate Specific Antigen (PSA) monitoring. If its level increases, the patient is referred for additional testing, including a multiparametric MRI. mpMRI is performed in various modes and sent to a radiologist for analysis. A special PI-RADS scale has been developed for MRI-based cancer diagnosis. This scale is used by specialists to verify detected changes in the prostate gland and determine their clinical significance [2].

When prostate cancer is detected in the late stages, MRI is used to plan surgical interventions. There are advanced developments in the use of artificial intelligence software for surgical planning. For example, AI can be used for 3D reconstruction of the prostate and the use of these data in planning tumor-removal surgery [3,4]. However, detecting cancer in the early stages is difficult because healthy and diseased tissues often have poor contrast.

When considering the application of AI in medicine in general, it is worth mentioning that during the COVID-19 pandemic, many studies were conducted on the use of AI to solve classification and segmentation problems [5,6]. In particular, the main emphasis in the analysis of lung diseases was placed on the analysis of CT studies and radiographs. Among the systems for analyzing CT scans of the lungs, the work of Jin, C., Chen, W., Cao, Y. et al. [7] is noteworthy. They obtained ROC AUC > 0.90 in the analysis of CT studies for lung damage from COVID-19 and pneumonia (viral and non-viral) and also used visualization of the locations to which the neural network responded using grad-cam [8].

Among all the problems in medicine that utilize AI, three fundamental image processing tasks are the most widely used: classification, detection, and segmentation. A classification model assumes the presence of an object that must be assigned to one of the established classes. A prime example of this is the detection of pathology, such as broken ribs in CT scans [9,10] or the work during the pandemic [11].

The detection task is to locate an object of a specified class in an image and construct a frame containing the object. In medicine, detection algorithms are used when it is necessary to precisely localize a detected object. Information about the object's localization potentially allows not only to establish the object's presence but also to obtain more precise data, such as assigning a broken rib to the right or left side of the body [10]. The use of detection has also been explored in detecting gastric cancer in endoscopic images [12].

Segmentation allows for pixel-by-pixel division of an image into classes. In medicine, neural networks for segmentation are primarily used to identify pathologies (for example, brain pathologies) or to identify organ boundaries in CT or MRI scans [13,14].

For most solutions, a single neural network architecture specific to the problem domain is sufficient. However, solving complex problems may require integrated approaches, including the use of multiple neural networks, including networks for specific tasks, such as image registration [15].

Our work aims to identify and describe an optimal method for segmenting prostate cancer on pelvic MRI. The primary objective is to segment the prostate and the cancer within its boundaries. The study used anonymized data from a dataset containing 400 studies. The data was collected at two medical institutions: the Mariinsky Hospital and the Petrov Research Institute. In all MRI studies, specialists with over 10 years of radiology experience mapped the prostate and the cancer.

Among the studies related to ours, we can highlight the study [14], in which a group of researchers used a two-stage approach for cancer segmentation. Wei C. et al., in their study, used a neural network to solve the problem of prostate segmentation (first stage) and cancer segmentation in a refined region of the prostate (second stage). The authors show that the use of a two-stage approach allows improving the results of cancer segmentation by 0.07 for U-Net and 0.18 for attention U-Net, relative to a one-stage approach, when cancer segmentation is performed without preliminary selection of the prostate region. Another example is the study [16] in which the authors used a combination of three parametric series of multi-parametric MRI as data. During the study, they obtained results comparable to an experienced person. Another study that evaluated the capabilities of trained models and compared their results with those of experienced specialists is [17], in which the authors compared the quality of prostate cancer detection of neural networks trained on more than 10,000 MRI studies collected from different locations and different devices, with the results of the work of experienced diagnostic doctors. The results indicate that, with enough data for training, trained neural networks can show results close to the quality of the work of specialists. In addition to the cited works, it is worth noting that [18] emphasizes the relevance of this research direction and the trend toward translating AI from abstract research into commercial medical applications.

In biomedical image analysis, clean masks of the prostate and intraprostatic cancer on MRI are a basic step that many clinical tools rely on. We treat this as an algorithm design problem and use a small, modular cascade: first detect the prostate, then segment the gland, and then segment cancer inside it. By focusing the models on a tight region of interest, the pipeline reduces background noise and false positives, gives more stable results across scanners and sites, and outputs masks that are easy for clinicians to check and reuse.

2. Materials and Methods

2.1. Image Segmentation

Semantic segmentation is an important approach when working with medical images, including problems of delineating individual organs or tumors. The essence of this method is to assign each pixel based on the context of the pixel relative to its environment. Currently, there are many variations of neural network architectures for segmenting medical images, ranging from the classic U-Net network with a U-shaped structure of the encoder and decoder at the core to the integration of Transformers for image processing - SWINUnet, which also has a self-attention mechanism. However, it is worth noting that with the transition from 2D medical image analysis to multi-series data (such as CT and MRI), implementations of architectures for full-fledged analysis of 3D data have appeared. This means creating architectures that can capture context not only within a 2D slice but also across the 3D image volume, which allows for a more accurate determination of the boundaries of organs or tumor inclusions since the spatial relationship between the slices is considered [19,20,21,22,23,24,25,26,27,28,29].

In our study, we used two segmentation settings. In the first, we split 3D MRI data into separate slices and passed them to a network for 2D slice segmentation. UNETR, UNET++, and Swin-UNETR are investigated as neural networks for segmentation. In the second case, training is carried out on loaded 3D data for the segmentation of MRI as a whole. The following architectures are used as trained networks: Swin-UNETR, UNETR, SegResNetDS, SegResNetVAE.

The main metric for segmentation models was the calculation of the Dice-Sorensen (Dice) coefficient. This coefficient quantifies the overlap between two areas. For the task of 2D segmentation, the coefficient is calculated between the true mask of the object (ground truth) and the neural network's prediction (Eq. 1). When implementing the metric in 3D segmentation, additional summation is introduced over all slices with at least one of the masks, and the average coefficient for a series of images is calculated (Eq. 2).

$$dice={\displaystyle \frac{2|\mathrm{X}\bigcap \mathrm{Y}|}{\left|X\right|+\left|Y\right|}}$$

(1)

$$dice={\displaystyle \frac{\sum _i^n\frac{2|X_i\bigcap Y_i|}{\left|X_i\right|+\left|Y_i\right|}}{n}}$$

(2)

where X and $X_i$ are the predicted mask for 2D segmentation and the i-th predicted mask for 3D segmentation, Y and $Y_i$ are the true mask for 2D segmentation, and the i-th true mask for 3D segmentation, n is the number of 3D image slices that contain at least one mask. An example of calculating the metric is shown in Figure 1 below.

2.2. Image Detection

Object detection in an image involves determining the location of an object of a given class if it is present in the image. The location of an object is determined by enclosing the object in certain boundaries, most often in a bounding box. Detection is often applied to medical problems to highlight an area of interest in the context of various tasks, for example, an area containing an organ or a tumor. The development of detection architectures has progressed from basic convolutional networks, such as R-CNN, to more complex ones, such as YOLO in its various versions [30,31]. An example of determining the prostate area in an image is shown in Figure 2.

Our work relied on YOLOv9 to solve the detection problem. A built-in validation set was used to assess the quality of the trained model, including the mAP50–95 metric. This metric shows the area under the Precision-Recall curve for threshold 0.5, 0.95.

2.3. Image Classification

Classification assigns objects to predefined classes. In medicine, neural classification networks are typically used to identify images with pathologies (for example, whether there is a pathology on an X-ray or not) or as a routing step (for example, to determine which part of the body an X-ray image belongs to before sending it to the appropriate analysis service). A full history is beyond the scope of this paper [32,33,34,35,36,37,38,39]. However, fully connected and convolutional neural networks laid down the basic principles, which later grew into such architectures as ResNet and DenseNet, Transformer.

To assess the quality of the model, standard metrics were used (the list is given in the next paragraph), as well as the Area Under the Receiver Operating Characteristic Curve (ROC AUC) metric. This metric is constructed in the coordinates of sensitivity and 1 - specificity, with different thresholds for classifying an object into a target class. This metric clearly shows how accurately the model performs classification. It is worth noting that the ROC AUC metric can also be used to evaluate segmentation models. To do this, you need to set the Dice metric threshold at which segmentation is considered successful and belongs to class 1, and everything below the threshold is classified as class 0. In this case, the ROC curve thresholds relate to the activation threshold for image pixels. In our experiments, a Dice score of 0.5 was chosen as the threshold for successful segmentation.

2.4. Additional Metrics

To comprehensively assess the quality of the trained segmentation models, the study used additional metrics listed in the table below. When calculating the metrics, a threshold of Dice ≥ 0.5 was set to classify the segmentation as “correct segmentation”, which corresponded to “1”.

Table 1. Description of the metrics used.

Metric	Description
Accuracy	Shows the overall correctness of the segmentation
Precision	Shows the proportion of positive pixels relative to all pixels classified as positive
Recall	Shows the proportion of positive pixels relative to all true positive pixels
Specificity	Shows the ability to identify negative pixels correctly
F1	based on precision and recall, and displays a balanced measure of segmentation accuracy

Table 1. Description of the metrics used.

Metric	Description
Accuracy	Shows the overall correctness of the segmentation
Precision	Shows the proportion of positive pixels relative to all pixels classified as positive
Recall	Shows the proportion of positive pixels relative to all true positive pixels
Specificity	Shows the ability to identify negative pixels correctly
F1	based on precision and recall, and displays a balanced measure of segmentation accuracy

2.5. Dataset

As a dataset for training the networks and assessing the quality of the results obtained, we used a closed dataset collected from two hospitals:

• Mariinsky Hospital - 183 MRI studies of the pelvis.
• Petrov Research Institute - 217 MRI studies of the pelvis

In total, the dataset contained 400 anonymized pelvic mpMRI studies from men aged 25–62 years. Each of the studies was reviewed by three specialists who identified signs of clinically significant prostate cancer of different stages. Also, each T2 series of studies from the dataset was marked by radiologists to highlight the prostate area and prostate cancer.

We used the T2 series of each study to train neural networks since they are the clearest and most easily interpreted by humans and are most suitable for analysis by neural network algorithms.

The division of all data into training, validation, and test sets was performed using the built-in tools of the MONAI framework with a fixed random seed in the following ratios:

• for segmentation (2D and 3D): train - 0.7, val - 0.2, test - 0.1;
• for detection: train - 0.85, test - 0.15;
• for classification: train - 0.7, val - 0.2, test - 0.1.

All splits were performed at the study level, such that all slices from a given mpMRI examination were assigned to the same set (train, validation, or test) for every task.

2.6. Hardware and Software Setup

Models were trained using the MONAI v.1.3.2 [40] and PyTorch v.2.4.0 [41] frameworks. The server hardware on which training and quality control of trained models were performed is shown in Table 2. The operating system was Windows Server 2019 Standard.

2.7. Method

Our main goal is to obtain accurate segmentation of prostate cancer. However, during the study, we encountered the fact that simply using a neural network for cancer segmentation did not produce acceptable results on our dataset, so we began searching for methods to improve prostate cancer segmentation. Based on the progress of the entire study and the results obtained, we can divide the research process into three stages: preliminary, main, and final.

During the preliminary stage, we used neural networks for prostate cancer segmentation, as well as an approach that preliminarily utilized a classifier trained to separate sections into two classes: "with cancer" and "without cancer." In the first case, we employed two approaches to cancer segmentation: segmentation of 2D data and segmentation of 3D data. In the second case, we first created a new dataset, retaining only those sections labeled by the classifier as containing cancer. The new dataset was then used to train neural networks for 2D and 3D segmentation. A schematic of the final algorithm, using a single study from this stage as an example, is shown in Figure 3.

To improve the results of the preliminary stage, we decided to use only the prostate region for cancer segmentation. This required us to extract only the prostate region and then train the model on the new data. The main phase of our study was based on prostate-region extraction and consisted of three stages. In the first stage, we used neural segmentation networks to extract the prostate from 2D and 3D data from entire studies, without first creating a new dataset.

To improve the results, we decided to try the idea from the preliminary stage: only feeding slices containing the prostate to the segmentation networks. To do this, we trained a neural network for binary classification (presence/absence of the prostate in the slice). We then created a new dataset. However, at this stage, we modified the rule for creating the new dataset: we now included only slices in the range from the lowest to the highest slice number. This approach is justified by the fact that the prostate is a continuous organ, thus eliminating possible classifier error in the midsection. The new dataset was split according to the proportions specified in 2.6 and used for 2D and 3D segmentation. A schematic representation of the algorithm using a single series as an example is shown in Figure 4.

Since some segmentation errors were found among the segmentation results, we decided to further refine the target region and use a neural network for detection as a preprocessing step. This network detected the prostate and enclosed it in a bounding box. To avoid cropping out part of the prostate (if the detection was inaccurate) and to ensure uniform data size (across the study), we calculated the maximum coordinates among all frames (maximum height and width) and increased them by 5% in each direction.

We then cropped the studies, thereby creating a parallelepiped smaller than the original MRI study but containing the target region with the prostate. Segmentation training was also performed on the new dataset. A schematic representation of the algorithm using a single series as an example is shown in Figure 5.

The final stage involved using the best prostate segmentation method to generate a modified dataset containing only the prostate region and training a cancer segmentation network with the best performance from the preliminary stage. A schematic representation of the entire final stage algorithm is shown in Figure 6.

3. Experiments

3.1. Preliminary Stage. Training 2D and 3D Cancer Segmentation Networks

At this stage, we trained neural networks for 2D cancer segmentation. The training dataset consisted of MRI scans. Three architectures were chosen as training neural networks: UNETR, UNET++, Swin-UNETR. The configuration parameters of the neural networks used, the training configuration, the metric used, and the error function for this and all other neural networks are presented in Tables S1 and S2.

To train neural networks for segmentation on 3D data (entire MRI scans), we chose four architectures: UNETR, Swin-UNETR, SegResNetDS, and SegResNetVAE. The training configurations and training graphs are presented below.

At this stage, we also trained the classification network. DenseNet-201 was chosen as the training architecture. After training DenseNet-201 (Figure 7 and Figure 8), we created a modified dataset for 2D and 3D segmentations. For the 2D data, it consisted of slices from MRI scans that the classification network classified as prostate cancer. For the 3D dataset, we created "stripped" MRI scans, which included only those slices labeled as prostate cancer, as well as one slice before and one slice after the classified sequence. This expansion was used to reduce the risk of classifier error.

Since the 2D dataset contained less data after the classifier, but the data itself remained unchanged, we did not retrain the classification networks; we only tested them on the new dataset. However, the networks for 3D segmentation were retrained because the data in the dataset had essentially changed (the number of slices, and therefore the third dimension in the transmitted tensor, had changed). The same set of networks was used for retraining. The configurations and training process are presented below.

3.2. Main Stage. Training 2D and 3D Prostate Segmentation Networks

To train networks for prostate segmentation, we decided to use the same architecture as in the preliminary stage. Prostate masks were used as data. The training dataset and the training process itself were prepared in a similar manner.

To train prostate segmentation on entire MRI studies (3D data), we also selected the architectures used for 3D prostate cancer segmentation. The training dataset and the training process itself were prepared in a similar manner.

The next step at this stage was to replicate the experience of using a classifier to restrict uninformative slices for segmentation. As in the preliminary stage, the DenseNet-201 neural network was chosen as the classifier. The configuration and training parameters are presented in Table S3.

After training the classifier, we updated the datasets for training 2D and 3D segmentations. The generation rule for the updated datasets was almost identical to the generation rule described in the preliminary stage, except that slices ranging from the minimum slice number classified as prostate to the maximum slice number were used to generate the new 3D data. This approach was adopted from the anatomy of the prostate: it is a continuous organ that is not divided into two or more parts, unlike cancer lesions, which can be multiple in different areas of the prostate gland.

Since the tasks of segmenting prostate cancer and prostate are essentially identical, we again faced the need to retrain the 3D segmentation networks, as the new 3D data was smaller than the original. The network configuration and training process are presented below. We achieved a slight improvement in the main Dice metric, as using the classifier required feeding the network fewer slices, resulting in fewer false positives. Our next step was to use detection instead of classification, as this would minimize the number of unnecessary inclusions and structures in the data passed to segmentation.

As a working approach, we developed the following algorithm: submitting the MRI scan for detection, obtaining bounding box coordinates, finding the maximum coordinates (for each of the four box vertices), incrementing each by 5%, rounding up, and cropping the original MRI scan to the resulting parallelepiped. In this approach, increasing the frame size (by 5% on each side) is necessary to avoid possible random errors in detection. Next, we needed to resize all the obtained data to a uniform size, since the resulting frames were smaller than the original image in any case. The standard image size was chosen to be 128 x 128 x 24 (in the case of 3D data), and resizing was done using PyTorch tools.

The popular and powerful YOLO architecture was used to train the detection network. The training parameters and standard set of post-training metrics are listed below in Table 3 and Figure 9.

Two architectures were chosen as trainable neural networks for working with 2D slices: UNET++ and Swin-UNETR. For training on entire MRI studies, we chose Swin-UNETR and SegResNetDS.

Among all the trained neural networks for segmentation, the Swin-UNETR architecture demonstrated the best results at the preliminary and main stages. Plots for the Dice metric, as well as plots for the error function, are presented below: for 2D segmentation in Figure 10 and Figure 11, respectively, and for 3D segmentation in Figure 12 and Figure 13.

3.3. Final Stage. Dataset Formation and Segmentation Training

The final stage consisted of preparing the dataset and training the neural network for cancer segmentation using new data. To generate the dataset, we used the best method from the main stage, namely, a combination of preliminary detection and subsequent 3D segmentation using the Swin-UNETR network. For training, we selected the architecture with the best 2D and 3D cancer segmentation performance from the preliminary stage—this was also the Swin-UNETR architecture in both cases. The training process is presented below.

Figure 14. Graphs of Dice changes during the training of the Swin-UNETR network at the final stage (a) Dice change from epoch to epoch on the training set; (b) Dice change from epoch to epoch on the validation set.

Figure 14. Graphs of Dice changes during the training of the Swin-UNETR network at the final stage (a) Dice change from epoch to epoch on the training set; (b) Dice change from epoch to epoch on the validation set.

Figure 15. Graphs of Dice changes during Swin-UNETR network training at the final stage: (a) Dice change from epoch to epoch on the training set; (b) Dice change from epoch to epoch on the validation set.

Figure 15. Graphs of Dice changes during Swin-UNETR network training at the final stage: (a) Dice change from epoch to epoch on the training set; (b) Dice change from epoch to epoch on the validation set.

4. Results

In this study, the best prostate cancer segmentation performance was achieved with the final-stage pipeline using 3D cancer segmentation. ROC AUC plots (Figure 16) for the trained segmentation models used in the final stage are presented below.

An example of segmentation of cancer foci by the final 2D cancer segmentation model is shown in Figure 17.

An example of segmentation of cancer foci by the final 3D cancer segmentation model is shown in Figure 18.

A summary of all metrics for each trained segmentation network is provided in Table 4. It does not include data for the classifier metrics or the YOLO detection network, as these networks were trained as a single instance for their task.

5. Discussion

In this paper, we explored methods for improving prostate cancer segmentation results. Although using cancer classification on a slice for further segmentation did not yield significant improvements, this prompted the idea of using the prostate region for cancer segmentation. This approach led to an analysis of conventional prostate segmentation and two ways to improve it: classification and detection. The results show that using detection to isolate the prostate region can increase the accuracy of organ segmentation, which can conceptually be transferred to other applied segmentation tasks.

A potential bottleneck of this approach to segmentation is the loss of algorithmic flexibility in dealing with variable data. During our study, we encountered 8 MRI studies that had significant deviations in contrast values relative to the rest of the dataset. These studies had the worst prostate and cancer segmentation results, even though prostate detection did not materially affect bounding-box localization. One of the future directions of our research will be to find methods to address this issue.

Another issue we encountered, which prompted us to use prostate segmentation for data improvement, was the problem of abrupt data changes when using the classifier. This problem arose when two cancer lesions had sections of healthy prostate tissue between them. With such lesions, when generating a modified dataset after the cancer classifier, sharp boundaries appeared within the segmented data, which was especially noticeable in 3D segmentation. In our work, we avoided this problem by using the prostate region as the transferred data, but it may arise in other segmentation applications using modified datasets.

Beyond technical performance, the proposed pipeline has several potential clinical applications. Robust prostate and cancer segmentation on MRI could support personalized care by assisting radiologists in lesion localization, enabling more precise targeting of biopsies and focal treatments, and providing quantitative tumor volumes for risk stratification and treatment response assessment at the individual-patient level. In radiotherapy, accurate organ and tumor masks are also essential for automated contouring and dose planning; improving segmentation quality may therefore translate into more consistent, patient-specific treatment plans. Prospective studies will be needed to confirm these benefits in real-world clinical workflows.

6. Conclusions

In this study, we explored a cascaded approach to improve prostate cancer segmentation on T2-weighted pelvic MRI. Starting from direct 2D and 3D segmentation of cancer, we incrementally refined the input data by incorporating prostate-focused preprocessing steps: slice-level classification, YOLO-based prostate detection, and 3D prostate segmentation. On a dataset of 400 annotated MRI studies from two clinical centers, the best prostate segmentation performance was obtained with a combination of prostate detection and 3D Swin-UNETR segmentation (Dice 76.09%). Using this refined prostate region as input for 3D cancer segmentation increased the cancer Dice coefficient from 55.03% for direct 3D Swin-UNETR segmentation to 67.11%, with a ROC AUC of 0.89, demonstrating that isolating the prostate region meaningfully improves cancer segmentation quality.

At the same time, the proposed solution increases pipeline complexity, as it relies on several sequential networks for classification, detection, and segmentation. The approach may be sensitive to variations in image acquisition: in our dataset, a small subset of studies with markedly different contrast characteristics showed the poorest segmentation performance despite accurate detection, highlighting robustness as an important limitation. Future work will focus on handling such heterogeneous studies, assessing generalization on additional external cohorts, and investigating whether the cascaded strategy can be simplified or adapted to multi-parametric inputs while preserving or further improving prostate and cancer segmentation accuracy.