cradhakr/uhura_unet_mss repository overview

⁠UHURA_UNet_MSS

The docker image contains pretrained UNet-MSS network for performing vessel segmentation on 3D nifti volumes. The network was trained with modifed UNET-MSS as in the paper "DS6, Deformation-aware Semi-supervised Learning: Application to Small Vessel Segmentation with Noisy Training Data"⁠. The project code can be found at https://github.com/ravi22342/DS6/tree/VSeg_with_MIP⁠.

⁠Docker RUN Instructions

The container is flexible and capable of performing several different experiments. However, the pretrained weights included in this docker image is the one resulting from training UNET-MSS baseline.

⁠Generating inference on 3D nifti volumes

To generate segmentations for a set of 3D nifti volumes please use

docker run --gpus all --rm -v <path-to-folder-containing-nifti-vols>:/data -v <path-to-op>:/output cradhakr/uhura_unet_mss -dataset_path /data -output_path /output -eval True

This creates a folder 'uhura_unet_mss' inside specified output path and this folder will contain the resulting segmentations with same name as input nifti volumes.

Additionally, you can modify the parameter defaults below depending on the GPU capability of the target system:

-batch_size 15 -num_worker 8 -patch_size 64 

Reduce the batch_size and patch_size in order to fit the volumes in available GPU memory. num_worker parameter allows creation of multiple worker threads for tqdm. Set this to 0 if you encounter Dataloader worker errors.

⁠Training the underlying UNet-MSS from scratch

To train the UNet-MSS baseline network from scratch, please use

docker run --gpus all --rm -v <path-to-dataset-folder>:/data -v <path-to-op>:/output cradhakr/uhura_unet_mss -dataset_path /data -output_path /output -train True -with_mip False -pre_train False

Arrange the dataset folder as shown below:
/<path-to-dataset-folder>
---/train
---/train_label
---/validate
---/validate_label
---/test
---/test_label
*All the folders mentioned above must contain atleast one 3D nifti volume

If you wish to perform training with k-fold cross validation include below args

-cross_validate True -k_folds 5

In this case, please arrange the daset folder as shown below:
/<path-to-dataset-folder>
---/train
---/train_label
---/test
---/test_label
*All the folders mentioned above must contain atleast one 3D nifti volume The network weights are stored at <path-to-op>/uhura_unet_mss/checkpoint/ as checkpointbest.pth after successful training.

⁠Tuning the pretrained UNet-MSS network with MIP loss

The network architecture is illustrated in the above mentioned papers. Please have a look at the network in https://github.com/ravi22342/DS6/tree/VSeg_with_MIP⁠ for more info. Once the pretrained weights are availble from training as specified above or if you wish to further tune the weights already available in the container, please use

docker run --gpus all --rm -v <path-to-dataset-folder>:/data -v <path-to-op>:/output -v <path-to-checkpointbest.pth>:/checkpoint cradhakr/uhura_unet_mss -dataset_path /data -output_path /output -train True -with_mip False -pre_train True -load_path /checkpoint

*Provide the checkpoint path if the training was done from scratch using -load_path. If you would like to use the existing weights, do not specify this arg.

⁠Comparing segmentation results against labels and obtaining diff MIP overlays

If you wish to compare the segmentations that will be generated using the container against segmentation labels, please use

docker run --gpus all --rm -v <path-to-dataset-folder>:/data -v <path-to-op>:/output -v <path-to-checkpointbest.pth>:/checkpoint cradhakr/uhura_unet_mss -dataset_path /data -output_path /output -test True -with_mip False -pre_train True -load_path /checkpoint

In this case the dataset folder must be arranged as below: /<path-to-dataset-folder>
---/test
---/test_label
A results folder will be created at <path-to-op>/uhura_unet_mss/ and this will contain resulting segmentations, MIPs of diff overlays of labels and their respective segmentations and a spreadsheet containing evaluated metrics (Dice Score and IoU) for each nifti volume in /test folder.

⁠Singularity RUN Instructions

In order to run this docker image inside a Singularity⁠ container, first we need to build the Singularity Image File(.sif) using appropriate build definitions from a Singularity file or a .def file. This is to accomodate for docker-to-singularity compatibility.

⁠Building .sif from .def

