Mitochondria are the primary energy providers for cell activities, thus essential for metabolism. Quantification of the size and geometry of mitochondria is not only crucial to basic neuroscience research, but also informative to the clinical studies of several diseases including bipolar disorder and diabetes.
This tutorial has two parts. In the first part, you will learn how to make pixel-wise class prediction on the widely used benchmark dataset released by Lucchi et al. in 2012. In the second part, you will learn how to predict the instance masks of individual mitochondrion from the large-scale MitoEM dataset released by Wei et al. in 2020.
This section provides step-by-step guidance for mitochondria segmentation with the EM benchmark datasets released by Lucchi et al. (2012). We consider the task as a semantic segmentation task and predict the mitochondria pixels with encoder-decoder ConvNets similar to the models used in affinity prediction in neuron segmentation. The evaluation of the mitochondria segmentation results is based on the F1 score and Intersection over Union (IoU).
Different from other EM connectomics datasets used in the tutorials, the dataset released by Lucchi et al. is an isotropic dataset, which means the spatial resolution along all three axes is the same. Therefore a completely 3D U-Net and data augmentation along z-x and z-y planes besides x-y planes are preferred.
All the scripts needed for this tutorial can be found at
pytorch_connectomics/scripts/. Need to pass the argument
--config-file configs/Lucchi-Mitochondria.yaml during training and inference to load the required configurations for this task.
The pytorch dataset class of lucchi data is
Qualitative results of the model prediction on the mitochondria segmentation dataset released by Lucchi et al., without any post-processing.
1 - Get the data¶
Download the dataset from our server:
For description of the data please check the author page.
2 - Run training¶
source activate py3_torch CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python -m torch.distributed.launch \ --nproc_per_node=4 --master_port=2345 scripts/main.py --distributed \ --config-file configs/Lucchi-Mitochondria.yaml
Similar to the neuron segmentation tutorial, we use distributed data-parallel training considering its high efficiency, and also to enable synchronized batch normalization (SyncBN).
3 - Visualize the training progress¶
tensorboard --logdir outputs/Lucchi_UNet/
4 - Inference on test data¶
source activate py3_torch CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python -u scripts/main.py \ --config-file configs/Lucchi-Mitochondria.yaml --inference \ --checkpoint outputs/Lucchi_UNet/volume_100000.pth.tar
5 - Run evaluation¶
Since the ground-truth label of the test set is public, we can run the evaluation locally:
from connectomics.utils.evaluation import get_binary_jaccard pred = pred / 255. # output is casted to uint8 with range [0,255]. gt = (gt!==0).astype(np.uint8) thres = [0.4, 0.6, 0.8] # evaluate at multiple thresholds. scores = get_binary_jaccard(pred, gt, thres)
The prediction can be further improved by conducting median filtering to remove noise:
from connectomics.utils.evaluate import get_binary_jaccard from connectomics.utils.process import binarize_and_median pred = pred / 255. # output is casted to uint8 with range [0,255]. pred = binarize_and_median(pred, size=(7,7,7), thres=0.8) gt = (gt!==0).astype(np.uint8) scores = get_binary_jaccard(pred, gt) # prediction is already binarized
Our pretained model achieves a foreground IoU and IoU of 0.892 and 0.943 on the test set, respectively. The results are better or on par with state-of-the-art approaches. Please check BENCHMARK.md for detailed performance comparison and the pre-trained models.
This section provides step-by-step guidance for mitochondria segmentation with our benchmark datasets MitoEM.
We consider the task as 3D instance segmentation task and provide three different confiurations of the model output.
The model is
UNet3D, similar to the one used in neuron segmentation.
The evaluation of the segmentation results is based on the AP-75 (average precision with an IoU threshold of 0.75).
Complex mitochondria in the MitoEM dataset:(a) mitochondria-on-a-string (MOAS), and (b) dense tangle of touching instances. Those challenging cases are prevalent but not covered in previous datasets.
The MitoEM dataset has two sub-datasets MitoEM-Rat and MitoEM-Human based on the source of the tissues. Three training configuration files on MitoEM-Rat
are provided in
pytorch_connectomics/configs/MitoEM/ for different learning setting as described in this paper.
Since the dataset is very large and can not be directly loaded into memory, we designed the
connectomics.data.dataset.TileDataset class that only
loads part of the whole volume each time by opening involved
1 - Dataset introduction¶
im: includes 1,000 single-channel
*.pngfiles (4096x4096) of raw EM images (with a spatial resolution of 30x8x8 nm). The 1,000 images are splited into 400, 100 and 500 slices for training, validation and inference, respectively.
mito_train/: includes 400 single-channel
*.pngfiles (4096x4096) of instance labels for training. Similarly, the
mito_val/folder contains 100 slices for validation. The ground-truth annotation of the test set (rest 500 slices) is not publicly provided but can be evaluated online at the MitoEM challenge page.
To run training, JSON files containing the metadata of the dataset (for both images and labels) need to be provided. Example
JSON files can be found in
2 - Model configuration¶
*.yaml files for different learning targets:
MitoEM-R-A.yaml: output 3 channels for predicting the affinty between voxels.
MitoEM-R-AC.yaml: output 4 channels for predicting both affinity and instance contour.
MitoEM-R-BC.yaml: output 2 channels for predicting both the binary foreground mask and instance contour. This configuration achieves the best overall performance according to our experiments.
3 - Run training¶
We show examples for running the training script for the U3D-BC model:
By default the path of images and labels are not specified. To
run the training scripts, please revise the
DATASET.INPUT_PATH options in
The options can also be given as command-line arguments without changing of the
yaml configuration files.
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python -m torch.distributed.launch \ --nproc_per_node=4 --master_port=4321 scripts/main.py --distributed \ --config-base configs/MitoEM/MitoEM-R-Base.yaml \ --config-file configs/MitoEM/MitoEM-R-BC.yaml
4 - Visualize the training progress¶
We use TensorBoard to visualize the training progress. For Harvard FASRC cluter users, more info can be found here.
tensorboard --logdir outputs/MitoEM_R_BC/
5 - Run inference¶
Run inference on validation/test image volumes (suppose the model is optimized for 100k iterations):
python -u scripts/main.py \ --config-base configs/MitoEM/MitoEM-R-Base.yaml \ --config-file configs/MitoEM/MitoEM-R-BC.yaml --inference \ --checkpoint outputs/MitoEM_R_BC/checkpoint_100000.pth.tar
Please change the
configs/MitoEM-R-*.yaml based on your own data path.
6 - Post-process¶
The post-processing step requires merging output volumes and applying watershed segmentation.
As mentioned before, the dataset is very large and can hardly be directly loaded into memory for
processing. Therefore our code run prediction on smaller chunks sequentially, which produces
*.h5 files with the coordinate information. To merge the chunks into a single volume
and apply the segmentation algorithm:
import glob import numpy as np from connectomics.data.utils import readvol from connectomics.utils.process import bc_watershed output_files = 'outputs/MitoEM_R_BC/test/*.h5' # output folder with chunks chunks = glob.glob(output_files) vol_shape = (2, 500, 4096, 4096) # MitoEM test set pred = np.ones(vol_shape, dtype=np.uint8) for x in chunks: pos = x.strip().split("/")[-1] print("process chunk: ", pos) pos = pos.split("_")[-1].split(".").split("-") pos = list(map(int, pos)) chunk = readvol(x) pred[:, pos:pos, pos:pos, pos:pos] = chunk # This function process the array in numpy.float64 format. # Please allocate enough memory for processing. segm = bc_watershed(pred, thres1=0.85, thres2=0.6, thres3=0.8, thres_small=1024)
Then the segmentation map should be ready to be submitted to the MitoEM challenge website for evaluation. Please note that this tutorial only take the MitoEM-Rat set as an example. The MitoEM-Human set also need to be segmented for online evaluation.