Global Voxel Transformer Networks for Augmented Microscopy

Publication

The paper is available at https://www.nature.com/articles/s42256-020-00283-x.

If using this code , please cite our paper.

@article{wang2021global,
  title={Global voxel transformer networks for augmented microscopy},
  author={Wang, Zhengyang and Xie, Yaochen and Ji, Shuiwang},
  journal={Nature Machine Intelligence},
  volume={3},
  number={2},
  pages={161--171},
  year={2021},
  publisher={Nature Publishing Group}
}

Global Voxel Transformer Networks (GVTNets)

Augmented Microscopy Tasks

System Requirements

We highly recommend using the Linux operating system. Using an nVIDIA GPU with >=11GB memory is highly recommended, although this tool can be used with CPU only.

We used Ubuntu 16.04.6 LTS (GNU/Linux 4.4.0-141-generic x86_64)) and an nVIDIA GeForce RTX 2080 Ti GPU with nVIDIA driver version 430.50 and CUDA version 10.1.

Installation

Environment setup

• We highly recommend installing Anaconda for a simple environment setup and installation.
• Download the code:

git clone https://github.com/zhengyang-wang/GVTNets.git
cd GVTNets

If you do not want to use provided pretrained models for reproduction of results in our paper, you may skip downloading save_dir.

• Create a virtual environment with required packages (it may take several minutes):

conda env create -f gvtnet.yml

• Activate the virtual environment:

conda activate gvtnet

or

source activate gvtnet

Choose whichever works for you.

Usage

Label-free prediction of 3D fluorescence images from transmitted-light microscopy

Preparation

• Download label-free datasets.

• Untar all the datasets under one folder. Suppose the folder is data/, it should contain 13 datasets named beta_actin, fibrillarin, etc. In addition, data/beta_actin/ should contain images only.

• Give execution permission to scripts:

