Flemme: A FLExible and Modular Learning Platform for MEdical Images

Flemme is a flexible and modular learning platform for medical images.

The above is an overview of Flemme. In Flemme, we separate encoders from the model architectures, enabling fast model construction via different combinations for medical image segmentation, reconstruction and generation.

We support encoders using convolution, transformer and SSM (state-space model) as backbones. For model architectures, we implement the base architecture following an encoder-decoder style. Several derived architectures are used for segmentation, reconstruction and generation tasks.

Now, Flemme also supports point cloud modeling! For example, an overview of models built with PointMamba2Encoder is illustrated in the following figure:

Now, let’s get Started with Flemme!

Install Flemme

Requirement list

You need to install some necessary packages before using flemme.

torch torchvision simpleitk nibabel matplotlib scikit-image scikit-learn tensorboard

The following packages are optional. You can modify Gloal Configuration to disable some components of Flemme so that you don’t need to install the corresponding required packages.

# For swin-transformer
einops
# For vision mamba
mamba-ssm (CUDA version >= 11.6)
# For io and evaluation of point cloud
POT plyfile KNN-CUDA

After installing the required packages, now you can install Flemme into your environment.

git clone git@github.com:wlsdzyzl/flemme.git
cd flemme
python setup.py install

Train a Model

It’s quite easy to build and train a deep learning model with Flemme, you don’t even need to write any codes. All things can be down through a yaml Configuration file.

Let’s start with a easy one: a simple U-Net for medical image segmentation. Firstly, you need to create a config file and copy the contents from train_unet_sem.yaml. Then, run the following command:

train_flemme path/to/config/train_unet_sem.yaml

To visualize the training process

tensorboard --logdir path/to/checkpoint/CVC-ClinicDB/UNet

To test the results, you need to create a test config file that looks like test_unet_sem.yaml and run:

test_flemme path/to/config/test_unet_sem.yaml

In the following modules, we elaborate supported architectures, encoders and other components of flemme.

Modules

• Model Architectures
  • Base Architectures
  • Classification Model
  • Segmentation Model
  • Auto-Encoder
  • Denoising Diffusion Probabilistic Model
  • Hierarchical Architectures
• Encoders
  • Image Encoders
  • Point Cloud Encoders
  • Graph Encoders
• Data Loading
  • Datasets
  • Data Augmentations
• Losses and Evaluation Metrics
  • Losses
  • Evaluation Metrics
• Sampler
• Global Configuration
• Scripts
  • Common Datasets
  • FastMRI Dataset
  • Performance Evaluation
  • Select Samples for Visualization
  • Render Points and Mesh
• Play with Flemme
  • Toy Example for Diffusion model
  • MINST
  • CIFA10
• MedPointS Dataset

Example configuration files

• train_unet_sem.yaml
• test_unet_sem.yaml
• train_unet_ddim.yaml
• train_unet_ddpm.yaml
• train_pointnet_clm.yaml
• test_pointnet_clm.yaml

More configuration files please refer to resources.