elektronn3.models.unet module

This is a modified version of the U-Net CNN architecture for biomedical image segmentation. U-Net was originally published in https://arxiv.org/abs/1505.04597 by Ronneberger et al.

A pure-3D variant of U-Net has been proposed by Çiçek et al. in https://arxiv.org/abs/1606.06650, but the below implementation is based on the original U-Net paper, with several improvements.

This code is based on https://github.com/jaxony/unet-pytorch (c) 2017 Jackson Huang, released under MIT License, which implements (2D) U-Net with user-defined network depth and a few other improvements of the original architecture.

Major differences of this version from Huang’s code:

  • Operates on 3D image data (5D tensors) instead of 2D data

  • Uses 3D convolution, 3D pooling etc. by default

  • planar_blocks architecture parameter for mixed 2D/3D convnets (see UNet class docstring for details)

  • Improved tests (see the bottom of the file)

  • Cleaned up parameter/variable names and formatting, changed default params

  • Updated for PyTorch 1.3 and Python 3.6 (earlier versions unsupported)

  • (Optional DEBUG mode for optional printing of debug information)

  • Extended documentation