Model Overview

A pre-trained model for volumetric (3D) segmentation of the prostate central gland and peripheral zone from the multi-contrast MRI (T2, ADC).

Model Architecture

This model is trained using the U-Net architecture [1].

The segmentation of prostate region is formulated as the voxel-wise 3-class classification. Each voxel is predicted as either foreground (prostate central gland, peripheral zone) or background. And the model is optimized with gradient descent method minimizing soft dice loss [2] between the predicted mask and ground truth segmentation.

Training

The training was performed with the following:

• Script: train_multi_gpu.sh
• GPU: 4 GPUs of at least 16GB of GPU memory
• Model Input: 96 x 96 x 32 for training, 192 x 192 x 64 for validation/testing
• AMP: True
• Optimizer: Adam
• Learning Rate: 5e-4
• Loss: DiceLoss

If out-of-memory or program crash occurs while caching the data set, please change the cache_rate in CacheDataset to a lower value in the range (0, 1).

Dataset

The training data is from the Medical Decathlon.

Target: Prostate Task: Segmentation Modality: MRI (T2, ADC) Size: 32 3D volumes (21 Training, 6 Validation, 5 Testing) Challenge: Segmenting two adjoint regions with large inter-subject variations The training dataset contains 21 images while the validation and testing datasets contain 6 and 5 images respectively.

The data was converted to resolution 1mm x 1mm x 1mm for training, using the following command:

medl-dataconvert -d ${SOURCE_IMAGE_ROOT} -r 1 -s .nii.gz -e .nii.gz -o ${DESTINATION_IMAGE_ROOT}

Performance

Dice score is used for evaluating the performance of the model. The trained model achieved average testing Dice score 0.6743 over 5 volumes (0.5495 for class 1, and 0.7992 for class 2).

Training

Validation

How to Use this Model

The model was validated with NVIDIA hardware and software. For hardware, the model can run on any NVIDIA GPU with memory greater than 16 GB. For software, this model is usable only as part of Transfer Learning & Annotation Tools in Clara Train SDK container. Find out more about Clara Train at the Clara Train Collections on NGC.

Full instructions for the training and validation workflow can be found in our documentation

Input

Input: 2 channel MRI image with intensity

Preprocessing:

1. Converting to channel first
2. Normalizing to unit Standard Deviation (std) with zero mean
3. Cropping foreground surrounding regions
4. Cropping random fixed sized regions of size 96 x 96 x 32 with the center being a foreground or background voxel at ratio 1 : 1

Augmentation in training:

1. Randomly flipping volumes
2. Randomly rotating volumes
3. Randomly shifting intensity of the volume

Output

Output: 3 channels

• Label 0: background
• Label 1: prostate central gland
• Label 2: prostate peripheral zone

Limitations

This training and inference pipeline was developed by NVIDIA. It is based on a segmentation model developed by NVIDIA researchers. This research use only software has not been cleared or approved by FDA or any regulatory agency. Clara pre-trained models are for developmental purposes only and cannot be used directly for clinical procedures.

References

[1] Çiçek, Özgün, et al. "3D U-Net: learning dense volumetric segmentation from sparse annotation." International conference on medical image computing and computer-assisted intervention. Springer, Cham, 2016. https://arxiv.org/abs/1606.06650.

[2] Milletari, Fausto, Nassir Navab, and Seyed-Ahmad Ahmadi. "V-net: Fully convolutional neural networks for volumetric medical image segmentation." 2016 fourth international conference on 3D vision (3DV). IEEE, 2016. https://arxiv.org/abs/1606.04797.

License

End User License Agreement is included with the product. Licenses are also available along with the model application zip file. By pulling and using the Clara Train SDK container and downloading models, you accept the terms and conditions of these licenses.