miTorch: Medical Imaging in PyTorch

By Mahdi Biparva (PhD in Computer Science)

This package implements deep learning modules for medical imaging application in PyTorch. It contains different modules in the data-pipeline such as the data-loaders, data-containers, transformations etc. In the model-pipeline, there are several segmentation neural networks, training logics, loss function, metrics etc.

License

TBC

Citing miTorch

If you find "miTorch: Medical Imaging in PyTorch" useful in your research, please consider citing the research paper:

TBC

Contents

1. Introduction
2. Features
3. Requirements: Software
4. Requirements: Hardware
5. Installation
6. Prerequisites
7. Preparation
8. Demo: 3D Medical Segmentation
9. Demo: Self-Supervised Pre-Training
10. Future Work
11. Contributors

Introduction

The primary goal is to have a solid, readable, modular, reliable, extendable PyTorch package for medical imaging application in deep learning. The input is 3D volumes of various modalities and the task could be segmentation, classification, and transfer learning.

To name a few, the learning tasks are:

• 3D Medical Segmentation:
  • Head-From-Brain / Skull-Stripping (HFB)
  • White-Matter-Hyperintensities (WMH)
• Robustness analysis test pipeline
• Self-supervised learning

Features

miTorch has currently the following capabilities and components:

• Robust 3D data loading modules for:
  • CT/MRI datasets:
    • Skull-stripping
    • White-matter hyperintensities
  • Electron-Microscopy datasets:
    • Neuron segmentation (counting)
    • Axon/Virus Tracing (Tractography)
    • Hippocampal Subfield Segmentation (multi-label)
• 3D data transformations and pipeline generation:
  • Spatial:
    • Cropping
    • Resizing
    • Axis Rotations
    • Flipping
    • Affine Transformations (Translation, Rotation, Scale, Shear)
  • Intensity:
    • Additive noise (Gaussian, Rice, etc)
    • Corrections (Gamma, Brightness, Contrast)
    • Bias Field
    • Blur
• Automatic Data transformation randomization
• Modular data-pipeline prototyping on-the-fly
• Seamless online patching and batching mechanisms
• Automatic Test Pipeline Generation and Evaluation
• Model Zoo containing:
  • Unet3D
  • Unet3D++ (NestedUnet3D)
  • CBAM
  • DUNet
  • DenseNet
  • SENet
  • VNet
  • DYNUnet
  • HighresNet
• Seamless 3D to 2D network conversion capability
• Various losses such as:
  • Dice
  • Focal
  • Hausdorff
  • Lovasz
  • MSE
• Weighted multi-loss training
• Various metrics such as
  • Jaccard
  • Hausdorff
  • Dice
  • F1
  • Relative Volume Difference
• Hyper-parameter optimization modes:
  • Manual grid search
  • Bayesian semi-automatic optimization search (GPyTorch|BoTorch|Ax)
  • Visualization and Logging (Tensorboard)
• Test Evaluation:
  • Automatic test transformation pipeline generation
  • Batch evaluation and result logging for analysis
• Checkpointing models
• Logging and Visualization (using Tensorboard)
• Automatic-Mixed-Precision (AMP) feature
• Data-Distributed feature (supporting AMP):
• Data Parallel (with GIL)
• Distributed Data Parallel (no GIL, multi node multi GPU)
• Model-Parallel (Under development)

Requirements: Software

Currently it relies on Python and PyTorch ecosystem.

Requirements: Hardware

GPU devices with CUDA capabilities are required.

Installation

There is no installation needed at this moment. You would simply need to call the main function.

Prerequisites

• Python 3.7
• PyTorch 1.4.0 (not tested on higher versions)
• CUDA 10.0 or higher

Preparation

TBC

Demo: 3D Medical Segmentation

TBC

Demo: Self-Supervised Pre-Training

TBC

Future Work

We are aiming to develop self-supervised learning modules to enhance the segmentation robustness.

Contributors

• Mahdi Biparva (core modeling and development)
• Parsa Esfahanian (self-supervised development)
• Braedyn Au (Tracing Segmentation development)
• Parisa Mojiri (experimentation)
• Lyndon Boone (experimentation)
• Maged Goubran (abstraction and methodologies)