3D-Unet: patched based Pytorch implementation for medical images segmentation

Important News -- Repository Maintenance

This repository will no longer be developed and improved. A new version for medical images segmentation is available at https://github.com/davidiommi/SALMON

3D-Unet pipeline is a computational toolbox for segmentation using neural networks.

Same script applied for prostate segmentation can be found here: https://github.com/davidiommi/3D_Prostate_segmentation_pytorch

The training and the inference are patch based: the script randomly extract corresponding patches of the images and labels and feed them to the network during training. The inference script extract, segment the sigle patches and automatically recontruct them in the original size.

Example images

Sample MR images from the sagittal and coronal views for carotid artery segmentation (the segmentation result is highlighted in green)

Requirements

pillow scikit-learn simpleITK keras scikit-image pandas pydicom nibabel tqdm torch>=0.4.1 torchvision>=0.2.1 dominate>=2.3.1 visdom>=0.1.8.

Python scripts and their function

organize_folder_structure.py: Organize the data in the folder structure for the network
NiftiDataset.py : They augment the data, extract the patches and feed them to the network (reads .nii files). NiftiDataset.py skeleton taken from https://github.com/jackyko1991/unet3d-pytorch
check_loader_patches: Shows example of patches fed to the network during the training
UNet.py: the architecture of the U-net.
utils.py : list of metrics and loss functions for the training
main.py: Runs the training and the prediction on the training and validation dataset.
predict.py: It launches the inference on training and validation data in the main.py
predict_single_image.py: It launches the inference on a single input image chosen by the user.

Usage

Use first organize_folder_structure.py to create organize the data in the following folder structure Modify the init.py to set the parameters and start the training/testing on the data:

Folder Structure:

.
├── Data_folder                   
|   ├── train_set              
|   |   ├── patient_1             # Training
|   |   |   ├── image             # Contains domain image 
|   |   |   └── label             # Contains domain label 
|   |   └── patient_2             
|   |   |   ├── image              
|   |   |   └── label              
|   ├── test_set               
|   |   ├── patient_3             # Testing
|   |   |   ├── image              
|   |   |   └── label              
|   |   └── patient_4             
|   |   |   ├── image              
|   |   |   └── label              
|   ├── validation_set               
|   |   ├── patient_5              # Validation
|   |   |   ├── image             
|   |   |   └── label              
|   |   └── patient_6             
|   |   |   ├── image              
|   |   |   └── label

Features

3D data processing ready
Augmented patching technique, requires less image input for training
one channel output (multichannel to be developed)
Generic image reader with SimpleITK support (Currently only support .nii/.nii.gz format for convenience, easy to expand to DICOM, tiff and jpg format)
Medical image pre-post processing with SimpleITK filters
Easy network replacement structure
Dice score similarity measurement as golden standard in medical image segmentation benchmarking