The corresponding .def file is available at https://github.com/ravi22342/DS6/tree/VSeg_with_MIP/Configs/singularity⁠. Download this file and execute the below command

sudo singularity build uhura_unet_mss.sif uhura_unet_mss.def

This will create a singularity container named uhura_unet_mss.sif by downloading the latest docker image > performing conda init > activating the conda environment with all the dependencies.
*Include --sandbox to build a sandbox container where you can access the container files.

⁠Generating inference on 3D nifti volumes

To generate segmentations for a set of 3D nifti volumes please use

singularity run --contain --nv --bind <path-to-folder-containing-nifti-vols>:/data,<path-to-op>:/output uhura_unet_mss.sif -eval True

This creates a folder 'uhura_unet_mss' inside specified output path and this folder will contain the resulting segmentations with same name as input nifti volumes.

Additionally, you can modify the parameter defaults below depending on the GPU capability of the target system:

-batch_size 15 -num_worker 8 -patch_size 64 

Reduce the batch_size and patch_size in order to fit the volumes in available GPU memory. num_worker parameter allows creation of multiple worker threads for tqdm. Set this to 0 if you encounter Dataloader worker errors.
The arg --contain specifies not to bind target system's /home and /tmp folders to the container which is the default behaviour of Singularity. This is necessary to avoid conflicts between host and the container packages. Additionally, you can use --clean-env to unlink host env variables from the container. The arg --nv is the flag to enable the container access to available NVIDIA GPU driver and corresponding CUDA libraries. Use --rocm to enable access to AMD Radeon GPUs. Finally, --bind specifies mounting of selected directories on the container and therefore exposing the container to only necessary files and directories.

⁠Training the underlying UNet-MSS from scratch

To train the UNet-MSS network with Semi-Supervised Deformation-Aware Learning from scratch, please use

singularity run --contain --nv --bind <path-to-folder-containing-nifti-vols>:/data,<path-to-op>:/output uhura_unet_mss.sif -train True -with_mip False -pre_train False

Arrange the dataset folder as shown below:
/<path-to-dataset-folder>
---/train
---/train_label
---/validate
---/validate_label
---/test
---/test_label
*All the folders mentioned above must contain atleast one 3D nifti volume

If you wish to perform training with k-fold cross validation include below args

-cross_validate True -k_folds 5

In this case, please arrange the daset folder as shown below:
/<path-to-dataset-folder>
---/train
---/train_label
---/test
---/test_label
*All the folders mentioned above must contain atleast one 3D nifti volume The network weights are stored at <path-to-op>/uhura_unet_mss/checkpoint/ as checkpointbest.pth after successful training. Although the Deformation-Aware learning can be used to optimize the pretrained network. So we recommend training UNet-MSS without deformation awareness and then tune the weights using DS6.

⁠Tuning the pretrained UNet-MSS network with MIP loss

The network architecture is illustrated in the above mentioned papers. Please have a look at the network in https://github.com/ravi22342/DS6/tree/VSeg_with_MIP⁠ for more info. Once the pretrained weights are availble from training as specified above or if you wish to further tune the weights already available in the container, please use

singularity run --contain --nv --bind <path-to-folder-containing-nifti-vols>:/data,<path-to-op>:/output,<path-to-checkpointbest.pth>:/checkpoint uhura_unet_mss.sif -train True -with_mip False -pre_train True -load_path /checkpoint

*Provide the checkpoint path if the training was done from scratch using -load_path. If you would like to use the existing weights, do not specify this arg.

⁠Comparing segmentation results against labels and obtaining diff MIP overlays

If you wish to compare the segmentations that will be generated using the container against segmentation labels, please use

singularity run --contain --nv --bind <path-to-folder-containing-nifti-vols>:/data,<path-to-op>:/output,<path-to-checkpointbest.pth>:/checkpoint uhura_unet_mss.sif -test True -with_mip False -pre_train True -load_path /checkpoint

In this case the dataset folder must be arranged as below:
/<path-to-dataset-folder>
---/test
---/test_label
A results folder will be created at <path-to-op>/uhura_unet_mss/ and this will contain resulting segmentations, MIPs of diff overlays of labels and their respective segmentations and a spreadsheet containing evaluated metrics (Dice Score and IoU) for each nifti volume in /test folder.

⁠**args

⁠References