volume-segmantics 0.4.0

A toolkit for semantic segmentation of volumetric data using PyTorch deep learning models

Volume Segmantics

A toolkit for semantic segmentation of volumetric data using PyTorch deep learning models.

Volume Segmantics provides a simple command-line interface and API that allows researchers to quickly train a variety of 2D PyTorch segmentation models (e.g. U-Net, U-Net++, FPN, DeepLabV3+) on their 3D datasets. These models use pre-trained encoders, enabling fast training on small datasets. Subsequently, the library enables using these trained models to segment larger 3D datasets, automatically merging predictions made in orthogonal planes and rotations to reduce artifacts that may result from predicting 3D segmentation using a 2D network.

Given a 3d image volume and corresponding dense labels (the segmentation), a 2d model is trained on image slices taken along the x, y, and z axes. The method is optimised for small training datasets, e.g a single dataset in between $128^3$ and $512^3$ pixels. To achieve this, all models use pre-trained encoders and image augmentations are used to expand the size of the training dataset.

This work utilises the abilities afforded by the excellent segmentation-models-pytorch library in combination with augmentations made available via Albumentations and from the MONAI library. Also the metrics and loss functions used make use of the hard work done by Adrian Wolny in his pytorch-3dunet repository.

Requirements

A machine capable of running CUDA enabled with a recent version of Pytorch (2.5 or greater is recommended). This generally means a reasonably modern NVIDIA GPU. The exact requirements differ according to operating system. For example on Windows you will need Visual Studio Build Tools as well as CUDA Toolkit installed see the CUDA docs for more details.

Installation

Conda/Virtualenv

Currently, the release version of Volume Segmantics (0.4.0) is not on PyPy. To install, create a new Conda/Miniconda environment (We recommend Python 3.11-3.12) and clone the Volume-Segmentics package directly from the online repo and go to the 'volume-segmantics' directory. Then:

pip install poetry
poetry install

To check that that the Pytorch version downloaded by the install is a CUDA version start a Python interpreter in the environment you've just installed Volume Segmantics into. Then:

import torch; torch.__version__

It should print out a version number (eg. 2.5.1+cu124) the "cu" indicates it's a CUDA version. If not ("cpu"), you will need to go to https://pytorch.org/get-started/locally/ and install a version of Pytorch that is suited to your CUDA version.

Docker/Apptainer Container (quay.io)

A container image with the latest published release is available on the Rosalind Franklin's quay.io instance. You can pull and run this using Apptainer:

apptainer run --nv docker://quay.io/rosalindfranklininstitute/volume-segmantics

Or docker:

docker run \
    --gpus all \
    --ipc=host \
    -v /path/to/data:/data
    quay.io/rosalindfranklininstitute/volume-segmantics

Note --ipc=host grants the container access to host shared memory, since the default allocation in Docker (64M) may be insufficient. More details can be found on the Volume EM Container documentation.

Configuration and command line use

After installation, two new commands will be available from your terminal whilst your environment is activated, model-train-2d and model-predict-2d.

These commands require access to some settings stored in YAML files. These need to be located in a directory named volseg-settings within the directory where you are running the commands. The settings files can be copied from here.

The file 2d_model_train_settings.yaml can be edited in order to change training parameters such as number of epochs, loss functions, evaluation metrics and also model and encoder architectures. The file 2d_model_predict_settings.yaml can be edited to change parameters such as the prediction "quality" e.g "low" quality refers to prediction of the volume segmentation by taking images along a single axis (images in the (x,y) plane). For "medium" and "high" quality, predictions are done along 3 axes and in 12 directions (3 axes, 4 rotations) respectively, before being combined by maximum probability. Check normalisation settings carefully for your intended use, different datasets may require different choices from the default. Double-check that the prediction normalisation settings match those used for training.

For training a 2d model on a 3d image volume and corresponding labels

Run the following command. Input files can be in HDF5 or multi-page TIFF format or MRC format.

