NGC | Catalog
Welcome Guest
CatalogModelsclara_pt_unetr_ct_btcv_segmentation

clara_pt_unetr_ct_btcv_segmentation

For downloads and more information, please view on a desktop device.
Logo for clara_pt_unetr_ct_btcv_segmentation

Description

A pre-trained model for volumetric (3D) multi-organ segmentation using CT images from Beyond the Cranial Vault (BTCV) Segmentation Challenge dataset.

Publisher

NVIDIA

Use Case

Segmentation

Framework

PyTorch

Latest Version

4.1

Modified

March 25, 2022

Size

708.47 MB

Model Overview

A pre-trained model for volumetric (3D) multi-organ segmentation using CT images from Beyond the Cranial Vault (BTCV) Segmentation Challenge dataset [1].

Note: The 4.1 version of this model is only compatible with the 4.1 version of the Clara Train SDK container

Model Architecture

This model is trained using the UNETR architecture [2]

UNETR workflow

The task of multi-organ segmentation is formulated as a voxel-wise classification problem. Each voxel is predicted as one of the 13 organ classes in BTCV dataset or background. And the model is optimized with gradient descent method minimizing Dice + cross entropy loss between the predicted mask and ground truth segmentation.

Training

The training was performed with the following:

  • Script: train.sh
  • GPU: At least 32GB of GPU memory.
  • Actual Model Input: 96 x 96 x 96
  • AMP: True
  • Optimizer: Adam
  • Learning Rate: 2e-4
  • Loss: DiceCELoss

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 BTCV dataset.

  • Target: Multi-organs
  • Task: Segmentation
  • Modality: CT
  • Size: 30 3D volumes (24 Training + 6 Testing)

Data Preparation

Please register and download the training data from the BTCV dataset.

NOTE: to match up with the default setting, we suggest that the image root match DATA_ROOT as defined in environment.json in this MMAR's config folder.

Performance

Dice score is used for evaluating the performance of the model. The trained model achieved a Dice score of 0.8270 on the test set (reproduce the results by running validate.sh).

Training

A graph showing the training loss for 5000 epochs (30000 iterations).

Graph that shows training loss over 5000 epochs

Validation

A graph showing the validation mean Dice for 5000 epochs (30000 iterations).

Graph that shows validation mean Dice over 5000 epochs

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 32 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: 1 channel CT image with intensity in HU and fixed spacing (1x1x1 mm)

  1. Add a channel
  2. Resample to resolution 1.5 x 1.5 x 2 mm
  3. Scale intensity
  4. Cropping foreground surrounding regions.
  5. Cropping random fixed sized regions of size [96,96,96] with the center being a foreground or background voxel at ratio 1 : 1.
  6. Randomly applying spatial flipping.
  7. Randomly applying spatial rotation.
  8. Randomly shifting intensity of the volume.

Output

Output: 14 channels

  • Label 0: background
  • Label Non-zero: body organs

Sliding-window Inference

Inference is performed on 3D volumes in a sliding window manner with a specified stride.

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] Landman B, et al. "MICCAI multi-atlas labeling beyond the cranial vault–workshop and challenge." In Proc. of the MICCAI Multi-Atlas Labeling Beyond Cranial Vault—Workshop Challenge 2015 Oct (Vol. 5, p. 12).

[2] Hatamizadeh A, et al. "Unetr: Transformers for 3d medical image segmentation." In Proc. of the IEEE/CVF Winter Conference on Applications of Computer Vision 2022 (pp. 574-584). https://arxiv.org/abs/2103.10504.

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.