chmod +x ./scripts/label-free/*.sh

• Change RAW_DATASET_DIR in train.sh, train_dic.sh, train_membrane.sh and predict.sh, predict_dic.sh, predict_membrane.sh to the path to your folder that saves all the untarred datasets.

Training

• Modify network_configure.py according to your design. For users who are not familiar with deep learning, simply copy the content of network_configures/gvtnet_label-free.py to network_configure.py.

• Train the GVTNets for datasets except dic_lamin_b1 and membrane_caax_63x:

./scripts/label-free/train.sh [dataset] [gpu_id] [model_name]

• For dic_lamin_b1 and membrane_caax_63x, use train_dic.sh and train_membrane.sh, respectively:

./scripts/label-free/train_dic.sh [gpu_id] [model_name]
./scripts/label-free/train_membrane.sh [gpu_id] [model_name]

Example:

If you want to train a GVTNet called your-gvtnet on beta_actin using the GPU #1, run:

./scripts/label-free/train.sh beta_actin 1 your-gvtnet

After training, you will find:

• Transformed datasets are saved under save_dir/label-free/beta_actin/datasets/. This process will only be performed for the first run.
• The content in network_configure.py is saved as network_configures/your-gvtnet.py.
• Model checkpoints are saved under save_dir/label-free/beta_actin/models/your-gvtnet/.

Note: Always give a different model_name when you use a different network_configure.py. This tool will used model_name to track different network configures.

Prediction

• Predict the testing set using saved model checkpoints for datasets except dic_lamin_b1 and membrane_caax_63x:

./scripts/label-free/predict.sh [dataset] [gpu_id] [model_name] [checkpoint_num]

• For dic_lamin_b1 and membrane_caax_63x, use predict_dic.sh and predict_membrane.sh, respectively:

./scripts/label-free/predict_dic.sh [gpu_id] [model_name] [checkpoint_num]
./scripts/label-free/predict_membrane.sh [gpu_id] [model_name] [checkpoint_num]

Example:

If you have trained a GVTNet called your-gvtnet on beta_actin, and want to make prediction for the testing set with the saved model checkpoints after training for 75,000 minibatch iterations, run:

./scripts/label-free/predict.sh beta_actin 1 your-gvtnet 75000

After prediction, you will find:

• Prediction results are saved under save_dir/label-free/beta_actin/results/your-gvtnet/checkpoint_75000/.

Note: If your GPU memory is limited, set gpu_id to -1 for CPU prediction.

Evaluation

• Evaluate the prediction on the testing set for your saved model checkpoints:

./scripts/label-free/evaluate_dir.sh [dataset] [model_name] [checkpoint_num]

• Users can use evaluate_file.sh to evaluate any single prediction.

Example:

To evaluate the predictions made in the last example, run:

./scripts/label-free/evaluate_dir.sh beta_actin your-gvtnet 75000

Use provided pretrained models for reproduction of results in our paper

• Make predictions for datasets except dic_lamin_b1 and membrane_caax_63x:

./scripts/label-free/predict.sh [dataset] [gpu_id] gvtnet_label-free_pretrained pretrained

• For dic_lamin_b1 and membrane_caax_63x, use predict_dic.sh and predict_membrane.sh, respectively:

./scripts/label-free/predict_dic.sh [gpu_id] gvtnet_label-free_pretrained pretrained
./scripts/label-free/predict_membrane.sh [gpu_id] gvtnet_label-free_pretrained pretrained

• Evaluate the results:

./scripts/label-free/evaluate_dir.sh [dataset] gvtnet_label-free_pretrained pretrained

You will obtain the exact number reported in the supplementary of our paper. The estimated running time can be found in the paper.

Content-aware 3D image denoising and 3D to 2D image projection (CARE)

Preparation

• Download Denoising_Planaria.tar.gz, Denoising_Tribolium.tar.gz, Projection_Flywing.tar.gz from CARE datasets.

• Extract the datasets. Each should contain a train_data folder and a test_data folder.

• Give execution permission to scripts:

chmod +x ./scripts/care_denoising/*.sh
chmod +x ./scripts/care_projection/*.sh

• Change NPZ_DATASET_DIR in train_[Planaria|Tribolium|Flywing].sh to the path to the corresponding train_data folder.

• Change TEST_DATASET_DIR in train_[Planaria|Tribolium|Flywing].sh to the path to the corresponding test_data folder.

Training

• Modify network_configure.py according to your design. For users who are not familiar with deep learning, simply copy the content of network_configures/gvtnet_care.py to network_configure.py.

• Train the GVTNets:

./scripts/care_denoising/train_[Planaria|Tribolium].sh [gpu_id] [model_name]
./scripts/care_projection/train_Flywing.sh [gpu_id] [model_name]

Note: Always give a different model_name when you use a different network_configure.py. This tool will used model_name to track different network configures.

Prediction

• Predict the testing set using saved model checkpoints:

./scripts/care_denoising/predict_[Planaria|Tribolium].sh [condition] [gpu_id] [model_name] [checkpoint_num]
./scripts/care_projection/predict_Flywing.sh [condition] [gpu_id] [model_name] [checkpoint_num]

Note: If your GPU memory is limited, set gpu_id to -1 for CPU prediction.

Evaluation

• Evaluate the prediction on the testing set for your saved model checkpoints:

./scripts/care_denoising/evaluate_[Planaria|Tribolium].sh [condition] [model_name] [checkpoint_num]
./scripts/care_projection/evaluate_Flywing.sh [condition] [model_name] [checkpoint_num]

Use provided pretrained models for reproduction of results in our paper

• Make predictions:

./scripts/care_denoising/predict_[Planaria|Tribolium].sh [condition] [gpu_id] gvtnet_care_pretrained pretrained
./scripts/care_projection/predict_Flywing.sh [condition] [gpu_id] gvtnet_care_pretrained pretrained

• Evaluate the results:

./scripts/care_denoising/evaluate_[Planaria|Tribolium].sh [condition] gvtnet_care_pretrained pretrained
./scripts/care_projection/evaluate_Flywing.sh [condition] gvtnet_care_pretrained pretrained

You will obtain the exact number reported in the supplementary of our paper. The estimated running time can be found in the paper.

To train and inference/test with your own datasets.

• To prepare your training dataset: (randomly) crop the training image pairs into two sets of patches and save them into npz file(s). The npz file(s) can be either a single npz file containing all training data structured as:

  {'X': (n_sample, n_channel, (depth,) height, width),
   'Y': (n_sample, n_channel, (depth,) height, width)}
  

  or multiple npz files where each one contains one training sample structured as:

  {'X': (n_channel, (depth,) height, width),
   'Y': (n_channel, (depth,) height, width)}
  

  If your data contains uncropped images with different sizes, use the later data structure. Check ./datasets/label-free/generate_npz_or_tiff.py for an example.

• To train with the dataset:

  python train.py [--args]
  

  You will need to specify the arguments, such as npz_dataset_dir, gpu_id. You can refer to the scripts for the example argument settings. You can also tune the model parameters by modifying network_configure.py.

Note: Always give a different model_name when you use a different network_configure.py. This tool will used model_name to track different network configures.

 Explanation to some arguments:
 ```
 --already_cropped: include it only when training images are already cropped to patches. If not, 
                    you need to specify the --train_patch_size and the image will be automatically 
                    cropped.
 --proj_model: whether to use ProjectionNet to project 3D images to 2D, only used in 3D-to-2D 
               transform task, e.g. CARE Flywings projection'.
 --offset: whether to add inputs to the outputs (so that the output is considered as an offset of 
           input image). It is applied in CARE models.
 --probalistic: whether to train with probalistic loss, used in CARE models.
 ```

• To prepare the prediction and evaluation data: the prediction and evaluation accept the tif/tiff files as inputs. Each tif/tiff file contains one image of shape

  [(depth,) height, width]
  

  The ground truth files used for evaluation should have the same names as their corresponding input files and be stored in a different directory to the result or the input files.

• To predict and evaluate the dataset: run the following command,

  python predict.py [--args]
  

  and then

  python evaluate.py [--args]
  

  You will need to specify the arguments for prediction and evaluation respectively, such as tiff_dataset_dir, gpu_id. You can refer to the scripts for the example argument settings.

  Explanation to some arguments:

  --cropped_prediction: suggested when having a GPU memory problem. The input images will be processed
                        patch by patch and assembled back together. If included, you also need to
                        specify the --predict_patch_size and --overlap.
  --CARE_normalize: include it when you need to use the percentile normalization used in CARE.