Segmentation Models Pytroch 3D

Python library with Neural Networks for Volume (3D) Segmentation based on PyTorch.

This library is based on famous Segmentation Models Pytorch library for images. Most of the documentation can be used directly from there.

Installation

• Type 1: pip install segmentation-models-pytorch-3d
• Type 2: Copy segmentation_models_pytorch_3d folder from this repository in your project folder.

Quick start

Segmentation model is just a PyTorch nn.Module, which can be created as easy as:

import segmentation_models_pytorch_3d as smp
import torch

model = smp.Unet(
    encoder_name="efficientnet-b0", # choose encoder, e.g. resnet34
    in_channels=1,                  # model input channels (1 for gray-scale volumes, 3 for RGB, etc.)
    classes=3,                      # model output channels (number of classes in your dataset)
)

# Shape of input (B, C, H, W, D). B - batch size, C - channels, H - height, W - width, D - depth
res = model(torch.randn(4, 1, 64, 64, 64)) 

Models

Architectures

• Unet [paper] [docs]
• Unet++ [paper] [docs]
• MAnet [paper] [docs]
• Linknet [paper] [docs]
• FPN [paper] [docs]
• PSPNet [paper] [docs]
• PAN [paper] [docs]
• DeepLabV3 [paper] [docs]

Encoders

The following is a list of supported encoders in the SMP. Select the appropriate family of encoders and click to expand the table and select a specific encoder and its pre-trained weights (encoder_name and encoder_weights parameters).

Timm 3D encoders

We now support encoders from timm_3d library. Full list available here. To use them add tu- before encoder name. Example:

encoder_name = 'tu-maxvit_base_tf_224.in21k'
model = smp.Unet(
    encoder_name=encoder_name,
    encoder_weights=None,
    in_channels=3,
    classes=1,
)

Notes for 3D version

Input size

Recommended input size for backbones can be calculated as: K = pow(N, 2/3). Where N - is size for input image for the same model in 2D variant.

For example for N = 224, K = 32. For N = 512, K = 64.

Strides

Typical strides for 2D case is 2 for H and W. It applied depth times (in almost all cases 5 times). So input image reduced from (224, 224) to (7, 7) on final layers. For 3D case because of very massive input, it's sometimes useful to control strides for every dimension independently. For this you can use input variable strides, which default values is: strides=((2, 2, 2), (2, 2, 2), (2, 2, 2), (2, 2, 2), (2, 2, 2)). Example:

Let's say you have input data of size: (224, 128, 12). You can use strides like that: ((2, 2, 2), (4, 2, 1), (2, 2, 2), (2, 2, 1), (1, 2, 3)). Output shape for these strides will be: (7, 4, 1)

import segmentation_models_pytorch_3d as smp
import torch

model = smp.Unet(
    encoder_name="resnet50",        
    in_channels=1,                  
    strides=((2, 2, 2), (4, 2, 1), (2, 2, 2), (2, 2, 1), (1, 2, 3)),
    classes=3, 
)

res = model(torch.randn(4, 1, 224, 128, 12)) 

Note: Strides currently supported by resnet-family and densenet models with Unet decoder only.

Related repositories

• https://github.com/qubvel/segmentation_models.pytorch - original segmentation 2D repo
• segmentation_models_3D - segmentation models in 3D for keras/tensorflow
• timm_3d - classification models in 3D for pytorch
• volumentations - 3D augmentations

Citation

If you find this code useful, please cite it as:

@article{solovyev20223d,
  title={3D convolutional neural networks for stalled brain capillary detection},
  author={Solovyev, Roman and Kalinin, Alexandr A and Gabruseva, Tatiana},
  journal={Computers in Biology and Medicine},
  volume={141},
  pages={105089},
  year={2022},
  publisher={Elsevier},
  doi={10.1016/j.compbiomed.2021.105089}
}

To Do List

• Support for strides for all encoders