model-train-2d --data path/to/image/data.tiff --labels path/to/corresponding/segmentation/labels.tiff

Paths to multiple data and label volumes can be added after the --data and --labels flags respectively. A model will be trained according to the settings defined in /volseg-settings/2d_model_train_settings.yaml and saved to your working directory. In addition, a figure showing "ground truth" segmentation vs model segmentation for some images in the validation set will be saved.

For 3d volume segmentation prediction using a 2d model

Run the following command. Input image files can be in HDF5 or multi-page TIFF format or MRC format.

model-predict-2d path/to/model_file.pytorch path/to/data_for_prediction.tiff

The input data will be segmented using the input model following the settings specified in volseg-settings/2d_model_predict_settings.yaml. Depending on whether OUTPUT_FORMAT in config.py is set to "tif" or "hdf" or "mrc", a TIFF or HDF5 or MRC file containing the segmented volume will be saved to your working directory.

Training features

Volume Segmantics supports training using a variety of U-net encoder-decoder architectures and encoder choices, including Transformer- based models such as DINO. It allows purely axial or tri-planar/multi-axis prediction as well as supporting multiple decoder and multiple head architectures and multi-task training through a configurable pipeline.yaml.

Volume Segmantics supports 2.5D slicing, which creates multi-channel images from adjacent slices in the volume. This feature can be enabled by setting use_2_5d_slicing: True in the training settings file. This approach provides the model with spatial context from adjacent slices. The encoder adjusts to use the number of input channels specified by the num_slices parameter in volseg-settings/2d_model_train_settings.yaml (when 2.5D slicing is enabled). When using 2.5D slicing for training, you must also enable 2.5D prediction for inference by setting use_2_5d_prediction: True in your 2d_model_predict_settings.yaml file and set the num_slices parameter there.

Tutorial using example data

A tutorial is available here that provides a walk-through of how to segment blood vessels from synchrotron X-ray micro-CT data collected on a sample of human placental tissue.

Contributing

We welcome contributions from the community. Please take a look at out contribution guidelines for more information.

Citation

If you use this package for your research, please cite:

King O.N.F, Bellos, D. and Basham, M. (2022). Volume Segmantics: A Python Package for Semantic Segmentation of Volumetric Data Using Pre-trained PyTorch Deep Learning Models. Journal of Open Source Software, 7(78), 4691. doi: 10.21105/joss.04691

@article{King2022,
    doi = {10.21105/joss.04691},
    url = {https://doi.org/10.21105/joss.04691},
    year = {2022},
    publisher = {The Open Journal},
    volume = {7},
    number = {78},
    pages = {4691},
    author = {Oliver N. F. King and Dimitrios Bellos and Mark Basham},
    title = {Volume Segmantics: A Python Package for Semantic Segmentation of Volumetric Data Using Pre-trained PyTorch Deep Learning Models},
    journal = {Journal of Open Source Software} }

References

Albumentations

Buslaev, A., Iglovikov, V.I., Khvedchenya, E., Parinov, A., Druzhinin, M., and Kalinin, A.A. (2020). Albumentations: Fast and Flexible Image Augmentations. Information 11. https://doi.org/10.3390/info11020125.

MONAI

Cardoso, M. Jorge, Wenqi Li, Richard Brown, Nic Ma, Eric Kerfoot, Yiheng Wang, Benjamin Murrey et al. (2022) "Monai: An open-source framework for deep learning in healthcare." arXiv preprint arXiv:2211.02701.

Segmentation Models PyTorch

Yakubovskiy, P. (2020). Segmentation Models Pytorch. GitHub.

PyTorch-3dUnet

Wolny, A., Cerrone, L., Vijayan, A., Tofanelli, R., Barro, A.V., Louveaux, M., Wenzl, C., Strauss, S., Wilson-Sánchez, D., Lymbouridou, R., et al. (2020). Accurate and versatile 3D segmentation of plant tissues at cellular resolution. ELife 9, e57613. https://doi.org/10.7554/eLife.57613.