Trained models

This repository contains pre-trained models for 3D neuroimaging data processing. These models can be used for their original purpose or for transfer learning on a new task. For example, a pre-trained brain extraction network can be trained on a tumor-labeling task. Models are included based on <org-name>/<model-name>/<version>/weights structure. Some models (such as kwyk and braingen) have various types which means there was different structural chracteristic during training that lead to different trained models. Therefor, the path for these models are changes as <org-name>/<model-name>/<version>/<model-type>/weights Instructions to add a model can be find here.

Neuronets organization

These models were trained using the Nobrainer framework, which wraps TensorFlow/Keras.

• brainy: 3D U-Net brain extraction model
• ams: automated meningioma segmentation model
• kwyk: bayesian neural network for brain parcellation and uncertainty estimation (Tensorflow/estimator)
• braingen: progressive generation of T1-weighted brain MR scans

UCL organization

• SynthSeg: 3D brain MRI segmentation model (Tensorflow/keras)
• SynthSR: 3D brain MRI (& CT) super resolution model (Tensorflow/keras)

DDIG Organization

• SynthMorph: contrast agnostic registration model (Tensorflow/keras)
• VoxelMorph: learning based registration model (Tensorflow/keras)

Laboratory for Computational Neuroscience (lcn) Organization

• ParcNet: cortical parcellation model (pytorch)

Downloading models

This repository is a datalad dataset. To get the models, you need to install datalad and datalad-osf to your environment.

datalad clone https://github.com/neuronets/trained-models
cd trained-models
git-annex enableremote osf-storage

to download all the models,

datalad get -s osf-storage .

to get a specific model you can pass the path of the model to the datalad get.

datalad get -s osf-storage neuronets/ams/0.1.0/weights/meningioma_T1wc_128iso_v1.h5

datalad get -s osf-storage neuronets/braingen/0.1.0

Using models for inference or training

You can use the Nobrainer-zoo toolbox for inference and re-training of the models without installing any additional model dependencies.

Loading models for training with python and tensorflow/keras

You can use tensorflow.keras module to load a tensorflow model.

import tensorflow as tf

model = tf.keras.models.load_model("neuronets/brainy/0.1.0/brain-extraction-unet-128iso-model.h5")
model.fit(...)

You can see a transfer learning example here, and an example of brain MRI generation using braingen models can be find here.

All models are available for re-training or transfer learning purposes except the kwyk model. The kwyk model weights are not available in a tf2 keras format (We are working to make it available in near future). The kwyk models can be loaded with tf.saved_model.load.

model = tf.saved_model.load(model_path)
predictor = model.signatures["serving_default"]

or you can use nobrainer predict_by_estimator function. check an example here.