GKalliatakis / camus-segmentation-pytorch

PyTorch implementation of the two U-Net-based architectures described in "Deep Learning for Segmentation using an Open Large-Scale Dataset in 2D Echocardiography"

CAMUS Segmentation: Semantic segmentation with PyTorch using an open large-scale dataset in 2D Echocardiography

Detailed architectures

The two U-Net implementations presented in the paper are summarised in tables I and II below. Both differ from the original U-Net proposed by Ronneberger et al. U-Net 1 is optimised for speed, while U-Net 2 is optimized for accuracy. U-Net 1 shows a more compact architecture, with the addition of one downsampling level. In the U-Net 2 design, the number of filters per convolutional layers increases and decreases linearly from an initial kernel size of 48, which makes for a wider net.

Usage

import torch
from camus_unet.camus_unet1 import CamusUnet1

if torch.cuda.is_available():
    device = torch.device('cuda')
else:
    device = torch.device('cpu')

model = CamusUnet1()  # initialise the U-NET 1 model

# move initialised model to chosen device
model = model.to(device)

# the usual training loop goes here...

References

• Deep Learning for Segmentation using an Open Large-Scale Dataset in 2D Echocardiography
• CAMUS project
• CAMUS Exploratory Data Analysis
• PyTorch U-Net
• PlotNeuralNet

License

Distributed under the MIT License. See LICENSE file for more information.